From 74777f44fb30764511c39ad1584803afb6a36d12 Mon Sep 17 00:00:00 2001 From: Vaibhav Jindal Date: Wed, 7 Oct 2026 20:52:57 +0000 Subject: [PATCH] =?UTF-8?q?fix(test):=20drop=20unsupported=20`device`=20kw?= =?UTF-8?q?arg=20from=20config-based=20RoPE=20con=E2=80=A6=20(#1514)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 <223556219+Copilot@users.noreply.github.com> --- benchmark/scripts/benchmark_llama4_rope.py | 2 +- benchmark/scripts/benchmark_rope.py | 2 +- src/liger_kernel/utils.py | 2 +- test/cutedsl/test_rope.py | 6 +++--- test/cutile/test_rope.py | 6 +++--- test/transformers/test_rope.py | 6 +++--- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/benchmark/scripts/benchmark_llama4_rope.py b/benchmark/scripts/benchmark_llama4_rope.py index 3e71dceb..a3b03239 100644 --- a/benchmark/scripts/benchmark_llama4_rope.py +++ b/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( diff --git a/benchmark/scripts/benchmark_rope.py b/benchmark/scripts/benchmark_rope.py index d7e1868c..b788cb21 100644 --- a/benchmark/scripts/benchmark_rope.py +++ b/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), diff --git a/src/liger_kernel/utils.py b/src/liger_kernel/utils.py index 8c11ac08..6421762f 100644 --- a/src/liger_kernel/utils.py +++ b/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 diff --git a/test/cutedsl/test_rope.py b/test/cutedsl/test_rope.py index 64d02a85..52002a0e 100644 --- a/test/cutedsl/test_rope.py +++ b/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) diff --git a/test/cutile/test_rope.py b/test/cutile/test_rope.py index 87afa085..5d2204e3 100644 --- a/test/cutile/test_rope.py +++ b/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 diff --git a/test/transformers/test_rope.py b/test/transformers/test_rope.py index 132ceb95..9b0a34a8 100644 --- a/test/transformers/test_rope.py +++ b/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)