Fix issue with flash linear attention in docker image

extensions_built_in/llm_models/src/qwen38_model.py

@@ -24,14 +24,31 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint
-from fla.modules import FusedRMSNormGated
-from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
 from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel
 
 from toolkit.util.convrot_w4a8_quant import apply_forward_hooks, module_hooks_only, w4a8_grouped_linear
 
 from .qwen38_kernels import _conv_step_op, _decode_attn_op, add_rmsnorm, decode_attention_fused_ok, rmsnorm, swiglu
 
+_FLA = None
+
+
+def _fla():
+    """flash-linear-attention, imported on first use so the extension registry loads
+    without it; the DeltaNet layers need its kernels at model build time."""
+    global _FLA
+    if _FLA is None:
+        try:
+            from fla.modules import FusedRMSNormGated
+            from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+        except ImportError as e:
+            raise ImportError(
+                "qwen38 needs the flash-linear-attention package (pip install flash-linear-attention)"
+            ) from e
+        _FLA = (FusedRMSNormGated, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
+    return _FLA
+
+
 # ------------------------------------------------------------------ norms / rope
 
 
@@ -173,7 +190,7 @@ class GatedDeltaNet(nn.Module):
         )
         self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
         self.A_log = nn.Parameter(torch.zeros(self.num_v_heads))
-        self.norm = FusedRMSNormGated(self.head_v_dim, eps=cfg.rms_norm_eps, activation="silu")
+        self.norm = _fla()[0](self.head_v_dim, eps=cfg.rms_norm_eps, activation="silu")
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
         self.in_proj_qkv = nn.Linear(self.hidden_size, self.conv_dim, bias=False)
         self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)
@@ -233,7 +250,7 @@ class GatedDeltaNet(nn.Module):
             cache.conv.copy_(F.pad(xt, (max(ks - t, 0), 0))[:, :, -ks:])
         conv = F.silu(F.conv1d(xt, self.conv1d.weight, None, padding=self.conv_kernel_size - 1, groups=self.conv_dim)[:, :, :t])
         q, k, v = self._heads(conv.transpose(1, 2))
-        core, final = chunk_gated_delta_rule(
+        core, final = _fla()[1](
             q, k, v, g=g, beta=beta, initial_state=None, output_final_state=cache is not None,
             use_qk_l2norm_in_kernel=True,
         )
@@ -249,7 +266,7 @@ class GatedDeltaNet(nn.Module):
         conv = _conv_step_op(cache.conv, mixed, self.conv1d.weight, ctx.commit)  # [B, M, C]
         q, k, v = self._heads(conv)
         # the decay / beta activations run inside the recurrent kernel
-        core, final = fused_recurrent_gated_delta_rule(
+        core, final = _fla()[2](
             q, k, v, g=a_raw.float(), beta=b_raw, initial_state=cache.rec, output_final_state=True,
             use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, A_log=self.A_log.float(),
             dt_bias=self.dt_bias.float(), use_beta_sigmoid_in_kernel=True,
@@ -268,7 +285,7 @@ class GatedDeltaNet(nn.Module):
         full = torch.cat([cache.conv, mixed.transpose(1, 2)], dim=-1)  # [B, C, ks-1+M]
         beta, g = self._gates(a_raw, b_raw)
         keep_h = keep.view(1, -1, 1).to(g.dtype)
-        _, final = fused_recurrent_gated_delta_rule(
+        _, final = _fla()[2](
             q, k, v, g=g * keep_h, beta=beta * keep_h.to(beta.dtype), initial_state=cache.rec,
             output_final_state=True, use_qk_l2norm_in_kernel=True,
         )