[KDA] Eagerly load Paddle training imports

fla/modules/conv/cp/ops.py

@@ -8,8 +8,12 @@
 import torch
 import torch.distributed as dist
 
+# initialize FLA ops before the convolution backend to avoid the modules/ops import cycle.
+# isort: off
 from fla.ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
 from fla.ops.utils import prepare_chunk_indices
+from fla.modules.conv.triton.ops import causal_conv1d_bwd, causal_conv1d_fwd
+# isort: on
 
 
 class CausalConv1dFunctionCP(torch.autograd.Function):
@@ -129,9 +133,6 @@ def forward(
         chunk_size: int | None,
         backend: str = 'triton',
     ):
-        # Import here to avoid circular dependency
-        from fla.modules.conv.triton.ops import causal_conv1d_fwd
-
         if cp_context is None:
             raise ValueError("cp_context must be provided for CausalConv1dFunctionCP")
         cu_seqlens = cp_context.cu_seqlens
@@ -150,6 +151,8 @@ def forward(
         )
 
         ctx.save_for_backward(x, weight, bias, initial_state)
+        ctx.has_bias = bias is not None
+        ctx.has_chunk_indices = chunk_indices is not None
         ctx.activation = activation
         ctx.cu_seqlens = cu_seqlens
         ctx.cu_seqlens_cpu = cu_seqlens_cpu
@@ -179,9 +182,6 @@ def forward(
 
     @staticmethod
     def backward(ctx, dy: torch.Tensor):
-        # Import here to avoid circular dependency
-        from fla.modules.conv.triton.ops import causal_conv1d_bwd
-
         x, weight, bias, initial_state = ctx.saved_tensors
         group = ctx.group
         W = ctx.W
@@ -212,7 +212,11 @@ def backward(ctx, dy: torch.Tensor):
             pre_num_conv_tokens=ctx.pre_num_conv_tokens,
         )
 
-        return dx, dw, db, None, None, None, None, None
+        return (
+            (dx, dw)
+            + ((db,) if ctx.has_bias else ())
+            + ((None,) if ctx.has_chunk_indices else ())
+        )
 
 
 def causal_conv1d_cp(

fla/modules/conv/short_conv.py

@@ -9,12 +9,16 @@
 
 from __future__ import annotations
 
+import os
 import warnings
 
 import torch
 import torch.nn as nn
 from einops import rearrange
 
+from fla.modules.conv.causal_conv1d import causal_conv1d
+from fla.modules.conv.triton.ops import causal_conv1d_update
+
 try:
     from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda
     from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda
@@ -85,7 +89,6 @@ def __init__(
                 "The `use_fast_conv1d` parameter is deprecated and will be ignored. "
                 "Please use the `backend` parameter instead.",
             )
-        import os
         self.backend = os.environ.get('FLA_CONV_BACKEND', backend)
         if backend not in ['cuda', 'triton']:
             raise ValueError(f"Invalid backend: {backend}, must be one of ['cuda', 'triton']")
@@ -151,9 +154,6 @@ def forward(
         Returns:
             Tensor of shape `[B, T, D]`.
         """
-        # Import here to avoid circular dependency
-        from fla.modules.conv.causal_conv1d import causal_conv1d
-
         B, T, *_ = x.shape
         N = B if cu_seqlens is None else len(cu_seqlens) - 1
         if mask is not None:
@@ -208,8 +208,6 @@ def step(
         output_final_state: bool = False,
         cu_seqlens: torch.LongTensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        from fla.modules.conv.triton.ops import causal_conv1d_update
-
         B, _, D, W = *x.shape, self.kernel_size[0]
         N = B if cu_seqlens is None else len(cu_seqlens) - 1
         # Always initialise cache when None so the Triton kernel never

tests/paddle/test_kimi_delta_attention.py

@@ -5,8 +5,10 @@
 # For a list of all contributors, visit:
 #   https://github.com/fla-org/flash-linear-attention/graphs/contributors
 
+import builtins
 import subprocess
 import sys
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import paddle
@@ -19,11 +21,35 @@
 paddle.disable_compat()
 
 from fla.modules import FusedRMSNormGated, ShortConvolution  # noqa: E402
+from fla.ops.cp import FLACPContext  # noqa: E402
 from fla.ops.kda import chunk_kda  # noqa: E402
 from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask  # noqa: E402
 from fla.utils import tensor_cache  # noqa: E402
 
 
+@contextmanager
+def _assert_no_fla_runtime_imports():
+    loaded_modules = {name for name in sys.modules if name == "fla" or name.startswith("fla.")}
+    runtime_imports = []
+    original_import = builtins.__import__
+
+    def monitored_import(name, globals_=None, locals_=None, fromlist=(), level=0):
+        caller = globals_.get("__name__", "") if globals_ else ""
+        if caller == "fla" or caller.startswith("fla."):
+            runtime_imports.append((caller, name, tuple(fromlist or ()), level))
+        return original_import(name, globals_, locals_, fromlist, level)
+
+    builtins.__import__ = monitored_import
+    try:
+        yield
+    finally:
+        builtins.__import__ = original_import
+
+    assert runtime_imports == []
+    current_modules = {name for name in sys.modules if name == "fla" or name.startswith("fla.")}
+    assert current_modules == loaded_modules
+
+
 @tensor_cache
 def _get_unpad_data(attention_mask: paddle.Tensor) -> tuple[paddle.Tensor, paddle.Tensor]:
     indices = paddle.nonzero(attention_mask.flatten()).flatten()
@@ -249,6 +275,32 @@ def test_fused_rms_norm_gated_forward_backward():
     paddle.testing.assert_close(layer.weight.grad, weight_ref.grad, rtol=1e-5, atol=1e-5)
 
 
+def test_short_convolution_cp_forward_backward_uses_eager_imports():
+    paddle.seed(42)
+    conv = ShortConvolution(64, 4, activation="silu")
+    x = paddle.randn([1, 32, 64], dtype="float32")
+    x.stop_gradient = False
+    cu_seqlens = paddle.to_tensor([0, 32], dtype="int32")
+    cp_context = FLACPContext(
+        cu_seqlens=cu_seqlens,
+        cu_seqlens_cpu=paddle.to_tensor([0, 32], dtype="int32", place=paddle.CPUPlace()),
+        is_first_rank=True,
+        conv1d_kernel_size=4,
+        pre_num_conv_tokens=0,
+    )
+
+    with _assert_no_fla_runtime_imports():
+        output, final_state = conv(x, cp_context=cp_context)
+        output.square().mean().backward()
+
+    assert final_state is None
+    assert paddle.isfinite(output).all().item()
+    assert x.grad is not None
+    assert paddle.isfinite(x.grad).all().item()
+    assert conv.weight.grad is not None
+    assert paddle.isfinite(conv.weight.grad).all().item()
+
+
 @pytest.mark.parametrize("use_padding_mask", [False, True])
 def test_kimi_delta_attention_training_forward_backward(use_padding_mask: bool):
     paddle.seed(42)
@@ -262,17 +314,18 @@ def test_kimi_delta_attention_training_forward_backward(use_padding_mask: bool):
             "gate_lower_bound": -5.0,
         },
     )
-    layer = _KimiDeltaAttentionTrainingHarness(config)
-    hidden_states = paddle.randn([2 if use_padding_mask else 1, 64, 128], dtype="float32")
-    hidden_states.stop_gradient = False
-    attention_mask = None
-    if use_padding_mask:
-        attention_mask = paddle.to_tensor([[1] * 53 + [0] * 11, [1] * 64], dtype="bool")
-
-    with paddle.amp.auto_cast(enable=True, dtype="bfloat16"):
-        output = layer(hidden_states, attention_mask)
-        loss = output.astype("float32").square().mean()
-    loss.backward()
+    with _assert_no_fla_runtime_imports():
+        layer = _KimiDeltaAttentionTrainingHarness(config)
+        hidden_states = paddle.randn([2 if use_padding_mask else 1, 64, 128], dtype="float32")
+        hidden_states.stop_gradient = False
+        attention_mask = None
+        if use_padding_mask:
+            attention_mask = paddle.to_tensor([[1] * 53 + [0] * 11, [1] * 64], dtype="bool")
+
+        with paddle.amp.auto_cast(enable=True, dtype="bfloat16"):
+            output = layer(hidden_states, attention_mask)
+            loss = output.astype("float32").square().mean()
+        loss.backward()
 
     assert output.shape == hidden_states.shape
     assert paddle.isfinite(output).all().item()