Add fused RoPE kernel

[amd-weisun]

Motivation

• FlyDSL implementation of fused Rotary Position Embedding + KV cache write, (AITER's Triton fused_qk_rope_reshape_and_cache kernel), used in GPT-OSS and some other models(Qwen3, Llama-3.1)

• Supports both flash and non-flash (ATOM production default, x-packed) KV cache layouts

• Computes rotation in native bf16 precision matching AITER/Triton — cross-validation (0 error)

Test Plan

Usage:

# Fast CI — correctness only (GPT-OSS 120B TP=8, 10 tests):
PYTHONPATH=./ pytest tests/kernels/test_fused_rope_cache.py -v -s

# All models × TPs (multi-model sweep):
FLYDSL_ALL_MODELS=1 PYTHONPATH=./ pytest tests/kernels/test_fused_rope_cache.py -v -s

# With benchmarking + optional AITER comparison:
FLYDSL_BENCH=1 AITER_REPO=../aiter PYTHONPATH=./ pytest tests/kernels/test_fused_rope_cache.py -v -s

# CLI — all models:
PYTHONPATH=./ python tests/kernels/test_fused_rope_cache.py --all-models

# CLI — with benchmark + AITER comparison:
FLYDSL_BENCH=1 AITER_REPO=../aiter PYTHONPATH=./ python tests/kernels/test_fused_rope_cache.py --all-models

Test Result

• Tested on MI350: 0 numerical error vs PyTorch reference. Performance: 1.4-1.8x faster than Triton (AITER) across all configs ( GPT-OSS-120B, Qwen3, Llama-3.1), both layouts verified. Cross-validated against AITER output.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Adds a new FlyDSL fused RoPE rotation + KV-cache write kernel (supporting both flash and ATOM non-flash cache layouts) along with a correctness/perf test harness and optional AITER cross-validation to align with GPT-OSS/ATOM usage.

Changes:

• Introduce build_fused_rope_cache_module() that emits two GPU kernels: Q RoPE, then K RoPE + KV-cache write.
• Add correctness tests for both cache layouts, plus an optional multi-model sweep and optional AITER perf/cross-check path.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 10 comments.

File	Description
kernels/fused_rope_cache_kernel.py	New fused RoPE + KV-cache kernel builder using @flyc.kernel + @flyc.jit with flash and non-flash cache layout support.
tests/kernels/test_fused_rope_cache.py	New GPU test validating Q/K outputs and KV-cache writes against a PyTorch reference, with optional AITER comparison.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.