From a5b79ab856a4d4e57a68ff48d1ecacd53ebd7b95 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 5 May 2026 19:56:08 -0700 Subject: [PATCH 1/7] fix gpt-oss Signed-off-by: jiqing-feng --- src/transformers/modeling_utils.py | 11 +++++++++-- src/transformers/models/gpt_oss/modeling_gpt_oss.py | 3 ++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 31d59d2f4862..6e5dede41b31 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4574,11 +4574,18 @@ def kernelize(self, mode=None): def attach_hidden_kernels(module): for name, fn in getattr(module, "_hidden_kernels", {}).items(): if name not in dict(module.named_children()): - module.register_module(name, fn) + if isinstance(fn, torch.nn.Module): + module._modules[name] = fn + else: + logger.warning_once( + f"Skipping kernel registration for '{name}': not an nn.Module. " + f"Consider decorating it with @use_kernel_func_from_hub." + ) def detach_hidden_kernels(module): for name in getattr(module, "_hidden_kernels", {}): - delattr(module, name) + if name in module._modules: + del module._modules[name] try: self.apply(attach_hidden_kernels) diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 55381a7e3c21..55291bfe8411 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -27,7 +27,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -234,6 +234,7 @@ def _apply_rotary_emb( return torch.cat((first_, second_), dim=-1) +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) From 152104611ee9f5f0106151fc28264e9a6bc6e9d9 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 5 May 2026 20:05:19 -0700 Subject: [PATCH 2/7] fix format Signed-off-by: jiqing-feng --- src/transformers/models/gpt_oss/modeling_gpt_oss.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 55291bfe8411..484b071f01bc 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -27,7 +27,12 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, From d05f8cb242f7d5935d541926f2b65099c0c9797f Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 5 May 2026 20:13:45 -0700 Subject: [PATCH 3/7] fix copy Signed-off-by: jiqing-feng --- src/transformers/models/gpt_oss/modular_gpt_oss.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/gpt_oss/modular_gpt_oss.py b/src/transformers/models/gpt_oss/modular_gpt_oss.py index 3354acef2196..cc9c42364511 100644 --- a/src/transformers/models/gpt_oss/modular_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modular_gpt_oss.py @@ -19,7 +19,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import ( MoeModelOutputWithPast, @@ -171,6 +171,7 @@ def _apply_rotary_emb( return torch.cat((first_, second_), dim=-1) +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) From 694088208a3428c3c1dd4b42b6a718d2671cdd20 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 5 May 2026 20:19:00 -0700 Subject: [PATCH 4/7] fix consistency Signed-off-by: jiqing-feng --- .../modeling_openai_privacy_filter.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py index 4082a3ee9bf4..e76bd3c35c1b 100644 --- a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py @@ -26,7 +26,12 @@ from torch.nn import functional as F from ... import initialization as init -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) from ...masking_utils import create_bidirectional_sliding_window_mask from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput @@ -137,6 +142,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) +@use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) From 8e657898a4d8e9af0e029c4abab2573f7da988c9 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Wed, 6 May 2026 18:36:37 -0700 Subject: [PATCH 5/7] fix comment Signed-off-by: jiqing-feng --- src/transformers/modeling_utils.py | 6 +++--- tests/models/gpt_oss/test_modeling_gpt_oss.py | 8 ++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 6e5dede41b31..fa6fef6d70ba 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4575,7 +4575,7 @@ def attach_hidden_kernels(module): for name, fn in getattr(module, "_hidden_kernels", {}).items(): if name not in dict(module.named_children()): if isinstance(fn, torch.nn.Module): - module._modules[name] = fn + module.register_module(name, fn) else: logger.warning_once( f"Skipping kernel registration for '{name}': not an nn.Module. " @@ -4584,8 +4584,8 @@ def attach_hidden_kernels(module): def detach_hidden_kernels(module): for name in getattr(module, "_hidden_kernels", {}): - if name in module._modules: - del module._modules[name] + if hasattr(module, name): + delattr(module, name) try: self.apply(attach_hidden_kernels) diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 45b6b0e87b37..691747561106 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -69,6 +69,14 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase): model_split_percents = [0.5, 0.6] model_tester_class = GptOssModelTester + @require_kernels + def test_kernelize_does_not_crash(self): + """Regression test: kernelize() should not crash with use_kernelized_func + use_kernel_func_from_hub.""" + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + model = GptOssModel(config).to(device=torch_device) + # This used to raise TypeError because apply_rotary_pos_emb was not wrapped as nn.Module + model.use_kernels = True + @require_kernels @pytest.mark.flash_attn_test @require_torch_gpu From db88dae2facd855393a513bf445e0b862fe2d6c4 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 7 May 2026 14:04:00 +0200 Subject: [PATCH 6/7] change to error instead of warning --- src/transformers/modeling_utils.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index fa6fef6d70ba..a914e750487e 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4574,18 +4574,17 @@ def kernelize(self, mode=None): def attach_hidden_kernels(module): for name, fn in getattr(module, "_hidden_kernels", {}).items(): if name not in dict(module.named_children()): - if isinstance(fn, torch.nn.Module): - module.register_module(name, fn) - else: - logger.warning_once( - f"Skipping kernel registration for '{name}': not an nn.Module. " - f"Consider decorating it with @use_kernel_func_from_hub." + if not isinstance(fn, nn.Module): + raise ValueError( + f"Attempted to register a kernel for {name}, but it was not a `torch.nn.Module`." + "This means the underlying function needs to be decorated with `@use_kernel_func_from_hub`." + "Please submit and issue to the transformers repo: `https://github.com/huggingface/transformers/issues`." ) + module.register_module(name, fn) def detach_hidden_kernels(module): for name in getattr(module, "_hidden_kernels", {}): - if hasattr(module, name): - delattr(module, name) + delattr(module, name) try: self.apply(attach_hidden_kernels) From 170b4fa48d2081accdee0991ac6b60ef350fe121 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 7 May 2026 14:18:22 +0200 Subject: [PATCH 7/7] fixup deregister on failure mode (as before) + more comments --- src/transformers/modeling_utils.py | 9 ++++++--- tests/models/gpt_oss/test_modeling_gpt_oss.py | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a914e750487e..d30239408c9d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4576,15 +4576,18 @@ def attach_hidden_kernels(module): if name not in dict(module.named_children()): if not isinstance(fn, nn.Module): raise ValueError( - f"Attempted to register a kernel for {name}, but it was not a `torch.nn.Module`." - "This means the underlying function needs to be decorated with `@use_kernel_func_from_hub`." + f"Attempted to register a kernel for {name}, but it was not a `torch.nn.Module`. " + "This means the underlying function needs to be decorated with `@use_kernel_func_from_hub`. " "Please submit and issue to the transformers repo: `https://github.com/huggingface/transformers/issues`." ) module.register_module(name, fn) def detach_hidden_kernels(module): for name in getattr(module, "_hidden_kernels", {}): - delattr(module, name) + # Skip deregistering if it failed to properly register, + # i.e. `ValueError` will be raised afterwards + if hasattr(module, name): + delattr(module, name) try: self.apply(attach_hidden_kernels) diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 691747561106..fa39405791f4 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -71,8 +71,8 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase): @require_kernels def test_kernelize_does_not_crash(self): - """Regression test: kernelize() should not crash with use_kernelized_func + use_kernel_func_from_hub.""" - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + """Regression test #45799: `kernelize` should not crash with `use_kernelized_func` + `use_kernel_func_from_hub`.""" + config, _ = self.model_tester.prepare_config_and_inputs_for_common() model = GptOssModel(config).to(device=torch_device) # This used to raise TypeError because apply_rotary_pos_emb was not wrapped as nn.Module model.use_kernels = True