A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 18:45 UTC
[Bugfix][ROCm] `ApplyRotaryEmb`: fall back to native when flash_attn rotary grid would exceed the HIP per-dim limit (#43684)
Signed-off-by: vLLM ROCm fix <noreply@example.com> Signed-off-by: amd-fuweiy <fuweiy@amd.com> Co-authored-by: vLLM ROCm fix <noreply@example.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Fuwei Yang committed
Jun 6, 2026 at 17:29 UTC
00d1fb77475567aef7d019881612e7ad564bb132
1 file changed
+23
vllm/model_executor/layers/rotary_embedding/common.py
+23
@@ -253,9 +253,32 @@ class ApplyRotaryEmb(CustomOp):
253
cos: torch.Tensor,
254
sin: torch.Tensor,
255
) -> torch.Tensor:
256
+ _HIP_MAX_GRID_DIM = 65535
257
+ """
258
+ HIP/ROCm has a per-dim grid limit of 65535 on gridY/gridZ. The
259
+ flash_attn triton rotary kernel uses
260
+ grid = (cdiv(nheads, BLOCK_H), cdiv(seq_len, BLOCK_M), batch)
261
+ with BLOCK_M=8 (rotary_dim<=128) or BLOCK_M=4 (otherwise) and
262
+ BLOCK_H=2. When the visual encoder packs many image patches into one
263
+ batch (e.g. vLLM profile_run with max_num_seqs images), gridY can
264
+ exceed 65535 and hipModuleLaunchKernel returns
265
+ `Triton Error [HIP]: Code: 1, invalid argument`. Fall back to the
266
+ native PyTorch implementation in that case.
267
+ """
268
if self.apply_rotary_emb_flash_attn is not None:
269
x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
270
271
+ seq_len = x.shape[-3]
272
+ batch = x.shape[0]
273
+ rotary_dim = cos.shape[-1] * 2
274
+ block_m = 8 if rotary_dim <= 128 else 4
275
+ grid_y = (seq_len + block_m - 1) // block_m
276
+ if grid_y > _HIP_MAX_GRID_DIM or batch > _HIP_MAX_GRID_DIM:
277
+ output = self.forward_static(
278
+ x, cos, sin, self.is_neox_style, self.enable_fp32_compute
279
+ )
280
+ return self._post_process(output, origin_shape, origin_dtype)
281
+
282
"""
283
Arguments of apply_rotary() in flash_attn:
284
x: [batch_size, seq_len, nheads, headdim]