[ROCm][DSv4][Perf] Fuse the inverse RoPE into the sparse decode reduce (#57451)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Fangzhou Ai committed
Sep 22, 2026 at 21:23 UTC
f9dce295c96fff3abfbf4d8fb34a34ef8c2c70b6
1 file changed
+27
-3
vllm/models/deepseek_v4/amd/rocm.py
+27
-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
)
@@ -1175,6 +1176,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1176
self.n_local_groups,
1177
self.o_lora_rank,
1178
self.wo_a,
1179
+ inverse_rope=False,
1180
)
1181
zf = z.flatten(1)
1182
if self._wo_b_scale is not None and zf.dim() == 2:
@@ -1245,9 +1247,15 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1247
attn_metadata=rocm_metadata,
1248
swa_metadata=swa_metadata,
1249
)
1250
+ # The fp8 wo_a path rotates inside inverse_rope_group_quant, so folding
1251
+ # the rotation into the decode reduce would apply it twice. Only the
1252
+ # BF16 einsum path hands its rotation off to the decode.
1253
+ fuse_inv_rope = self._wo_a_fp8_weight is None
1254
+ rotated = 0
1255
if num_decodes > 0:
1249
- self._forward_decode(
1256
+ rotated = self._forward_decode(
1257
q=q[:num_decode_tokens],
1258
+ positions=positions[:num_decode_tokens] if fuse_inv_rope else None,
1259
kv_cache=self_kv_cache,
1260
swa_metadata=swa_metadata,
1261
attn_metadata=rocm_metadata,
@@ -1261,17 +1269,31 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1269
and rocm_metadata.for_cudagraph_capture
1270
),
1271
)
1272
+ if fuse_inv_rope:
1273
+ # Only the decode reduce rotates its own rows, and only the leading
1274
+ # `rotated` of them; prefill rows and any decode path that did not
1275
+ # fuse still owe the standalone pass. Settle that here rather than
1276
+ # in _o_proj: the split is batch-dependent and _o_proj runs
1277
+ # compiled, where such a value freezes at its trace-time value.
1278
+ rocm_inverse_rope_rows_(
1279
+ output[rotated:, : self.n_local_heads, :],
1280
+ positions[rotated:],
1281
+ self.rotary_emb.cos_sin_cache,
1282
+ self.rope_head_dim,
1283
+ )
1284
1285
def _forward_decode(
1286
self,
1287
q: torch.Tensor,
1288
+ positions: torch.Tensor | None,
1289
kv_cache: torch.Tensor | None,
1290
swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
1291
attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
1292
swa_only: bool,
1293
output: torch.Tensor,
1294
adaptive_splits: bool,
1274
- ) -> None:
1295
+ ) -> int:
1296
+ """Returns how many leading rows the decode epilogue inverse-RoPE'd."""
1297
num_decodes = swa_metadata.num_decodes
1298
num_decode_tokens = swa_metadata.num_decode_tokens
1299
@@ -1303,7 +1325,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1325
topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
1326
topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr
1327
1306
- rocm_sparse_attn_decode(
1328
+ return rocm_sparse_attn_decode(
1329
q=q,
1330
kv_cache=kv_cache,
1331
swa_k_cache=self.swa_cache_layer.kv_cache,
@@ -1323,6 +1345,8 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1345
rope_head_dim=self.rope_head_dim,
1346
output=output,
1347
adaptive_splits=adaptive_splits,
1348
+ inv_rope_positions=positions,
1349
+ inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
1350
extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
1351
self.kv_cache_dtype,
1352
self._has_kv_transfer,