[ROCm][CI] add Aiter ops tests - #52208
Conversation
Add tests for ROCm aiter operations including rms_norm, rms_norm2d_with_add, and various fused quantization ops. Only adds two new methods to rocm_aiter_ops class: - rms_norm: wraps aiter.rmsnorm2d_fwd - rms_norm2d_with_add: wraps aiter.rmsnorm2d_fwd_with_add Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
AndreasKaratzas
left a comment
AndreasKaratzas
left a comment
There was a problem hiding this comment.
Overall LGTM, but let's address these few comments carefully.
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
|
/ci run |
AndreasKaratzas
added
the
ready
|
✅ Triggered Buildkite CI #83971 for commit |
|
✅ @divakar-amd, CI is now available for this PR.
|
| cos_pos = cos_half[positions] # [num_tokens, half_dim] | ||
| sin_pos = sin_half[positions] # [num_tokens, half_dim] | ||
|
|
||
| def apply_rope_ref(t: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
t.float() upcasts the operands to float32 before doing the unit RoPE op, which does not match aiter triton (keeps native dtype): https://github.com/ROCm/aiter/blob/main/aiter/ops/triton/_triton_kernels/rope/rope.py#L858
This has been a point of divergence before, because even vLLM's C++ and Triton RoPE impls do not fully agree on upcasting behavior for RoPE.
vLLM C++ always upcasts to fp32: https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/pos_encoding_kernels.cu#L18
vLLM Triton keeps native dtype by default unless enable_fp32_compute is set: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding/common.py#L151
There was a problem hiding this comment.
Thanks for sharing this @Rohan138. Opened a new PR to handle the fix.
Signed-off-by: Divakar Verma <divakar.verma@amd.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
3 participants
Add ROCm AITER ops tests and enable
Kernels Core Operation Testfor the MI355 mirrorTests cover:
New ops added to rocm_aiter_ops: