From 887c82fa0d1bbe609410087a9f3901c9c81285d8 Mon Sep 17 00:00:00 2001 From: huangjiyi Date: Thu, 30 Jul 2026 20:38:29 +0800 Subject: [PATCH] [Fix] Stringify torch.device annotations for Paddle --- fla/layers/mamba3.py | 4 ++-- fla/layers/yoco.py | 2 +- fla/modules/conv/short_conv.py | 2 +- fla/modules/feature_map.py | 2 +- fla/modules/fused_norm_gate.py | 16 ++++++++-------- fla/modules/layernorm.py | 12 ++++++------ fla/modules/layernorm_gated.py | 4 ++-- fla/modules/rotary.py | 2 +- fla/ops/attnres/backends/triton_ascend/fused.py | 2 +- fla/ops/common/intracard_cp.py | 2 +- .../backends/triton_ascend/wy_fast.py | 2 +- fla/ops/utils/index.py | 2 +- 12 files changed, 26 insertions(+), 26 deletions(-) diff --git a/fla/layers/mamba3.py b/fla/layers/mamba3.py index 58f700c8ac..b81bfea32a 100644 --- a/fla/layers/mamba3.py +++ b/fla/layers/mamba3.py @@ -81,7 +81,7 @@ def __init__( use_bias: bool = False, norm_eps: float = 1e-5, layer_idx: int | None = None, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> Mamba3: super().__init__() @@ -449,7 +449,7 @@ def allocate_inference_cache( self, batch_size: int, max_seqlen: int, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ): device = device or self.in_proj.weight.device diff --git a/fla/layers/yoco.py b/fla/layers/yoco.py index 4fc7c53f59..3211f6f37e 100644 --- a/fla/layers/yoco.py +++ b/fla/layers/yoco.py @@ -38,7 +38,7 @@ def __init__( scale_base: float | None = None, interleaved: bool = False, pos_idx_in_fp32: bool = True, - device: torch.device | None = None, + device: "torch.device | None" = None, rope_inv_freq: str = 'fla', ): self.rope_inv_freq = rope_inv_freq diff --git a/fla/modules/conv/short_conv.py b/fla/modules/conv/short_conv.py index 254159c37e..545d4f1281 100644 --- a/fla/modules/conv/short_conv.py +++ b/fla/modules/conv/short_conv.py @@ -62,7 +62,7 @@ def __init__( bias: bool = False, activation: str | None = 'silu', backend: str | None = 'triton', - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, **kwargs, ): diff --git a/fla/modules/feature_map.py b/fla/modules/feature_map.py index 2de94ffae8..1c61650c65 100644 --- a/fla/modules/feature_map.py +++ b/fla/modules/feature_map.py @@ -20,7 +20,7 @@ @functools.cache -def _triu_indices(n: int, offset: int, device: torch.device) -> torch.Tensor: +def _triu_indices(n: int, offset: int, device: "torch.device") -> torch.Tensor: # cache the upper-triangular gather indices per (size, offset, device) to avoid rebuilding # them and copying host -> device on every forward return torch.triu_indices(n, n, offset, device=device) diff --git a/fla/modules/fused_norm_gate.py b/fla/modules/fused_norm_gate.py index 55bf75f182..0d5014efd2 100644 --- a/fla/modules/fused_norm_gate.py +++ b/fla/modules/fused_norm_gate.py @@ -933,7 +933,7 @@ def __init__( bias: bool = False, activation: str = "swish", eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedLayerNormGated: factory_kwargs = {"device": device, "dtype": dtype} @@ -999,7 +999,7 @@ def __init__( elementwise_affine: bool = True, eps: float = 1e-5, activation: str = "swish", - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedRMSNormGated: factory_kwargs = {"device": device, "dtype": dtype} @@ -1062,7 +1062,7 @@ def __init__( elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedLayerNormSwishGate: super().__init__( @@ -1081,7 +1081,7 @@ def __init__( hidden_size: int, elementwise_affine: bool = True, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedRMSNormSwishGate: super().__init__( @@ -1099,7 +1099,7 @@ def __init__( hidden_size: int, elementwise_affine: bool = True, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedLayerNormGatedLinear: factory_kwargs = {"device": device, "dtype": dtype} @@ -1159,7 +1159,7 @@ def __init__( hidden_size: int, elementwise_affine: bool = True, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedLayerNormSwishGateLinear: super().__init__( @@ -1177,7 +1177,7 @@ def __init__( hidden_size, elementwise_affine: bool = True, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedRMSNormGatedLinear: factory_kwargs = {"device": device, "dtype": dtype} @@ -1236,7 +1236,7 @@ def __init__( hidden_size: int, elementwise_affine: bool = True, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> FusedRMSNormSwishGateLinear: super().__init__( diff --git a/fla/modules/layernorm.py b/fla/modules/layernorm.py index ef57bbba7c..2d76ab300f 100644 --- a/fla/modules/layernorm.py +++ b/fla/modules/layernorm.py @@ -980,7 +980,7 @@ def __init__( elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> LayerNorm: factory_kwargs = {"device": device, "dtype": dtype} @@ -1035,7 +1035,7 @@ def __init__( bias: bool = False, eps: float = 1e-5, is_rms_norm: bool = False, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> GroupNorm: factory_kwargs = {"device": device, "dtype": dtype} @@ -1097,7 +1097,7 @@ def __init__( elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> RMSNorm: factory_kwargs = {"device": device, "dtype": dtype} @@ -1255,7 +1255,7 @@ def __init__( elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> LayerNormLinear: factory_kwargs = {"device": device, "dtype": dtype} @@ -1313,7 +1313,7 @@ def __init__( bias: bool = False, eps: float = 1e-5, is_rms_norm: bool = False, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> GroupNormLinear: factory_kwargs = {"device": device, "dtype": dtype} @@ -1377,7 +1377,7 @@ def __init__( elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ) -> RMSNormLinear: factory_kwargs = {"device": device, "dtype": dtype} diff --git a/fla/modules/layernorm_gated.py b/fla/modules/layernorm_gated.py index 3d9f9e1a4f..a3e92bda2b 100644 --- a/fla/modules/layernorm_gated.py +++ b/fla/modules/layernorm_gated.py @@ -475,7 +475,7 @@ def __init__( eps: float = 1e-5, group_size: int | None = None, norm_before_gate: bool = True, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ): """If group_size is not None, we do GroupNorm with each group having group_size elements. @@ -510,7 +510,7 @@ def __init__( eps: float = 1e-5, group_size: int | None = None, norm_before_gate: bool = False, - device: torch.device | None = None, + device: "torch.device | None" = None, dtype: torch.dtype | None = None, ): """If group_size is not None, we do GroupNorm with each group having group_size elements. diff --git a/fla/modules/rotary.py b/fla/modules/rotary.py index ccf6b649a2..2a685c027d 100644 --- a/fla/modules/rotary.py +++ b/fla/modules/rotary.py @@ -342,7 +342,7 @@ def __init__( scale_base: float | None = None, interleaved: bool = False, pos_idx_in_fp32: bool = True, - device: torch.device | None = None, + device: "torch.device | None" = None, ): """ interleaved: diff --git a/fla/ops/attnres/backends/triton_ascend/fused.py b/fla/ops/attnres/backends/triton_ascend/fused.py index fa58d3b8fc..5115ee9f1d 100644 --- a/fla/ops/attnres/backends/triton_ascend/fused.py +++ b/fla/ops/attnres/backends/triton_ascend/fused.py @@ -388,7 +388,7 @@ def _get_o_mix( logit: torch.Tensor, lse: torch.Tensor, scale: float, - device: torch.device, + device: "torch.device", o_pre: torch.Tensor | None = None, o_mix: torch.Tensor | None = None, ) -> torch.Tensor: diff --git a/fla/ops/common/intracard_cp.py b/fla/ops/common/intracard_cp.py index a1be749c00..929eb241ad 100644 --- a/fla/ops/common/intracard_cp.py +++ b/fla/ops/common/intracard_cp.py @@ -297,7 +297,7 @@ def intracard_merge( num_non_first: int, merge_seq_offsets: list[int], merge_init_offsets: list[int], - device: torch.device, + device: "torch.device", initial_state: torch.Tensor | None = None, state_v_first: bool = False, ) -> tuple[torch.Tensor | None, int]: diff --git a/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py b/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py index fe6625bfb1..39c7f79c24 100644 --- a/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py +++ b/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py @@ -50,7 +50,7 @@ def _beta_npu_arg(beta: torch.Tensor, HV: int) -> tuple[torch.Tensor, bool]: def _t_npu_buf( - B: int, T: int, HV: int, *, dtype: torch.dtype, device: torch.device, + B: int, T: int, HV: int, *, dtype: torch.dtype, device: "torch.device", ) -> tuple[torch.Tensor, bool]: if HV == 1: return torch.empty(B, T, HV, dtype=dtype, device=device), False diff --git a/fla/ops/utils/index.py b/fla/ops/utils/index.py index 7b5597378d..bb9021fdfd 100644 --- a/fla/ops/utils/index.py +++ b/fla/ops/utils/index.py @@ -72,7 +72,7 @@ def prepare_split_cu_seqlens( split_size: int | None = None, cu_seqlens: torch.LongTensor | None = None, dtype: torch.dtype | None = torch.int32, - device: torch.device | None = torch.device('cpu'), + device: "torch.device | None" = torch.device('cpu'), ) -> torch.LongTensor: """Sub-split a (optionally packed) batch along the token axis.