diff --git a/docker/Dockerfile b/docker/Dockerfile index 8d0d0fb..b15709a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -78,6 +78,12 @@ RUN pip install --no-cache-dir --break-system-packages -U \ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl && \ python -c "import flash_attn, natten, torchcodec; print('accelerators OK:', flash_attn.__version__, natten.__version__, torchcodec.__version__)" +# flash-linear-attention (DeltaNet kernels for Qwen3.5/3.8), matching manager/spec.py FLA. +# Pure-Python Triton package installed bare and WITHOUT -U: its unpinned transformers +# requirement must not upgrade the pinned transformers from requirements_base.txt. +RUN pip install --no-cache-dir --break-system-packages flash-linear-attention==0.5.2 && \ + python -c "import fla, transformers; print('fla OK:', fla.__version__, 'transformers', transformers.__version__)" + # Install Node dependencies (only re-runs when package.json / package-lock.json change) COPY ui/package.json ui/package-lock.json /app/ai-toolkit/ui/ RUN cd /app/ai-toolkit/ui && npm ci diff --git a/extensions_built_in/llm_models/src/qwen38_model.py b/extensions_built_in/llm_models/src/qwen38_model.py index 16da7dc..804e642 100644 --- a/extensions_built_in/llm_models/src/qwen38_model.py +++ b/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, ) diff --git a/version.py b/version.py index 9d6b13d..31c9518 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.33" +VERSION = "0.13.34"