diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 2ecc31ae54cf..5d6fc98f1f25 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -80,6 +80,7 @@ "replace_kernel_forward_from_hub", "use_kernel_forward_from_hub", "use_kernel_func_from_hub", + "use_kernel_func_from_hub_with_fallback", "use_kernelized_func", ], "integration_utils": [ @@ -241,6 +242,7 @@ replace_kernel_forward_from_hub, use_kernel_forward_from_hub, use_kernel_func_from_hub, + use_kernel_func_from_hub_with_fallback, use_kernelized_func, ) from .integration_utils import ( diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 25885dbb2b05..c69528c001c1 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools +import importlib +import inspect import os import re import sys @@ -31,6 +33,7 @@ is_kernels_available, is_rocm_platform, is_torch_available, + resolve_internal_import, ) from .flash_attention import flash_attention_forward @@ -531,8 +534,6 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ else: # Try to import is_{kernel_name}_available from ..utils - import importlib - new_kernel_name = kernel_name.replace("-", "_") func_name = f"is_{new_kernel_name}_available" @@ -655,6 +656,32 @@ def new_init(self, *args, **kwargs): return decorator +def use_kernel_func_from_hub_with_fallback(package: str, func_name: str, internal_path: str | None = None): + # TODO: change when we sync to 0.16.x+ + kernel_wrapper_decorator = use_kernel_func_from_hub(func_name) + + def decorator(torch_function: Callable) -> Callable: + implementation = None + try: + module = importlib.import_module(package) + implementation = resolve_internal_import(module, internal_path or func_name) + except Exception: + implementation = torch_function + finally: + implementation = torch_function if implementation is None else implementation + + applicable_params = inspect.signature(implementation).parameters + + @functools.wraps(torch_function) + def wrapped(*args, **kwargs): + kwargs = {k: v for k, v in kwargs.items() if k in applicable_params} + return implementation(*args, **kwargs) + + return kernel_wrapper_decorator(wrapped) + + return decorator + + # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community` ALLOW_ALL_KERNELS = False diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 1f695935030c..dad105e8b954 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -32,7 +32,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -58,18 +58,13 @@ maybe_autocast, merge_with_config_defaults, ) -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available +from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import capture_outputs from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids from ..auto.modeling_auto import AutoModel from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig -if is_causal_conv1d_available(): - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update -else: - causal_conv1d_update, causal_conv1d_fn = None, None - if is_flash_linear_attention_available(): from fla.modules import FusedRMSNormGated from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule @@ -216,17 +211,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): return hidden_states -is_fast_path_available = all( - (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule) -) +is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)) -def torch_causal_conv1d_update( - hidden_states, - conv_state, - weight, - bias=None, - activation=None, +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") +def causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] @@ -234,9 +228,33 @@ def torch_causal_conv1d_update( hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = F.silu(out[:, :, -seq_len:]) - out = out.to(hidden_states.dtype) - return out + out = out[:, :, -seq_len:] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") +def causal_conv1d_fn( + hidden_states: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, + **kwargs, +): + _, hidden_size, seq_len = hidden_states.shape + padding = weight.shape[-1] - 1 + + out = F.conv1d( + hidden_states.to(weight.dtype), + weight=weight.unsqueeze(1), + bias=bias, + padding=padding, + groups=hidden_size, + )[:, :, :seq_len] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): @@ -385,7 +403,6 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx self.activation = config.hidden_act - self.act = ACT2FN[config.hidden_act] self.layer_norm_epsilon = config.rms_norm_eps # QKV @@ -418,16 +435,13 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.causal_conv1d_fn = causal_conv1d_fn - self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule if not is_fast_path_available: logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" + "The fast path is not available because the required library is not installed. Falling back to " + "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation" ) self.layer_type = config.layer_types[layer_idx] @@ -472,7 +486,7 @@ def forward( if use_precomputed_states and seq_len == 1: # Single-token cached decode: the fused per-step kernel updates the conv state in-place. - mixed_qkv = self.causal_conv1d_update( + mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), @@ -489,16 +503,15 @@ def forward( if cache_params is not None: new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) cache_params.update_conv_state(new_conv_state, self.layer_idx) - if self.causal_conv1d_fn is not None: - mixed_qkv = self.causal_conv1d_fn( - x=mixed_qkv, - weight=self.conv1d.weight.squeeze(1), - bias=self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - else: - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + seq_idx=kwargs.get("seq_idx"), + ) + if use_precomputed_states: mixed_qkv = mixed_qkv[:, :, -seq_len:] diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 2b88938d533d..d81edf5015a3 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -46,6 +46,8 @@ Qwen3NextPreTrainedModel, Qwen3NextRMSNorm, apply_mask_to_padding_states, + causal_conv1d_fn, + causal_conv1d_update, ) from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig from ..qwen3_vl.modeling_qwen3_vl import ( @@ -251,7 +253,7 @@ def forward( if use_precomputed_states and seq_len == 1: # Single-token cached decode: the fused per-step kernel updates the conv state in-place. - mixed_qkv = self.causal_conv1d_update( + mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), @@ -268,16 +270,15 @@ def forward( if cache_params is not None: new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) cache_params.update_conv_state(new_conv_state, self.layer_idx) - if self.causal_conv1d_fn is not None: - mixed_qkv = self.causal_conv1d_fn( - x=mixed_qkv, - weight=self.conv1d.weight.squeeze(1), - bias=self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - else: - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + seq_idx=kwargs.get("seq_idx"), + ) + if use_precomputed_states: mixed_qkv = mixed_qkv[:, :, -seq_len:] diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 643af3afd48b..2ad29909ec05 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -32,7 +32,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -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_with_fallback, +) from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -55,18 +59,13 @@ maybe_autocast, merge_with_config_defaults, ) -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available +from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import OutputRecorder, capture_outputs from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids from ..auto.modeling_auto import AutoModel from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5MoeVisionConfig -if is_causal_conv1d_available(): - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update -else: - causal_conv1d_update, causal_conv1d_fn = None, None - if is_flash_linear_attention_available(): from fla.modules import FusedRMSNormGated from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule @@ -213,17 +212,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): return hidden_states -is_fast_path_available = all( - (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule) -) +is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)) -def torch_causal_conv1d_update( - hidden_states, - conv_state, - weight, - bias=None, - activation=None, +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") +def causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] @@ -231,9 +229,33 @@ def torch_causal_conv1d_update( hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = F.silu(out[:, :, -seq_len:]) - out = out.to(hidden_states.dtype) - return out + out = out[:, :, -seq_len:] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") +def causal_conv1d_fn( + hidden_states: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, + **kwargs, +): + _, hidden_size, seq_len = hidden_states.shape + padding = weight.shape[-1] - 1 + + out = F.conv1d( + hidden_states.to(weight.dtype), + weight=weight.unsqueeze(1), + bias=bias, + padding=padding, + groups=hidden_size, + )[:, :, :seq_len] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): @@ -382,7 +404,6 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx self.activation = config.hidden_act - self.act = ACT2FN[config.hidden_act] self.layer_norm_epsilon = config.rms_norm_eps # QKV @@ -415,16 +436,13 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.causal_conv1d_fn = causal_conv1d_fn - self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule if not is_fast_path_available: logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" + "The fast path is not available because the required library is not installed. Falling back to " + "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation" ) self.layer_type = config.layer_types[layer_idx] @@ -469,7 +487,7 @@ def forward( if use_precomputed_states and seq_len == 1: # Single-token cached decode: the fused per-step kernel updates the conv state in-place. - mixed_qkv = self.causal_conv1d_update( + mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), @@ -486,16 +504,15 @@ def forward( if cache_params is not None: new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) cache_params.update_conv_state(new_conv_state, self.layer_idx) - if self.causal_conv1d_fn is not None: - mixed_qkv = self.causal_conv1d_fn( - x=mixed_qkv, - weight=self.conv1d.weight.squeeze(1), - bias=self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - else: - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + seq_idx=kwargs.get("seq_idx"), + ) + if use_precomputed_states: mixed_qkv = mixed_qkv[:, :, -seq_len:] diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 95c506aaf8e1..e618ad1462e8 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -29,7 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation +from ...integrations import use_experts_implementation, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -45,16 +45,11 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available +from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_qwen3_next import Qwen3NextConfig -if is_causal_conv1d_available(): - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update -else: - causal_conv1d_update, causal_conv1d_fn = None, None - if is_flash_linear_attention_available(): from fla.modules import FusedRMSNormGated from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule @@ -342,17 +337,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): return hidden_states -is_fast_path_available = all( - (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule) -) +is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)) -def torch_causal_conv1d_update( - hidden_states, - conv_state, - weight, - bias=None, - activation=None, +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") +def causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] @@ -360,9 +354,33 @@ def torch_causal_conv1d_update( hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = F.silu(out[:, :, -seq_len:]) - out = out.to(hidden_states.dtype) - return out + out = out[:, :, -seq_len:] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") +def causal_conv1d_fn( + hidden_states: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, + **kwargs, +): + _, hidden_size, seq_len = hidden_states.shape + padding = weight.shape[-1] - 1 + + out = F.conv1d( + hidden_states.to(weight.dtype), + weight=weight.unsqueeze(1), + bias=bias, + padding=padding, + groups=hidden_size, + )[:, :, :seq_len] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): @@ -496,6 +514,7 @@ def torch_recurrent_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class Qwen3NextGatedDeltaNet(nn.Module): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__() @@ -510,7 +529,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx self.activation = config.hidden_act - self.act = ACT2FN[config.hidden_act] self.layer_norm_epsilon = config.rms_norm_eps # QKV @@ -549,16 +567,13 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.causal_conv1d_fn = causal_conv1d_fn - self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule if not is_fast_path_available: logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" + "The fast path is not available because the required library is not installed. Falling back to " + "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation" ) self.layer_type = config.layer_types[layer_idx] @@ -626,7 +641,7 @@ def forward( if use_precomputed_states and seq_len == 1: # Single-token cached decode: the fused per-step kernel updates the conv state in-place. - mixed_qkv = self.causal_conv1d_update( + mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), @@ -643,16 +658,15 @@ def forward( if cache_params is not None: new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) cache_params.update_conv_state(new_conv_state, self.layer_idx) - if self.causal_conv1d_fn is not None: - mixed_qkv = self.causal_conv1d_fn( - x=mixed_qkv, - weight=self.conv1d.weight.squeeze(1), - bias=self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - else: - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + seq_idx=kwargs.get("seq_idx"), + ) + if use_precomputed_states: mixed_qkv = mixed_qkv[:, :, -seq_len:] diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index 86f31ad1431d..6e9e7bffba37 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -23,6 +23,7 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache +from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -31,10 +32,7 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, logging from ...utils.generic import merge_with_config_defaults, no_inherit_decorator -from ...utils.import_utils import ( - is_causal_conv1d_available, - is_flash_linear_attention_available, -) +from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import OutputRecorder, capture_outputs from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding @@ -55,11 +53,6 @@ from .configuration_qwen3_next import Qwen3NextConfig -if is_causal_conv1d_available(): - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update -else: - causal_conv1d_update, causal_conv1d_fn = None, None - if is_flash_linear_attention_available(): from fla.modules import FusedRMSNormGated from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule @@ -68,9 +61,7 @@ FusedRMSNormGated = None -is_fast_path_available = all( - (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule) -) +is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)) logger = logging.get_logger(__name__) @@ -188,12 +179,13 @@ def forward( return attn_output, attn_weights -def torch_causal_conv1d_update( - hidden_states, - conv_state, - weight, - bias=None, - activation=None, +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") +def causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] @@ -201,9 +193,33 @@ def torch_causal_conv1d_update( hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = F.silu(out[:, :, -seq_len:]) - out = out.to(hidden_states.dtype) - return out + out = out[:, :, -seq_len:] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") +def causal_conv1d_fn( + hidden_states: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, + **kwargs, +): + _, hidden_size, seq_len = hidden_states.shape + padding = weight.shape[-1] - 1 + + out = F.conv1d( + hidden_states.to(weight.dtype), + weight=weight.unsqueeze(1), + bias=bias, + padding=padding, + groups=hidden_size, + )[:, :, :seq_len] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): @@ -337,6 +353,7 @@ def torch_recurrent_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class Qwen3NextGatedDeltaNet(nn.Module): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__() @@ -351,7 +368,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx self.activation = config.hidden_act - self.act = ACT2FN[config.hidden_act] self.layer_norm_epsilon = config.rms_norm_eps # QKV @@ -390,16 +406,13 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.causal_conv1d_fn = causal_conv1d_fn - self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule if not is_fast_path_available: logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" + "The fast path is not available because the required library is not installed. Falling back to " + "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation" ) self.layer_type = config.layer_types[layer_idx] @@ -467,7 +480,7 @@ def forward( if use_precomputed_states and seq_len == 1: # Single-token cached decode: the fused per-step kernel updates the conv state in-place. - mixed_qkv = self.causal_conv1d_update( + mixed_qkv = causal_conv1d_update( mixed_qkv, conv_state, self.conv1d.weight.squeeze(1), @@ -484,16 +497,15 @@ def forward( if cache_params is not None: new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) cache_params.update_conv_state(new_conv_state, self.layer_idx) - if self.causal_conv1d_fn is not None: - mixed_qkv = self.causal_conv1d_fn( - x=mixed_qkv, - weight=self.conv1d.weight.squeeze(1), - bias=self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - else: - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + seq_idx=kwargs.get("seq_idx"), + ) + if use_precomputed_states: mixed_qkv = mixed_qkv[:, :, -seq_len:]