From 2839b2bd0c3b67b877ca2d7ff9fb7e2f183a9d5e Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 12 Jun 2026 02:51:58 +0000 Subject: [PATCH 1/3] fuse o prj Signed-off-by: Fangzhou Ai --- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 151 ++++++++++++++++-- 1 file changed, 141 insertions(+), 10 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 8104e808f670..9106e4906e6e 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -924,22 +924,122 @@ def _apply_inv_rope_ref( return _apply_gptj_inv_rope_ref(x, positions, rotary_emb.cos_sin_cache, rope_dim) -def rocm_inv_rope_einsum( +@triton.jit +def _inverse_rope_gptj_kernel( + o_ptr, # [T, H, D] input + out_ptr, # [T, H, D] bf16 output + pos_ptr, # [T] positions + cos_sin_ptr, # [P, rope_dim] fp32 (cos[:half] | sin[half:]) + s_t, + s_h, # input row strides (last dim contiguous) + os_t, + os_h, # output row strides + cs_stride, # cos_sin_cache row stride + NOPE: tl.constexpr, # non-rope head dims (passed through) + HALF: tl.constexpr, # rope_dim // 2 + BLOCK_NOPE: tl.constexpr, + BLOCK_HALF: tl.constexpr, +): + """Fused inverse GPT-J RoPE on the trailing rope_dim of each (token, head). + + Mirrors ``DeepseekV4ScalingRotaryEmbedding.forward_native(inverse=True)`` + for the GPT-J (non-neox) layout, writing bf16 directly. Replaces the + clone + index_select + repeat_interleave + neg + stack + cat + cast chain + (~10 small kernels) with a single launch. + """ + t = tl.program_id(0) + h = tl.program_id(1) + in_base = t * s_t + h * s_h + out_base = t * os_t + h * os_h + + # NoPE lanes pass through unchanged (only cast to bf16). + n = tl.arange(0, BLOCK_NOPE) + nmask = n < NOPE + vals = tl.load(o_ptr + in_base + n, mask=nmask) + tl.store(out_ptr + out_base + n, vals.to(tl.bfloat16), mask=nmask) + + # RoPE lanes: out_even = a*cos + b*sin, out_odd = b*cos - a*sin + # (a = even lane, b = odd lane; sin negated for the inverse rotation). + pos = tl.load(pos_ptr + t).to(tl.int64) + k = tl.arange(0, BLOCK_HALF) + kmask = k < HALF + a = tl.load(o_ptr + in_base + NOPE + 2 * k, mask=kmask).to(tl.float32) + b = tl.load(o_ptr + in_base + NOPE + 2 * k + 1, mask=kmask).to(tl.float32) + cos = tl.load(cos_sin_ptr + pos * cs_stride + k, mask=kmask) + sin = tl.load(cos_sin_ptr + pos * cs_stride + HALF + k, mask=kmask) + out_even = a * cos + b * sin + out_odd = b * cos - a * sin + tl.store(out_ptr + out_base + NOPE + 2 * k, out_even.to(tl.bfloat16), mask=kmask) + tl.store(out_ptr + out_base + NOPE + 2 * k + 1, out_odd.to(tl.bfloat16), mask=kmask) + + +def _can_use_fused_inv_rope( rotary_emb: torch.nn.Module, + o: torch.Tensor, + rope_head_dim: int, +) -> bool: + """Guard: only fuse the exactly-validated GPT-J standard layout.""" + cos_sin_cache = getattr(rotary_emb, "cos_sin_cache", None) + return ( + rope_head_dim > 0 + and rope_head_dim % 2 == 0 + and o.dim() == 3 + and o.stride(-1) == 1 + and not getattr(rotary_emb, "is_neox_style", False) + and getattr(rotary_emb, "rotary_dim", rope_head_dim) == rope_head_dim + and cos_sin_cache is not None + and cos_sin_cache.shape[-1] == rope_head_dim + ) + + +def _fused_inverse_rope_gptj( o: torch.Tensor, positions: torch.Tensor, + cos_sin_cache: torch.Tensor, rope_head_dim: int, +) -> torch.Tensor: + """bf16 inverse GPT-J RoPE via a single fused Triton kernel.""" + num_tokens, num_heads, head_dim = o.shape + out = torch.empty( + (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device + ) + if num_tokens == 0: + return out + _inverse_rope_gptj_kernel[(num_tokens, num_heads)]( + o, + out, + positions, + cos_sin_cache, + o.stride(0), + o.stride(1), + out.stride(0), + out.stride(1), + cos_sin_cache.stride(0), + NOPE=head_dim - rope_head_dim, + HALF=rope_head_dim // 2, + BLOCK_NOPE=triton.next_power_of_2(head_dim - rope_head_dim), + BLOCK_HALF=triton.next_power_of_2(rope_head_dim // 2), + ) + return out + + +def _get_cached_wo_a_bf16( + wo_a: torch.nn.Module, n_local_groups: int, o_lora_rank: int, - wo_a: torch.nn.Module, + hidden_dim: int, ) -> torch.Tensor: - """Reference inverse-RoPE + WO_A einsum path used on ROCm.""" - o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to( - torch.bfloat16 - ) - o_ref = o_ref.view(o.shape[0], n_local_groups, -1) + """Dequantize wo_a to bf16 once and cache it on the module. - hidden_dim = o_ref.shape[-1] + wo_a weights are static, so the fp8 -> fp32 -> (* block scale) -> bf16 + dequant only needs to run once. Recomputing it every decode step shows up + in the profile as the largest copy/mul kernels (``direct_copy float`` ~55us + and ``MulFunctor float`` ~31us per two layers). SGLang / ATOM keep wo_a in + bf16 and feed a plain bf16 GEMM; this mirrors that. + """ + cached = getattr(wo_a, "_dsv4_wo_a_bf16", None) + if cached is not None: + return cached if hasattr(wo_a, "weight_scale_inv"): wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to( torch.float32 @@ -951,11 +1051,42 @@ def rocm_inv_rope_einsum( o_lora_rank, hidden_dim, ) - wo_a_weight = (wo_a_weight * wo_a_scale).to(torch.bfloat16) + cached = (wo_a_weight * wo_a_scale).to(torch.bfloat16) else: - wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to( + cached = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to( torch.bfloat16 ) + wo_a._dsv4_wo_a_bf16 = cached + return cached + + +def rocm_inv_rope_einsum( + rotary_emb: torch.nn.Module, + o: torch.Tensor, + positions: torch.Tensor, + rope_head_dim: int, + n_local_groups: int, + o_lora_rank: int, + wo_a: torch.nn.Module, +) -> torch.Tensor: + """Inverse-RoPE + WO_A bmm path used on ROCm. + + Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16 + wo_a weight so the per-step dequant disappears. + """ + if _can_use_fused_inv_rope(rotary_emb, o, rope_head_dim): + o_ref = _fused_inverse_rope_gptj( + o, positions, rotary_emb.cos_sin_cache, rope_head_dim + ) + else: + o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to( + torch.bfloat16 + ) + o_ref = o_ref.view(o.shape[0], n_local_groups, -1) + + wo_a_weight = _get_cached_wo_a_bf16( + wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1] + ) return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight) From 789938890c8043b796261e04780c008cff5ab6e2 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 12 Jun 2026 02:51:58 +0000 Subject: [PATCH 2/3] add fused o proj unit test Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 300 ++++++++++++++++++ 1 file changed, 300 insertions(+) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index f328f339332e..c6d6ee5fbc16 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -515,3 +515,303 @@ def test_sparse_attn_decode_split_k_kernel( ) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +# --------------------------------------------------------------------------- +# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) +# --------------------------------------------------------------------------- + + +def _make_cos_sin_cache( + max_pos: int, rope_dim: int, device: torch.device +) -> torch.Tensor: + """fp32 cos_sin_cache laid out as [P, rope_dim] = cos[:half] | sin[half:].""" + half = rope_dim // 2 + angles = torch.rand(max_pos, half, dtype=torch.float32, device=device) * 6.2831853 + return torch.cat((angles.cos(), angles.sin()), dim=-1) + + +class _FakeGPTJRotary(torch.nn.Module): + """Minimal stand-in for DeepseekV4ScalingRotaryEmbedding (GPT-J layout). + + No ``forward_native`` so the reference path falls back to the in-module + pure-PyTorch ``_apply_gptj_inv_rope_ref``. + """ + + def __init__(self, cos_sin_cache: torch.Tensor, rotary_dim: int) -> None: + super().__init__() + self.cos_sin_cache = cos_sin_cache + self.is_neox_style = False + self.rotary_dim = rotary_dim + + +class _FakeWoA(torch.nn.Module): + """Stand-in for the wo_a linear layer holding the (optionally fp8) weight.""" + + def __init__( + self, weight: torch.Tensor, weight_scale_inv: torch.Tensor | None = None + ) -> None: + super().__init__() + self.weight = weight + if weight_scale_inv is not None: + self.weight_scale_inv = weight_scale_inv + + +@pytest.mark.parametrize("num_tokens", [1, 7, 64]) +@pytest.mark.parametrize("num_heads", [1, 8]) +@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64]) +@torch.inference_mode() +def test_fused_inverse_rope_gptj_matches_reference( + num_tokens: int, num_heads: int, pos_dtype: torch.dtype +) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_gptj_inv_rope_ref, + _fused_inverse_rope_gptj, + ) + + device = torch.device("cuda") + torch.manual_seed(0) + max_pos = 4096 + cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) + o = torch.randn( + num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + positions = torch.randint(0, max_pos, (num_tokens,), dtype=pos_dtype, device=device) + + actual = _fused_inverse_rope_gptj(o, positions, cos_sin_cache, ROPE_HEAD_DIM) + expected = _apply_gptj_inv_rope_ref(o, positions, cos_sin_cache, ROPE_HEAD_DIM).to( + torch.bfloat16 + ) + + assert actual.dtype == torch.bfloat16 + assert actual.shape == o.shape + # NoPE lanes are a pure bf16 passthrough -> must be bit-exact. + assert torch.equal(actual[..., :NOPE_HEAD_DIM], expected[..., :NOPE_HEAD_DIM]) + # RoPE lanes: tolerate at most ~1 bf16 ulp from fp32 fma ordering. + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@torch.inference_mode() +def test_fused_inverse_rope_gptj_empty() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj + + device = torch.device("cuda") + cos_sin_cache = _make_cos_sin_cache(16, ROPE_HEAD_DIM, device) + o = torch.empty(0, 8, HEAD_DIM, dtype=torch.bfloat16, device=device) + positions = torch.empty(0, dtype=torch.int32, device=device) + + out = _fused_inverse_rope_gptj(o, positions, cos_sin_cache, ROPE_HEAD_DIM) + assert out.shape == (0, 8, HEAD_DIM) + assert out.dtype == torch.bfloat16 + + +@torch.inference_mode() +def test_rocm_inv_rope_einsum_fused_matches_reference() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_gptj_inv_rope_ref, + _can_use_fused_inv_rope, + rocm_inv_rope_einsum, + ) + + device = torch.device("cuda") + torch.manual_seed(2) + num_tokens, num_heads = 5, 8 + n_local_groups = num_heads + o_lora_rank = 16 + hidden_dim = num_heads * HEAD_DIM // n_local_groups # 512 + max_pos = 4096 + + cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) + rotary_emb = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + o = ( + torch.randn( + num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + * 0.125 + ) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int32, device=device + ) + weight = ( + torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125 + ).to(torch.bfloat16) + wo_a = _FakeWoA(weight) + + # Sanity: this configuration must actually take the fused path. + assert _can_use_fused_inv_rope(rotary_emb, o, ROPE_HEAD_DIM) + + actual = rocm_inv_rope_einsum( + rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a + ) + + o_ref = _apply_gptj_inv_rope_ref(o, positions, cos_sin_cache, ROPE_HEAD_DIM).to( + torch.bfloat16 + ) + o_ref = o_ref.view(num_tokens, n_local_groups, -1) + wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16) + expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref) + + assert actual.shape == (num_tokens, n_local_groups, o_lora_rank) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@torch.inference_mode() +def test_rocm_inv_rope_einsum_fallback_matches_reference() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_inv_rope_ref, + _can_use_fused_inv_rope, + rocm_inv_rope_einsum, + ) + + device = torch.device("cuda") + torch.manual_seed(3) + num_tokens, num_heads = 4, 8 + n_local_groups = num_heads + o_lora_rank = 16 + hidden_dim = num_heads * HEAD_DIM // n_local_groups + max_pos = 4096 + + cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) + rotary_emb = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + # Non-contiguous last dim disqualifies the fused kernel -> PyTorch fallback. + o = ( + torch.randn( + num_tokens, HEAD_DIM, num_heads, dtype=torch.bfloat16, device=device + ) + * 0.125 + ).transpose(1, 2) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int32, device=device + ) + weight = ( + torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125 + ).to(torch.bfloat16) + wo_a = _FakeWoA(weight) + + assert not _can_use_fused_inv_rope(rotary_emb, o, ROPE_HEAD_DIM) + + actual = rocm_inv_rope_einsum( + rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a + ) + + o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, ROPE_HEAD_DIM).to( + torch.bfloat16 + ) + o_ref = o_ref.view(num_tokens, n_local_groups, -1) + wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16) + expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@torch.inference_mode() +def test_get_cached_wo_a_bf16_plain_caches() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 + + device = torch.device("cuda") + torch.manual_seed(4) + n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8 + weight = torch.randn( + n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device + ) + wo_a = _FakeWoA(weight) + + out1 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) + expected = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16) + assert out1.shape == (n_local_groups, o_lora_rank, hidden_dim) + torch.testing.assert_close(out1, expected, atol=0, rtol=0) + assert hasattr(wo_a, "_dsv4_wo_a_bf16") + + # Mutate the source weight: the cached tensor must be returned unchanged + # (proving the dequant is not recomputed per call). + wo_a.weight.zero_() + out2 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) + assert out2 is out1 + torch.testing.assert_close(out2, expected, atol=0, rtol=0) + + +@torch.inference_mode() +def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 + + device = torch.device("cuda") + torch.manual_seed(5) + n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8 + row_block, col_block = 2, 2 + row_blocks = o_lora_rank // row_block + col_blocks = hidden_dim // col_block + + fp8_dtype = current_platform.fp8_dtype() + weight_f32 = ( + torch.randn( + n_local_groups, o_lora_rank, hidden_dim, dtype=torch.float32, device=device + ) + * 0.1 + ) + weight_fp8 = weight_f32.to(fp8_dtype) + scale = ( + torch.rand( + n_local_groups, row_blocks, col_blocks, dtype=torch.float32, device=device + ) + * 0.5 + + 0.5 + ) + wo_a = _FakeWoA( + weight_fp8.reshape(n_local_groups * o_lora_rank, hidden_dim), + weight_scale_inv=scale.reshape(n_local_groups * row_blocks, col_blocks), + ) + + out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) + + scale_full = scale.repeat_interleave(row_block, dim=-2).repeat_interleave( + col_block, dim=-1 + ) + expected = (weight_fp8.to(torch.float32) * scale_full).to(torch.bfloat16) + assert out.shape == (n_local_groups, o_lora_rank, hidden_dim) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + # Second call returns the same cached object. + assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out + + +@torch.inference_mode() +def test_can_use_fused_inv_rope_guard() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _can_use_fused_inv_rope + + device = torch.device("cuda") + cos_sin_cache = _make_cos_sin_cache(32, ROPE_HEAD_DIM, device) + o = torch.randn(4, 8, HEAD_DIM, dtype=torch.bfloat16, device=device) + + good = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + assert _can_use_fused_inv_rope(good, o, ROPE_HEAD_DIM) + + # neox-style layout is not supported by the fused kernel. + neox = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + neox.is_neox_style = True + assert not _can_use_fused_inv_rope(neox, o, ROPE_HEAD_DIM) + + # rotary_dim must equal rope_head_dim (only the trailing dims are rotated). + mismatched = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM + 2) + assert not _can_use_fused_inv_rope(mismatched, o, ROPE_HEAD_DIM) + + # Missing cos_sin_cache disables fusion. + no_cache = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + no_cache.cos_sin_cache = None + assert not _can_use_fused_inv_rope(no_cache, o, ROPE_HEAD_DIM) + + # cos_sin_cache width must match rope_head_dim. + bad_cache = _FakeGPTJRotary( + _make_cos_sin_cache(32, ROPE_HEAD_DIM + 2, device), ROPE_HEAD_DIM + ) + assert not _can_use_fused_inv_rope(bad_cache, o, ROPE_HEAD_DIM) + + # Odd rope_head_dim is rejected. + odd_cache = _FakeGPTJRotary( + _make_cos_sin_cache(32, ROPE_HEAD_DIM + 1, device), ROPE_HEAD_DIM + 1 + ) + assert not _can_use_fused_inv_rope(odd_cache, o, ROPE_HEAD_DIM + 1) + + # Non-3D / non-contiguous-last-dim inputs are rejected. + assert not _can_use_fused_inv_rope(good, o.reshape(4, 8 * HEAD_DIM), ROPE_HEAD_DIM) + o_t = torch.randn(4, HEAD_DIM, 8, dtype=torch.bfloat16, device=device).transpose( + 1, 2 + ) + assert not _can_use_fused_inv_rope(good, o_t, ROPE_HEAD_DIM) From e10bf23d53f253c90b42e75b9d66a9993471183b Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 12 Jun 2026 02:51:58 +0000 Subject: [PATCH 3/3] Address review: make fused inverse-RoPE the only path, test against official rotary - Remove _can_use_fused_inv_rope and the PyTorch fallback: the fused Triton kernel is the only path for rocm_inv_rope_einsum (DSv4-only, guard was always true). Replace the runtime guard with asserts on the kernel's layout contract. - Drop the now-unused _apply_inv_rope_ref / _apply_gptj_inv_rope_ref reference helpers. - Unit tests now validate against the official DeepseekV4ScalingRotaryEmbedding.forward_native(inverse=True) instead of a hand-rolled fake, so semantic changes to the rotary class are caught. Co-Authored-By: Claude Fable 5 Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 199 +++++------------- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 90 ++------ 2 files changed, 70 insertions(+), 219 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index c6d6ee5fbc16..daf73b82e614 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -515,32 +515,52 @@ def test_sparse_attn_decode_split_k_kernel( ) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- -def _make_cos_sin_cache( - max_pos: int, rope_dim: int, device: torch.device -) -> torch.Tensor: - """fp32 cos_sin_cache laid out as [P, rope_dim] = cos[:half] | sin[half:].""" - half = rope_dim // 2 - angles = torch.rand(max_pos, half, dtype=torch.float32, device=device) * 6.2831853 - return torch.cat((angles.cos(), angles.sin()), dim=-1) +# Cache rows = max_position_embeddings * scaling_factor. +_ROTARY_MAX_POS = 1024 +_ROTARY_SCALING_FACTOR = 4.0 +_ROTARY_CACHE_LEN = int(_ROTARY_MAX_POS * _ROTARY_SCALING_FACTOR) -class _FakeGPTJRotary(torch.nn.Module): - """Minimal stand-in for DeepseekV4ScalingRotaryEmbedding (GPT-J layout). +def _make_dsv4_rotary(device: torch.device): + """The official DSv4 rotary embedding, sized down for unit tests.""" + from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + DeepseekV4ScalingRotaryEmbedding, + ) - No ``forward_native`` so the reference path falls back to the in-module - pure-PyTorch ``_apply_gptj_inv_rope_ref``. - """ + # The model loader constructs layers under a default-device context; + # mirror that so the fp32 cos_sin_cache lands on the GPU. + with torch.device(device): + rotary_emb = DeepseekV4ScalingRotaryEmbedding( + head_size=ROPE_HEAD_DIM, + rotary_dim=ROPE_HEAD_DIM, + max_position_embeddings=_ROTARY_MAX_POS, + base=10000, + is_neox_style=False, + scaling_factor=_ROTARY_SCALING_FACTOR, + dtype=torch.bfloat16, + mscale=1.0, + mscale_all_dim=1.0, + ) + rotary_emb = rotary_emb.to(device) + assert rotary_emb.cos_sin_cache.shape == (_ROTARY_CACHE_LEN, ROPE_HEAD_DIM) + return rotary_emb - def __init__(self, cos_sin_cache: torch.Tensor, rotary_dim: int) -> None: - super().__init__() - self.cos_sin_cache = cos_sin_cache - self.is_neox_style = False - self.rotary_dim = rotary_dim + +def _inv_rope_via_rotary_native( + rotary_emb: torch.nn.Module, + o: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """Reference: the official ``forward_native(inverse=True)`` path.""" + expected, _ = rotary_emb.forward_native(positions, o.clone(), None, inverse=True) + return expected.to(torch.bfloat16) class _FakeWoA(torch.nn.Module): @@ -559,27 +579,25 @@ def __init__( @pytest.mark.parametrize("num_heads", [1, 8]) @pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64]) @torch.inference_mode() -def test_fused_inverse_rope_gptj_matches_reference( - num_tokens: int, num_heads: int, pos_dtype: torch.dtype +def test_fused_inverse_rope_gptj_matches_rotary_native( + num_tokens: int, num_heads: int, pos_dtype: torch.dtype, default_vllm_config ) -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _apply_gptj_inv_rope_ref, - _fused_inverse_rope_gptj, - ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj device = torch.device("cuda") torch.manual_seed(0) - max_pos = 4096 - cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) + rotary_emb = _make_dsv4_rotary(device) o = torch.randn( num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device ) - positions = torch.randint(0, max_pos, (num_tokens,), dtype=pos_dtype, device=device) + positions = torch.randint( + 0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=pos_dtype, device=device + ) - actual = _fused_inverse_rope_gptj(o, positions, cos_sin_cache, ROPE_HEAD_DIM) - expected = _apply_gptj_inv_rope_ref(o, positions, cos_sin_cache, ROPE_HEAD_DIM).to( - torch.bfloat16 + actual = _fused_inverse_rope_gptj( + o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM ) + expected = _inv_rope_via_rotary_native(rotary_emb, o, positions) assert actual.dtype == torch.bfloat16 assert actual.shape == o.shape @@ -590,26 +608,24 @@ def test_fused_inverse_rope_gptj_matches_reference( @torch.inference_mode() -def test_fused_inverse_rope_gptj_empty() -> None: +def test_fused_inverse_rope_gptj_empty(default_vllm_config) -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj device = torch.device("cuda") - cos_sin_cache = _make_cos_sin_cache(16, ROPE_HEAD_DIM, device) + rotary_emb = _make_dsv4_rotary(device) o = torch.empty(0, 8, HEAD_DIM, dtype=torch.bfloat16, device=device) positions = torch.empty(0, dtype=torch.int32, device=device) - out = _fused_inverse_rope_gptj(o, positions, cos_sin_cache, ROPE_HEAD_DIM) + out = _fused_inverse_rope_gptj( + o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM + ) assert out.shape == (0, 8, HEAD_DIM) assert out.dtype == torch.bfloat16 @torch.inference_mode() -def test_rocm_inv_rope_einsum_fused_matches_reference() -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _apply_gptj_inv_rope_ref, - _can_use_fused_inv_rope, - rocm_inv_rope_einsum, - ) +def test_rocm_inv_rope_einsum_matches_rotary_native(default_vllm_config) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum device = torch.device("cuda") torch.manual_seed(2) @@ -617,10 +633,8 @@ def test_rocm_inv_rope_einsum_fused_matches_reference() -> None: n_local_groups = num_heads o_lora_rank = 16 hidden_dim = num_heads * HEAD_DIM // n_local_groups # 512 - max_pos = 4096 - cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) - rotary_emb = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) + rotary_emb = _make_dsv4_rotary(device) o = ( torch.randn( num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device @@ -628,23 +642,18 @@ def test_rocm_inv_rope_einsum_fused_matches_reference() -> None: * 0.125 ) positions = torch.randint( - 0, max_pos, (num_tokens,), dtype=torch.int32, device=device + 0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=torch.int32, device=device ) weight = ( torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125 ).to(torch.bfloat16) wo_a = _FakeWoA(weight) - # Sanity: this configuration must actually take the fused path. - assert _can_use_fused_inv_rope(rotary_emb, o, ROPE_HEAD_DIM) - actual = rocm_inv_rope_einsum( rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a ) - o_ref = _apply_gptj_inv_rope_ref(o, positions, cos_sin_cache, ROPE_HEAD_DIM).to( - torch.bfloat16 - ) + o_ref = _inv_rope_via_rotary_native(rotary_emb, o, positions) o_ref = o_ref.view(num_tokens, n_local_groups, -1) wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16) expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref) @@ -653,55 +662,6 @@ def test_rocm_inv_rope_einsum_fused_matches_reference() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) -@torch.inference_mode() -def test_rocm_inv_rope_einsum_fallback_matches_reference() -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _apply_inv_rope_ref, - _can_use_fused_inv_rope, - rocm_inv_rope_einsum, - ) - - device = torch.device("cuda") - torch.manual_seed(3) - num_tokens, num_heads = 4, 8 - n_local_groups = num_heads - o_lora_rank = 16 - hidden_dim = num_heads * HEAD_DIM // n_local_groups - max_pos = 4096 - - cos_sin_cache = _make_cos_sin_cache(max_pos, ROPE_HEAD_DIM, device) - rotary_emb = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) - # Non-contiguous last dim disqualifies the fused kernel -> PyTorch fallback. - o = ( - torch.randn( - num_tokens, HEAD_DIM, num_heads, dtype=torch.bfloat16, device=device - ) - * 0.125 - ).transpose(1, 2) - positions = torch.randint( - 0, max_pos, (num_tokens,), dtype=torch.int32, device=device - ) - weight = ( - torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125 - ).to(torch.bfloat16) - wo_a = _FakeWoA(weight) - - assert not _can_use_fused_inv_rope(rotary_emb, o, ROPE_HEAD_DIM) - - actual = rocm_inv_rope_einsum( - rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a - ) - - o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, ROPE_HEAD_DIM).to( - torch.bfloat16 - ) - o_ref = o_ref.view(num_tokens, n_local_groups, -1) - wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16) - expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref) - - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) - - @torch.inference_mode() def test_get_cached_wo_a_bf16_plain_caches() -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 @@ -770,48 +730,3 @@ def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None: # Second call returns the same cached object. assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out - - -@torch.inference_mode() -def test_can_use_fused_inv_rope_guard() -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _can_use_fused_inv_rope - - device = torch.device("cuda") - cos_sin_cache = _make_cos_sin_cache(32, ROPE_HEAD_DIM, device) - o = torch.randn(4, 8, HEAD_DIM, dtype=torch.bfloat16, device=device) - - good = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) - assert _can_use_fused_inv_rope(good, o, ROPE_HEAD_DIM) - - # neox-style layout is not supported by the fused kernel. - neox = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) - neox.is_neox_style = True - assert not _can_use_fused_inv_rope(neox, o, ROPE_HEAD_DIM) - - # rotary_dim must equal rope_head_dim (only the trailing dims are rotated). - mismatched = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM + 2) - assert not _can_use_fused_inv_rope(mismatched, o, ROPE_HEAD_DIM) - - # Missing cos_sin_cache disables fusion. - no_cache = _FakeGPTJRotary(cos_sin_cache, ROPE_HEAD_DIM) - no_cache.cos_sin_cache = None - assert not _can_use_fused_inv_rope(no_cache, o, ROPE_HEAD_DIM) - - # cos_sin_cache width must match rope_head_dim. - bad_cache = _FakeGPTJRotary( - _make_cos_sin_cache(32, ROPE_HEAD_DIM + 2, device), ROPE_HEAD_DIM - ) - assert not _can_use_fused_inv_rope(bad_cache, o, ROPE_HEAD_DIM) - - # Odd rope_head_dim is rejected. - odd_cache = _FakeGPTJRotary( - _make_cos_sin_cache(32, ROPE_HEAD_DIM + 1, device), ROPE_HEAD_DIM + 1 - ) - assert not _can_use_fused_inv_rope(odd_cache, o, ROPE_HEAD_DIM + 1) - - # Non-3D / non-contiguous-last-dim inputs are rejected. - assert not _can_use_fused_inv_rope(good, o.reshape(4, 8 * HEAD_DIM), ROPE_HEAD_DIM) - o_t = torch.randn(4, HEAD_DIM, 8, dtype=torch.bfloat16, device=device).transpose( - 1, 2 - ) - assert not _can_use_fused_inv_rope(good, o_t, ROPE_HEAD_DIM) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 9106e4906e6e..c38a4780f784 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -874,56 +874,6 @@ def _expand_2d_block_scales( return scale -def _apply_gptj_inv_rope_ref( - x: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - rope_dim: int, -) -> torch.Tensor: - if rope_dim == 0 or x.numel() == 0: - return x - half_rot = rope_dim // 2 - nope_dim = x.shape[-1] - rope_dim - dtype = x.dtype - x = x.to(torch.float32) - cache = cos_sin_cache.index_select(0, positions.to(torch.long)) - cos = cache[:, :half_rot].to(torch.float32) - sin = cache[:, half_rot : 2 * half_rot].to(torch.float32) - view_shape = (positions.shape[0],) + (1,) * (x.dim() - 2) + (half_rot,) - cos = cos.view(view_shape) - sin = sin.view(view_shape) - rope = x[..., nope_dim:] - y_even = rope[..., 0::2] - y_odd = rope[..., 1::2] - rope_out = torch.stack( - (y_even * cos + y_odd * sin, y_odd * cos - y_even * sin), - dim=-1, - ).flatten(-2) - x = x.clone() - x[..., nope_dim:] = rope_out - return x.to(dtype) - - -def _apply_inv_rope_ref( - rotary_emb: torch.nn.Module, - x: torch.Tensor, - positions: torch.Tensor, - rope_dim: int, -) -> torch.Tensor: - if hasattr(rotary_emb, "forward_native"): - try: - query, _ = rotary_emb.forward_native( - positions, - x.clone(), - None, - inverse=True, - ) - return query - except TypeError: - pass - return _apply_gptj_inv_rope_ref(x, positions, rotary_emb.cos_sin_cache, rope_dim) - - @triton.jit def _inverse_rope_gptj_kernel( o_ptr, # [T, H, D] input @@ -973,25 +923,6 @@ def _inverse_rope_gptj_kernel( tl.store(out_ptr + out_base + NOPE + 2 * k + 1, out_odd.to(tl.bfloat16), mask=kmask) -def _can_use_fused_inv_rope( - rotary_emb: torch.nn.Module, - o: torch.Tensor, - rope_head_dim: int, -) -> bool: - """Guard: only fuse the exactly-validated GPT-J standard layout.""" - cos_sin_cache = getattr(rotary_emb, "cos_sin_cache", None) - return ( - rope_head_dim > 0 - and rope_head_dim % 2 == 0 - and o.dim() == 3 - and o.stride(-1) == 1 - and not getattr(rotary_emb, "is_neox_style", False) - and getattr(rotary_emb, "rotary_dim", rope_head_dim) == rope_head_dim - and cos_sin_cache is not None - and cos_sin_cache.shape[-1] == rope_head_dim - ) - - def _fused_inverse_rope_gptj( o: torch.Tensor, positions: torch.Tensor, @@ -999,6 +930,16 @@ def _fused_inverse_rope_gptj( rope_head_dim: int, ) -> torch.Tensor: """bf16 inverse GPT-J RoPE via a single fused Triton kernel.""" + assert o.dim() == 3 and o.stride(-1) == 1, ( + "_fused_inverse_rope_gptj expects a [T, H, D] input with a contiguous last dim" + ) + assert rope_head_dim > 0 and rope_head_dim % 2 == 0, ( + f"_fused_inverse_rope_gptj expects an even rope_head_dim, got {rope_head_dim}" + ) + assert cos_sin_cache.shape[-1] == rope_head_dim, ( + "_fused_inverse_rope_gptj expects cos_sin_cache laid out as " + f"[P, {rope_head_dim}] = cos | sin, got {tuple(cos_sin_cache.shape)}" + ) num_tokens, num_heads, head_dim = o.shape out = torch.empty( (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device @@ -1074,14 +1015,9 @@ def rocm_inv_rope_einsum( Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16 wo_a weight so the per-step dequant disappears. """ - if _can_use_fused_inv_rope(rotary_emb, o, rope_head_dim): - o_ref = _fused_inverse_rope_gptj( - o, positions, rotary_emb.cos_sin_cache, rope_head_dim - ) - else: - o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to( - torch.bfloat16 - ) + o_ref = _fused_inverse_rope_gptj( + o, positions, rotary_emb.cos_sin_cache, rope_head_dim + ) o_ref = o_ref.view(o.shape[0], n_local_groups, -1) wo_a_weight = _get_cached_wo_a_bf16(