fix(test): drop unsupported device kwarg from config-based RoPE con… (#1514)

## Summary

Fixes the cuTile / cuTeDSL / correctness RoPE CI jobs that started
failing across both H100 and B200 (e.g. run 37538343198) with:

```
TypeError: LlamaRotaryEmbedding.__init__() got an unexpected keyword argument 'device'
```

## Root cause

Newer `transformers` (observed on 5.19.0 in CI) removed the `device`
keyword from the **config-based** `*RotaryEmbedding.__init__`
(`__init__(self, config)`). Our `transformers_version_dispatch(...)`
helper still forwarded `device` in `after_kwargs`, so every RoPE test
that builds a rotary embedding raised the `TypeError` **at setup**,
before any kernel ran — which is why all cuTile + cuTeDSL jobs failed on
both architectures at once. The on-pod/mirror `transformers` (<= 5.15.0)
still accepted `device`, so it only surfaced in CI where public PyPI
resolves a newer version.

## Fix

The config-based constructor derives device from the input tensor in
`forward` (`inv_freq` is moved to `x.device`), so passing `device` at
construction is redundant. This PR drops `"device": device` from every
`after_kwargs`, while keeping it in `before_kwargs` for the pre-4.48.0
positional-`dim` path.

Changed call sites:
- `test/transformers/test_rope.py` (3)
- `test/cutile/test_rope.py` (3)
- `test/cutedsl/test_rope.py` (3)
- `benchmark/scripts/benchmark_rope.py` (1)
- `benchmark/scripts/benchmark_llama4_rope.py` (1)
- `src/liger_kernel/utils.py` — docstring example updated

## Verification

- `test/transformers/test_rope.py`: **81 passed** on transformers 5.15.0
(backward-compatible).
- Simulated the 5.19.0 strict signature (reject `device`): the **old**
pattern raises the exact CI `TypeError`; the **new** pattern constructs
successfully and `forward` returns `cos`/`sin` on the input's CUDA
device.
- `ruff check` + `ruff format --check`: clean.
- cuTile/cuTeDSL backends aren't installed on the dev pod, but all three
test files use the identical dispatch pattern, so the fix applies
uniformly.

Co-authored-by: Copilot <[email protected]>

benchmark/scripts/benchmark_llama4_rope.py

@@ -52,7 +52,7 @@ def setup_llama4_rope(input: SingleBenchmarkRunInput):
         Llama4TextRotaryEmbedding,
         Llama4TextRotaryEmbedding,
         before_kwargs={"config": config, "device": device},
-        after_kwargs={"config": config, "device": device},
+        after_kwargs={"config": config},
     )
 
     q = torch.randn(

benchmark/scripts/benchmark_rope.py

@@ -48,7 +48,7 @@ def setup_rope(input: SingleBenchmarkRunInput):
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
     q = torch.randn(
         (1, seq_len, num_q_heads, head_dim),

src/liger_kernel/utils.py

@@ -183,7 +183,7 @@ def transformers_version_dispatch(
         ...     before_args=(head_dim,),
         ...     after_args=(LlamaConfig(head_dim=head_dim),),
         ...     before_kwargs={'device': device},
-        ...     after_kwargs={'device': device}
+        ...     after_kwargs={}  # config-based constructor derives device from input; newer transformers rejects `device`
         ... )
     """
     from packaging import version

test/cutedsl/test_rope.py

@@ -70,7 +70,7 @@ def test_cutedsl_matches_triton_and_hf(
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
@@ -148,7 +148,7 @@ def test_cutedsl_backward_contiguous_grad(bsz, seq_len, num_q_heads, num_kv_head
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
@@ -211,7 +211,7 @@ def test_cutedsl_vision_2d_cos_sin(seq_len, num_heads, head_dim, dtype, atol, rt
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((seq_len, num_heads, head_dim), device=device).to(dtype)

test/cutile/test_rope.py

@@ -80,7 +80,7 @@ def test_cutile_matches_triton_and_hf(
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
@@ -162,7 +162,7 @@ def test_cutile_backward_contiguous_grad(bsz, seq_len, num_q_heads, num_kv_heads
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
@@ -247,7 +247,7 @@ def test_cutile_vision_2d_cos_sin(seq_len, num_heads, head_dim, dtype, atol, rto
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim)},
     )
 
     # Vision layout: q, k arrive as (seq_len, num_heads, head_dim); the wrapper

test/transformers/test_rope.py

@@ -67,7 +67,7 @@ def test_correctness(
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     _tensor_q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
@@ -152,7 +152,7 @@ def test_functional_correctness(
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim)},
     )
 
     pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0)
@@ -224,7 +224,7 @@ def test_vision_2d_cos_sin(seq_len, num_heads, head_dim, dtype, atol, rtol):
         LlamaRotaryEmbedding,
         LlamaRotaryEmbedding,
         before_kwargs={"dim": head_dim, "device": device},
-        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim), "device": device},
+        after_kwargs={"config": LlamaConfig(num_kv_heads=num_heads, head_dim=head_dim)},
     )
 
     _q = torch.randn((seq_len, num_heads, head_dim), device=device).to(dtype)