feat: add FlashQLA backend for Qwen GDN and skip selected comm memory checks

build_conda.sh

@@ -42,6 +42,8 @@ MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation
 pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
 pip install --no-build-isolation "transformer_engine[pytorch]==2.10.0"
 pip install flash-linear-attention==0.4.1
+# FlashQLA: optional GDN backend for Qwen3.5/Qwen3-Next (--qwen-gdn-backend flashqla; requires SM90+)
+pip install git+https://github.com/QwenLM/FlashQLA.git --no-build-isolation
 NVCC_APPEND_FLAGS="--threads 4" \
   pip -v install --disable-pip-version-check --no-cache-dir \
   --no-build-isolation \

docker/Dockerfile

@@ -36,6 +36,8 @@ RUN git clone https://github.com/Dao-AILab/flash-attention.git && \
 RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
 
 RUN pip install flash-linear-attention==0.4.1
+# FlashQLA: optional GDN backend for Qwen3.5/Qwen3-Next (--qwen-gdn-backend flashqla; requires SM90+)
+RUN pip install git+https://github.com/QwenLM/FlashQLA.git --no-build-isolation
 RUN pip install tilelang -f https://tile-ai.github.io/whl/nightly/cu128/
 
 # TE does not have wheel on cuda 13 yet, thus need to install from source

docker/Dockerfile.gb10

@@ -168,6 +168,18 @@ RUN pip install --no-deps \
         flash-linear-attention==0.4.1 \
         tilelang==0.1.8
 
+# FlashQLA (--qwen-gdn-backend flashqla): Qwen GDN kernels built for Hopper (SM90).
+# This image targets GB10 (sm_121a, aarch64). FlashQLA is not validated here;
+# compiling from source often fails and --qwen-gdn-backend flashqla is unsupported
+# at runtime even if install succeeds. Use the default backend fla (FLA above).
+# Opt in only for experiments: docker build --build-arg INSTALL_FLASHQLA=1 ...
+ARG INSTALL_FLASHQLA=0
+RUN if [ "${INSTALL_FLASHQLA}" = "1" ]; then \
+      pip install git+https://github.com/QwenLM/FlashQLA.git --no-build-isolation; \
+    else \
+      echo "Skipping FlashQLA (INSTALL_FLASHQLA=0; use --qwen-gdn-backend fla on GB10)"; \
+    fi
+
 RUN pip install --no-deps --no-build-isolation \
         "git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl" \
         "nvidia-modelopt>=0.37.0"

docs/zh/developer_guide/install_flashqla.md

@@ -0,0 +1,47 @@
+# 安装 FlashQLA
+
+FlashQLA 是 Qwen GDN kernel 的可选运行后端。安装 FlashQLA 后，仍需要在训练命令中显式加入：
+
+```bash
+--qwen-gdn-backend flashqla
+```
+
+如果不传该参数，Qwen GDN 仍使用默认的 FLA 后端。
+
+## 环境要求
+
+使用 `--qwen-gdn-backend flashqla` 前，请确认训练节点满足：
+
+- PyTorch 2.8 或更新版本。
+- CUDA 12.8 或更新版本。
+- NVIDIA SM90 或更新架构 GPU。
+- 所有训练节点都安装了同一套 FlashQLA Python 包。
+
+## Conda / 本地环境
+
+`build_conda.sh` 会默认安装 FlashQLA：
+
+```bash
+bash build_conda.sh
+```
+
+## Docker 镜像
+
+标准 CUDA Docker 镜像会默认安装 FlashQLA：
+
+```bash
+docker build \
+  -f docker/Dockerfile \
+  -t slime:flashqla .
+```
+
+## GB10 镜像
+
+`docker/Dockerfile.gb10` 默认不安装 FlashQLA。GB10 环境需要自行验证 FlashQLA 编译和运行行为；如需实验性安装，可显式传入：
+
+```bash
+docker build \
+  --build-arg INSTALL_FLASHQLA=1 \
+  -f docker/Dockerfile.gb10 \
+  -t slime:gb10-flashqla .
+```

slime/utils/arguments.py

@@ -112,6 +112,13 @@ def add_train_arguments(parser):
                 default="thd",
                 help="The qkv layout for Megatron backend.",
             )
+            parser.add_argument(
+                "--qwen-gdn-backend",
+                type=str,
+                choices=["fla", "flashqla"],
+                default="fla",
+                help="GDN implementation backend for Qwen linear-attention layers.",
+            )
             parser.add_argument(
                 "--train-env-vars",
                 type=json.loads,

slime/utils/reloadable_process_group.py

@@ -11,6 +11,21 @@
 
 old_new_group_dict = {}
 
+_COMM_MEMORY_CHECK_SKIP_OPS = {
+    "all_gather_into_tensor",
+    "allgather_into_tensor_coalesced",
+    "barrier",
+    "broadcast_object_list",
+    "reduce_scatter_tensor",
+    "all_to_all_single",
+    "isend",
+    "irecv",
+}
+
+
+def _should_check_memory_for_comm(op_name):
+    return op_name not in _COMM_MEMORY_CHECK_SKIP_OPS
+
 
 def monkey_patch_torch_dist():
     pid = os.getpid()
@@ -57,13 +72,14 @@ def new_function(*args, **kwargs):
 
         return new_function
 
-    def get_new_comm_function(func):
+    def get_new_comm_function(func, op_name=None):
         """Wrap communication functions with memory check."""
 
         def new_function(*args, **kwargs):
             args = tuple([arg.group if isinstance(arg, ReloadableProcessGroup) else arg for arg in args])
             kwargs = {k: (v.group if isinstance(v, ReloadableProcessGroup) else v) for k, v in kwargs.items()}
-            with _wrap_low_level_call():
+            check_memory = True if op_name is None else _should_check_memory_for_comm(op_name)
+            with _wrap_low_level_call(check_memory=check_memory):
                 return func(*args, **kwargs)
 
         return new_function
@@ -77,18 +93,18 @@ def new_function(*args, **kwargs):
 
     dist.all_reduce = get_new_comm_function(dist.all_reduce)
     dist.all_gather = get_new_comm_function(dist.all_gather)
-    dist.all_gather_into_tensor = get_new_comm_function(dist.all_gather_into_tensor)
+    dist.all_gather_into_tensor = get_new_comm_function(dist.all_gather_into_tensor, "all_gather_into_tensor")
     dist.all_gather_object = get_new_comm_function(dist.all_gather_object)
     dist.all_to_all = get_new_comm_function(dist.all_to_all)
-    dist.all_to_all_single = get_new_comm_function(dist.all_to_all_single)
+    dist.all_to_all_single = get_new_comm_function(dist.all_to_all_single, "all_to_all_single")
     dist.broadcast = get_new_comm_function(dist.broadcast)
-    dist.broadcast_object_list = get_new_comm_function(dist.broadcast_object_list)
+    dist.broadcast_object_list = get_new_comm_function(dist.broadcast_object_list, "broadcast_object_list")
     dist.reduce = get_new_comm_function(dist.reduce)
     dist.reduce_scatter = get_new_comm_function(dist.reduce_scatter)
-    dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor)
+    dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor, "reduce_scatter_tensor")
     dist.scatter = get_new_comm_function(dist.scatter)
     dist.gather = get_new_comm_function(dist.gather)
-    dist.barrier = get_new_comm_function(dist.barrier)
+    dist.barrier = get_new_comm_function(dist.barrier, "barrier")
     dist.send = get_new_comm_function(dist.send)
     dist.recv = get_new_comm_function(dist.recv)
     dist._coalescing_manager = get_new_comm_function(dist._coalescing_manager)
@@ -97,8 +113,8 @@ def new_function(*args, **kwargs):
     old_isend = dist.isend
     old_irecv = dist.irecv
 
-    dist.isend = get_new_comm_function(dist.isend)
-    dist.irecv = get_new_comm_function(dist.irecv)
+    dist.isend = get_new_comm_function(dist.isend, "isend")
+    dist.irecv = get_new_comm_function(dist.irecv, "irecv")
 
     def get_new_p2pop_function(func):
         def new_function(*args, **kwargs):
@@ -191,7 +207,7 @@ def _fwd(self, method, *args, **kwargs):
         inner = self.group
         if inner is None:
             raise RuntimeError("ReloadableProcessGroup: inner PG is None, call reload() first.")
-        with _wrap_low_level_call():
+        with _wrap_low_level_call(check_memory=_should_check_memory_for_comm(method)):
             return getattr(inner, method)(*args, **kwargs)
 
     def _fwd_query(self, method, *args, **kwargs):
@@ -293,11 +309,12 @@ def reload_process_groups():
 
 
 @contextmanager
-def _wrap_low_level_call():
+def _wrap_low_level_call(check_memory=True):
     try:
-        mem_info = available_memory()
-        if mem_info["free_GB"] < 3:
-            clear_memory()
+        if check_memory:
+            mem_info = available_memory()
+            if mem_info["free_GB"] < 3:
+                clear_memory()
         yield
     except Exception as e:
         mem_info = print_memory("after torch distributed error")

slime_plugins/models/qwen3_5.py

@@ -11,11 +11,11 @@
 
 try:
     from fla.modules import FusedRMSNormGated, ShortConvolution
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 except ImportError:
     pass
 
 from .hf_attention import HuggingfaceAttention, _load_hf_config
+from .qwen_gdn_backend import get_chunk_gated_delta_rule
 
 
 def _get_text_config(hf_config):
@@ -33,8 +33,10 @@ class Qwen3_5GatedDeltaNet(nn.Module):
     separate in_proj_qkv (for Q,K,V) and in_proj_z (for Z).
     """
 
-    def __init__(self, config, layer_idx: int):
+    def __init__(self, config, layer_idx: int, args=None):
         super().__init__()
+        self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla")
+        self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend)
         self.hidden_size = config.hidden_size
         self.num_v_heads = config.linear_num_value_heads
         self.num_k_heads = config.linear_num_key_heads
@@ -118,7 +120,14 @@ def forward(
             query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
             key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
 
-        core_attn_out, last_recurrent_state = chunk_gated_delta_rule(
+        if self.gdn_backend == "flashqla":
+            query = query.contiguous()
+            key = key.contiguous()
+            value = value.contiguous()
+            g = g.contiguous()
+            beta = beta.contiguous()
+
+        core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
             query,
             key,
             value,
@@ -162,7 +171,7 @@ def __init__(
         self.hf_config = _get_text_config(self.hf_config)
         self.hf_config._attn_implementation = "flash_attention_2"
 
-        self.linear_attn = Qwen3_5GatedDeltaNet(self.hf_config, self.hf_layer_idx)
+        self.linear_attn = Qwen3_5GatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args)
 
         # Use a simple RMSNorm
         try:

slime_plugins/models/qwen3_next.py

@@ -13,12 +13,12 @@
 
 try:
     from fla.modules import FusedRMSNormGated, ShortConvolution
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
     from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm
 except ImportError:
     pass
 
 from .hf_attention import HuggingfaceAttention
+from .qwen_gdn_backend import get_chunk_gated_delta_rule
 
 
 # adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564
@@ -27,8 +27,10 @@ class Qwen3NextGatedDeltaNet(nn.Module):
     Qwen3NextGatedDeltaNet with varlen support
     """
 
-    def __init__(self, config, layer_idx: int):
+    def __init__(self, config, layer_idx: int, args=None):
         super().__init__()
+        self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla")
+        self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend)
         self.hidden_size = config.hidden_size
         self.num_v_heads = config.linear_num_value_heads
         self.num_k_heads = config.linear_num_key_heads
@@ -140,7 +142,14 @@ def forward(
             query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
             key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
 
-        core_attn_out, last_recurrent_state = chunk_gated_delta_rule(
+        if self.gdn_backend == "flashqla":
+            query = query.contiguous()
+            key = key.contiguous()
+            value = value.contiguous()
+            g = g.contiguous()
+            beta = beta.contiguous()
+
+        core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
             query,
             key,
             value,
@@ -183,7 +192,7 @@ def __init__(
         if Qwen3NextAttention is None:
             raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.")
 
-        self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx)
+        self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args)
         self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps)
 
     def hf_forward(self, hidden_states, packed_seq_params):

slime_plugins/models/qwen_gdn_backend.py

@@ -0,0 +1,46 @@
+import torch
+
+
+def _parse_version(version):
+    version = version.split("+", 1)[0]
+    parts = version.split(".")
+    major = int(parts[0])
+    minor = int(parts[1]) if len(parts) > 1 else 0
+    return major, minor
+
+
+def _validate_flashqla_runtime():
+    if _parse_version(torch.__version__) < (2, 8):
+        raise RuntimeError(f"FlashQLA backend requires PyTorch 2.8 or newer, got PyTorch {torch.__version__}.")
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("FlashQLA backend requires CUDA and an NVIDIA SM90 GPU.")
+
+    major, minor = torch.cuda.get_device_capability()
+    if (major, minor) < (9, 0):
+        raise RuntimeError(f"FlashQLA backend requires NVIDIA SM90 or newer, got sm{major}{minor}.")
+
+    cuda_version = torch.version.cuda
+    if cuda_version is not None and _parse_version(cuda_version) < (12, 8):
+        raise RuntimeError(f"FlashQLA backend requires CUDA 12.8 or newer, got CUDA {cuda_version}.")
+
+
+def get_chunk_gated_delta_rule(backend: str):
+    if backend == "fla":
+        try:
+            from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+        except ImportError as exc:
+            raise ImportError("Qwen GDN backend 'fla' requires flash-linear-attention.") from exc
+        return chunk_gated_delta_rule
+
+    if backend == "flashqla":
+        try:
+            from flash_qla import chunk_gated_delta_rule
+        except ImportError as exc:
+            raise ImportError(
+                "Qwen GDN backend 'flashqla' requires FlashQLA. " "Install it from https://github.com/QwenLM/FlashQLA."
+            ) from exc
+        _validate_flashqla_runtime()
+        return chunk_gated_delta_rule
+
+    raise ValueError(f"Unsupported Qwen GDN backend: {backend}")