[Fix] Support Paddle 3.4 compat imports (#5)

Co-authored-by: huangjiyi <huangjiyi@users.noreply.github.com>

Author: huangjiyi

fla/modules/conv/cp/ops.py

@@ -35,7 +35,7 @@ def _prepare_initial_state_for_cp(
         weight: torch.Tensor,
         cu_seqlens: torch.Tensor | None,
         context: FLACPContext,
-        group: dist.ProcessGroup | None,
+        group: "dist.ProcessGroup | None",
     ) -> torch.Tensor | None:
         """Prepare initial_state for CP forward pass by communicating with previous rank.
 
@@ -81,7 +81,7 @@ def _correct_dx_for_cp(
         dx: torch.Tensor,
         dh0: torch.Tensor | None,
         W: int,
-        group: dist.ProcessGroup | None,
+        group: "dist.ProcessGroup | None",
         is_first_rank: bool,
         pre_num_conv_tokens: int = 0,
     ) -> None:

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" = "cpu",
 ) -> torch.LongTensor:
     """Sub-split a (optionally packed) batch along the token axis.
 

tests/paddle/test_kimi_delta_attention.py

@@ -176,7 +176,7 @@ def test_huggingface_kda_training_import_contract():
 from fla.modules.backends import dispatch as modules_dispatch
 from fla.ops.backends import dispatch as ops_dispatch
 from fla.ops.kda import chunk_kda
-from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask
+from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask, prepare_split_cu_seqlens
 from fla.utils import tensor_cache
 
 assert FusedRMSNormGated is not None
@@ -185,6 +185,7 @@ def test_huggingface_kda_training_import_contract():
 assert callable(chunk_kda)
 assert callable(prepare_cu_seqlens_from_mask)
 assert callable(prepare_lens_from_mask)
+assert prepare_split_cu_seqlens(batch_size=2, seq_len=5, split_size=3).tolist() == [0, 3, 5, 8, 10]
 assert callable(tensor_cache)
 assert fla.modules is sys.modules["fla.modules"]
 assert fla.modules.conv.cp is sys.modules["fla.modules.conv.cp"]