[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>
This commit is contained in:
Fangzhou Ai
co-authored by
Claude Fable 5
Claude Fable 5
parent
c90650088d
commit
badddd254f
@@ -515,3 +515,218 @@ 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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# 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 _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):
|
||||
"""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_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 _fused_inverse_rope_gptj
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
rotary_emb = _make_dsv4_rotary(device)
|
||||
o = torch.randn(
|
||||
num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
positions = torch.randint(
|
||||
0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=pos_dtype, device=device
|
||||
)
|
||||
|
||||
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
|
||||
# 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(default_vllm_config) -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
|
||||
|
||||
device = torch.device("cuda")
|
||||
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, 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_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)
|
||||
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
|
||||
|
||||
rotary_emb = _make_dsv4_rotary(device)
|
||||
o = (
|
||||
torch.randn(
|
||||
num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
* 0.125
|
||||
)
|
||||
positions = torch.randint(
|
||||
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)
|
||||
|
||||
actual = rocm_inv_rope_einsum(
|
||||
rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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_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
|
||||
|
||||
@@ -874,72 +874,113 @@ 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)
|
||||
@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 _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)
|
||||
|
||||
|
||||
def rocm_inv_rope_einsum(
|
||||
rotary_emb: torch.nn.Module,
|
||||
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."""
|
||||
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
|
||||
)
|
||||
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 +992,37 @@ 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.
|
||||
"""
|
||||
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(
|
||||
wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1]
|
||||
)
|
||||
|
||||
return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user