[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>

vllm/model_executor/layers/rotary_embedding/common.py

@@ -253,9 +253,32 @@ class ApplyRotaryEmb(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
+        _HIP_MAX_GRID_DIM = 65535
+        """
+        HIP/ROCm has a per-dim grid limit of 65535 on gridY/gridZ. The
+        flash_attn triton rotary kernel uses
+        grid = (cdiv(nheads, BLOCK_H), cdiv(seq_len, BLOCK_M), batch)
+        with BLOCK_M=8 (rotary_dim<=128) or BLOCK_M=4 (otherwise) and
+        BLOCK_H=2. When the visual encoder packs many image patches into one
+        batch (e.g. vLLM profile_run with max_num_seqs images), gridY can
+        exceed 65535 and hipModuleLaunchKernel returns
+        `Triton Error [HIP]: Code: 1, invalid argument`. Fall back to the
+        native PyTorch implementation in that case.
+        """
         if self.apply_rotary_emb_flash_attn is not None:
             x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
 
+            seq_len = x.shape[-3]
+            batch = x.shape[0]
+            rotary_dim = cos.shape[-1] * 2
+            block_m = 8 if rotary_dim <= 128 else 4
+            grid_y = (seq_len + block_m - 1) // block_m
+            if grid_y > _HIP_MAX_GRID_DIM or batch > _HIP_MAX_GRID_DIM:
+                output = self.forward_static(
+                    x, cos, sin, self.is_neox_style, self.enable_fp32_compute
+                )
+                return self._post_process(output, origin_shape, origin_dtype)
+
             """
             Arguments of apply_rotary() in flash_attn:
                 x: [batch_size, seq_len, nheads, headdim]