From 0a49c65df411432f9c28d74fc74a8c6dee646310 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 31 Aug 2026 19:23:48 +0800 Subject: [PATCH 1/3] Build the inference rotary embedding through LlamaConfig transformers 4.48 replaced LlamaRotaryEmbedding(dim, base=..., device=...) with a constructor that reads both off a config, and changed forward from taking a token count to taking position_ids. The fallback attention path used both of the old shapes, so InferenceContext.get_rotary raised TypeError: LlamaRotaryEmbedding.__init__() got an unexpected keyword argument 'base' before any kernel ran, on every transformers in the supported range. head_dim carries the rotary width and rope_theta the base; transformers 5 folds rope_theta into rope_parameters itself, so the config form works on both. The position_ids handed to the rotary are the ones already passed to apply_rotary_pos_emb two lines below. Signed-off-by: alanhuangyoo --- .../inference/op_binding/softmax_context.py | 4 +- .../inference/op_binding/workspace.py | 10 ++- .../inference/test_rotary_embedding.py | 63 +++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) create mode 100644 tests/unit/ops/transformer/inference/test_rotary_embedding.py diff --git a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py index d745df678e93..216abb6b22b3 100644 --- a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py +++ b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py @@ -80,7 +80,9 @@ def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotat from transformers.models.llama.modeling_llama import apply_rotary_pos_emb rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device) - cos, sin = rotary(bat_0213_value, InferenceContext.Instance().get_max_tokens_num()) + # LlamaRotaryEmbedding.forward takes position_ids, not a token count. These are + # the same ids handed to apply_rotary_pos_emb two lines down. + cos, sin = rotary(bat_0213_value, position_ids) bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids) bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt, diff --git a/deepspeed/ops/transformer/inference/op_binding/workspace.py b/deepspeed/ops/transformer/inference/op_binding/workspace.py index 0565666559bf..efd9104a57b4 100644 --- a/deepspeed/ops/transformer/inference/op_binding/workspace.py +++ b/deepspeed/ops/transformer/inference/op_binding/workspace.py @@ -137,9 +137,17 @@ def get_kv_cache(self): def get_rotary(self, rotary_dim, rope_theta, device=None): if self.rotary is None: + from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding - self.rotary = LlamaRotaryEmbedding(rotary_dim, base=rope_theta, device=device) + # transformers 4.48 replaced the (dim, base=..., device=...) signature with one + # that reads both off a config. head_dim carries the rotary width and rope_theta + # the base; transformers 5 folds rope_theta into rope_parameters itself, so + # passing it this way works on both. + config = LlamaConfig(head_dim=rotary_dim, rope_theta=rope_theta) + self.rotary = LlamaRotaryEmbedding(config) + if device is not None: + self.rotary = self.rotary.to(device) return self.rotary diff --git a/tests/unit/ops/transformer/inference/test_rotary_embedding.py b/tests/unit/ops/transformer/inference/test_rotary_embedding.py new file mode 100644 index 000000000000..16cbad53f45d --- /dev/null +++ b/tests/unit/ops/transformer/inference/test_rotary_embedding.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""InferenceContext.get_rotary has to build a rotary embedding the installed transformers accepts. + +transformers 4.48 replaced `LlamaRotaryEmbedding(dim, base=..., device=...)` with a +config-taking constructor, and its forward went from a token count to `position_ids`. The +fallback attention path in `softmax_context.py` used both of the old shapes, so it raised +`TypeError: __init__() got an unexpected keyword argument 'base'` before reaching any kernel. + +No accelerator needed: this covers the construction and the rope values, and the rest of the +fallback path is unchanged. +""" + +import pytest +import torch + +from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext + + +def _reference_cos_sin(rotary_dim, rope_theta, seq_len): + """cos/sin straight from the rope definition, independent of transformers.""" + inv_freq = 1.0 / (rope_theta**(torch.arange(0, rotary_dim, 2).float() / rotary_dim)) + freqs = torch.outer(torch.arange(seq_len).float(), inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +@pytest.fixture +def context(): + ctx = InferenceContext.Instance() + ctx.rotary = None + yield ctx + ctx.rotary = None + + +@pytest.mark.parametrize("rotary_dim", [32, 64, 128]) +@pytest.mark.parametrize("rope_theta", [10000.0, 500000.0]) +def test_get_rotary_matches_the_rope_definition(context, rotary_dim, rope_theta): + seq_len = 12 + rotary = context.get_rotary(rotary_dim, rope_theta) + + position_ids = torch.arange(seq_len).unsqueeze(0) + cos, sin = rotary(torch.zeros(1, 1, seq_len, rotary_dim), position_ids) + + expected_cos, expected_sin = _reference_cos_sin(rotary_dim, rope_theta, seq_len) + assert cos.shape == (1, seq_len, rotary_dim) + torch.testing.assert_close(cos[0], expected_cos) + torch.testing.assert_close(sin[0], expected_sin) + + +def test_get_rotary_uses_rotary_dim_not_the_config_default(context): + """rotary_dim has to reach the embedding, not the LlamaConfig hidden_size // heads default.""" + rotary = context.get_rotary(32, 10000.0) + + assert rotary.inv_freq.numel() == 16 + + +def test_get_rotary_is_cached(context): + first = context.get_rotary(64, 10000.0) + + assert context.get_rotary(64, 10000.0) is first From 30baf0ea663b34dbc7e58c67a10e4475cb4c930f Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 16:25:27 +0800 Subject: [PATCH 2/3] Stop passing position_ids into apply_rotary_pos_emb's fifth slot The constructor fix landed but the call below it still used the pre-5.0 signature. transformers 5.0 removed the deprecated position_ids parameter, so the fifth positional argument is unsqueeze_dim: 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) Passing position_ids there reaches unsqueeze(dim=...) as a tensor, so the fallback attention path still raised on 5.x after the constructor was fixed: TypeError: unsqueeze(): argument 'dim' (position 1) must be int, not Tensor The ids only ever entered through cos/sin, which rotary() already receives, so four arguments is both correct and version-independent. Signed-off-by: alanhuangyoo --- .../inference/op_binding/softmax_context.py | 12 +++++-- .../inference/test_rotary_embedding.py | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py index 216abb6b22b3..a199f49d8528 100644 --- a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py +++ b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py @@ -80,10 +80,16 @@ def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotat from transformers.models.llama.modeling_llama import apply_rotary_pos_emb rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device) - # LlamaRotaryEmbedding.forward takes position_ids, not a token count. These are - # the same ids handed to apply_rotary_pos_emb two lines down. + # LlamaRotaryEmbedding.forward takes position_ids, not a token count. cos, sin = rotary(bat_0213_value, position_ids) - bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids) + # apply_rotary_pos_emb takes them only through cos/sin. transformers 5.0 + # dropped the deprecated position_ids parameter, so the fifth positional + # slot is unsqueeze_dim there: + # 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) + # 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) + # Passing four arguments is correct on both, and leaves unsqueeze_dim at its + # default rather than handing it a tensor. + bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin) bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt, bat_0213_key, bat_0213_value) diff --git a/tests/unit/ops/transformer/inference/test_rotary_embedding.py b/tests/unit/ops/transformer/inference/test_rotary_embedding.py index 16cbad53f45d..0ecefc3d8b8e 100644 --- a/tests/unit/ops/transformer/inference/test_rotary_embedding.py +++ b/tests/unit/ops/transformer/inference/test_rotary_embedding.py @@ -13,6 +13,8 @@ fallback path is unchanged. """ +import inspect + import pytest import torch @@ -61,3 +63,34 @@ def test_get_rotary_is_cached(context): first = context.get_rotary(64, 10000.0) assert context.get_rotary(64, 10000.0) is first + + +def test_rotary_is_applied_through_cos_sin_not_a_fifth_argument(): + """The fallback hands `apply_rotary_pos_emb` four arguments, and has to. + + transformers 5.0 dropped the deprecated `position_ids` parameter, so the fifth + positional slot became `unsqueeze_dim`: + + 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) + 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) + + Passing position_ids there reaches `unsqueeze(dim=...)` as a tensor. + """ + llama = pytest.importorskip("transformers.models.llama.modeling_llama") + apply_rotary_pos_emb = llama.apply_rotary_pos_emb + + seq_len, rotary_dim = 8, 16 + q = torch.randn(1, 4, seq_len, rotary_dim) + k = torch.randn(1, 4, seq_len, rotary_dim) + cos = torch.randn(1, seq_len, rotary_dim) + sin = torch.randn(1, seq_len, rotary_dim) + position_ids = torch.arange(seq_len).unsqueeze(0) + + rotated_q, rotated_k = apply_rotary_pos_emb(q, k, cos, sin) + assert rotated_q.shape == q.shape + assert rotated_k.shape == k.shape + + fifth = list(inspect.signature(apply_rotary_pos_emb).parameters)[4] + if fifth == "unsqueeze_dim": + with pytest.raises(TypeError): + apply_rotary_pos_emb(q, k, cos, sin, position_ids) From a29573ad68b6c8940f9a0ead19e4579be12bbde8 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 19:17:30 +0800 Subject: [PATCH 3/3] Pin the rotary call site instead of the transformers signature The previous test called apply_rotary_pos_emb directly, so it asserted what the library does, not what softmax_context_fallback does; reverting the fix left it green. Patch apply_rotary_pos_emb and drive the fallback instead, stopping at the rotary block so no workspace is needed. Signed-off-by: alanhuangyoo --- .../inference/test_rotary_embedding.py | 52 ++++++++++++------- 1 file changed, 34 insertions(+), 18 deletions(-) diff --git a/tests/unit/ops/transformer/inference/test_rotary_embedding.py b/tests/unit/ops/transformer/inference/test_rotary_embedding.py index 0ecefc3d8b8e..f728e60fabfd 100644 --- a/tests/unit/ops/transformer/inference/test_rotary_embedding.py +++ b/tests/unit/ops/transformer/inference/test_rotary_embedding.py @@ -13,11 +13,11 @@ fallback path is unchanged. """ -import inspect - import pytest import torch +from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig +from deepspeed.ops.transformer.inference.op_binding.softmax_context import SoftmaxContextOp from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext @@ -65,8 +65,12 @@ def test_get_rotary_is_cached(context): assert context.get_rotary(64, 10000.0) is first -def test_rotary_is_applied_through_cos_sin_not_a_fifth_argument(): - """The fallback hands `apply_rotary_pos_emb` four arguments, and has to. +class _StopAfterRotary(Exception): + """Ends the fallback at the call under test, before it wants a workspace.""" + + +def test_the_fallback_passes_apply_rotary_pos_emb_four_arguments(monkeypatch, context): + """Asserts this repo's call, not the library's signature. transformers 5.0 dropped the deprecated `position_ids` parameter, so the fifth positional slot became `unsqueeze_dim`: @@ -74,23 +78,35 @@ def test_rotary_is_applied_through_cos_sin_not_a_fifth_argument(): 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) - Passing position_ids there reaches `unsqueeze(dim=...)` as a tensor. + Passing `position_ids` there reaches `unsqueeze(dim=...)` as a tensor. Checking the + library's own signature would not catch that, since the mistake is in the caller; the + recorded call has to come from `softmax_context_fallback` itself. + + The recorder raises so execution stops at the rotary block. `update_cache` sits two + lines below and needs a workspace, which is not what this is about. """ llama = pytest.importorskip("transformers.models.llama.modeling_llama") - apply_rotary_pos_emb = llama.apply_rotary_pos_emb - seq_len, rotary_dim = 8, 16 - q = torch.randn(1, 4, seq_len, rotary_dim) - k = torch.randn(1, 4, seq_len, rotary_dim) - cos = torch.randn(1, seq_len, rotary_dim) - sin = torch.randn(1, seq_len, rotary_dim) + recorded = {} + + def recorder(*args, **kwargs): + recorded["args"], recorded["kwargs"] = args, kwargs + raise _StopAfterRotary + + monkeypatch.setattr(llama, "apply_rotary_pos_emb", recorder) + + heads, head_dim, seq_len, rotary_dim = 4, 16, 8, 16 + query_key_value = torch.randn(1, seq_len, 3 * heads * head_dim) position_ids = torch.arange(seq_len).unsqueeze(0) - rotated_q, rotated_k = apply_rotary_pos_emb(q, k, cos, sin) - assert rotated_q.shape == q.shape - assert rotated_k.shape == k.shape + config = DeepSpeedInferenceConfig(hidden_size=heads * head_dim, heads=heads, dtype=torch.float32) + op = SoftmaxContextOp.__new__(SoftmaxContextOp) + op.config = config + + with pytest.raises(_StopAfterRotary): + op.softmax_context_fallback(query_key_value, None, rotary_dim, True, False, heads, heads, 1.0, False, False, 0, + False, 0, 1, None, 10000.0, True, 0, position_ids) - fifth = list(inspect.signature(apply_rotary_pos_emb).parameters)[4] - if fifth == "unsqueeze_dim": - with pytest.raises(TypeError): - apply_rotary_pos_emb(q, k, cos, sin, position_ids) + assert len(recorded["args"]) == 4, \ + f"the fallback passed {len(recorded['args'])} positional arguments; the fifth is unsqueeze_dim" + assert not recorded["kwargs"], f"unexpected keyword arguments: {sorted(recorded['kwargs'])}"