I'm unsure if registering with the same IR ops is the best approach.

PyTorch custom op registration does not support None (i.e., Optional[Tensor]) in return types. So when key is None, I fall it to static implement.

(EngineCore_DP0 pid=115375)   File "/home/wangxiaoshuang/vllm/.venv/lib/python3.11/site-packages/torch/_library/infer_schema.py", line 71, in error_fn
(EngineCore_DP0 pid=115375)     raise ValueError(f"infer_schema(func): {what} Got func with signature {sig})")
(EngineCore_DP0 pid=115375) ValueError: infer_schema(func): Return has unsupported type tuple[torch.Tensor, torch.Tensor | None]. The valid types are: {<class 'torch.Tensor'>: 'Tensor', typing.List[torch.Tensor]: 'Tensor[]', list[torch.Tensor]: 'Tensor[]', <class 'int'>: 'SymInt', <class 'float'>: 'float', <class 'bool'>: 'bool', int | float | bool: 'Scalar'}. Got func with signature (positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor | None, head_size: int, cos_sin_cache: torch.Tensor, is_neox_style: bool) -> tuple[torch.Tensor, torch.Tensor | None])
[rank0]:[W318 06:36:34.959129788 ProcessGroupNCCL.cpp:1553] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

The standard format implementation for XPU kernels does not check if offsets is None. If offsets is provided with the standard format, it will be silently ignored by the underlying torch.ops._C.rotary_embedding call, leading to incorrect results. Please add a check for offsets is None in the supports_args or within the implementation.

The rotary_no_offsets_16bit_only check does not verify that rotary_dim == head_size. The rocm_aiter_triton_rotary_embedding kernel is called with head_size but without rotary_dim, which suggests it might not support partial rotation (where rotary_dim < head_size). If this is the case, the implementation should only be selected when rotary_dim == head_size to avoid incorrect results for models with partial RoPE.

The rotary_no_offsets check should verify that rotary_dim == head_size if the underlying torch.ops._C.rotary_embedding kernel does not support partial rotation. Since the kernel is called with head_size as the 4th argument and no rotary_dim is provided, it is likely to rotate the entire head. This would be incorrect for models where rotary_dim < head_size.