fix npu qwen3.5 fla patch and use fla GDN

[addsubmuldiv]

PR Type

• Bug Fix
• New Feature
• Documentation Updates
• Support for More Models or Datasets

PR Information

Background

When Qwen3.5 training enables packing / padding-free mode, the linear-attention layers need to invoke the causal-convolution and gated-delta-rule operators provided by FLA.

This path currently has two compatibility issues on Ascend NPUs:

1. Transformers’ is_flash_linear_attention_available() check includes CUDA-related conditions. Even when the relevant FLA operators can be imported and run correctly on NPU, Transformers may still treat FLA as unavailable.
2. MindSpeed initialization and runtime repatch() may replace chunk_gated_delta_rule at FLA’s public entry point. In addition, mcore-bridge may cache the replaced callable early, preventing stable use of FLA’s native Triton-Ascend GDN implementation during training.

Previously, ms-swift also maintained its own chunk_gated_delta_rule copied from MindSpeed. This implementation duplicates functionality already provided upstream by FLA and increases future maintenance and version-adaptation costs.

Main Changes

1. Add FLA availability checks for Ascend NPUs

A NPU-specific FLA availability check is added. It verifies whether FLA is available by attempting to import the following public entry points:

• fla.modules.convolution.causal_conv1d
• fla.ops.gated_delta_rule.chunk_gated_delta_rule

When both entry points can be imported successfully, the FLA fast path is enabled for Transformers’ Qwen3.5 and Qwen3.5-MoE modules.

This only supplements the availability check in NPU environments and does not change the existing Transformers behavior on GPUs or other devices.

2. Use upstream FLA operators in the Qwen3.5 packing path

• The Qwen3.5 packing/sequence-parallel forward path retries importing the FLA causal-convolution and GDN operators, preventing None values from being permanently cached when the initial module import failed because of the availability check.
• The causal-convolution path uses FLA’s public causal_conv1d interface.
• Remove the built-in swift.model.chunk_gated_delta_rule implementation, so ms-swift no longer maintains a copy of MindSpeed’s GDN code.
• The NPU packing path prioritizes the Triton-Ascend backend provided upstream by FLA.

The expected call chain is:

Qwen3.5
    -> fla.modules.convolution.causal_conv1d
    -> fla.ops.gated_delta_rule.chunk_gated_delta_rule
    -> FLA Triton-Ascend backend

3. Restore FLA GDN after MindSpeed patching

MindSpeed initialization and training-argument repatch() may replace:

fla.ops.gated_delta_rule.chunk_gated_delta_rule

This PR adds restoration logic:

• Prefer the original FLA callable saved as orig_func by MindSpeedPatchesManager;
• If the patch manager does not contain a usable original callable, re-import it directly from FLA’s chunk module;
• Perform one restoration during Megatron environment initialization;
• Perform restoration again after MindSpeed training arguments call repatch();
• If mcore-bridge has already been imported and cached the GDN callable, refresh its cached reference as well;
• After successful restoration, log the final implementation module and source path to make the actual runtime path easy to verify.

Expected log output:

Using upstream FLA chunk_gated_delta_rule for Megatron GDN:
module=fla.ops.gated_delta_rule.chunk
source=fla/ops/gated_delta_rule/chunk.py

4. Keep FLA as an optional dependency for regular Megatron training

FLA is the preferred implementation for the current Qwen3.5 NPU packing path, but it does not become a mandatory dependency for all NPU Megatron training workloads.

If FLA GDN cannot be imported or restored:

• Log a warning with the specific failure reason;
• Do not proactively overwrite the current MindSpeed/Megatron GDN implementation;
• Preserve the existing runtime dispatch behavior;
• If the current implementation does not support packed cu_seqlens, fail at the actual GDN invocation.

Therefore, regular non-packing Megatron training will not exit during initialization solely because FLA is unavailable. Whether packing/variable-length GDN is available depends on the implementation actually selected at runtime.

5. Keep Qwen3.5’s native normalization implementation on NPU

FLA’s FusedRMSNormGated currently contains:

torch.cuda.current_device()

This triggers CUDA-related initialization issues in NPU environments.

Therefore, this PR continues to use Qwen3.5’s native PyTorch gated RMSNorm implementation on NPU. This only affects normalization and does not prevent GDN from using FLA’s Triton-Ascend backend.

6. Documentation updates

• Update the Chinese and English NPU support documentation;
• Update the Qwen3.5 best-practices documentation;
• Explain the runtime dispatch relationship among FLA, MindSpeed, and mcore-bridge;
• Explain the behavior when FLA is unavailable and the current GDN implementation is retained;
• Add the validation environment, verified configurations, and current feature limitations;
• Recommend that Ascend NPU users install the latest version from FLA’s main branch to obtain Triton-Ascend GDN support.

Validation

Transformers/FSDP Training

The validated configuration is as follows:

• Hardware: 8 × Atlas 900 A2
• Model: Qwen3.5-4B
• Training method: LoRA + FSDP
• Precision: BF16
• Dataset: alpaca-gpt4-data-zh
• packing=true
• max_length=512
• per_device_train_batch_size=1
• gradient_accumulation_steps=1
• TP=1
• CP=1
• Training duration: 300 steps

During training, both loss and gradient norm remained finite, and model saving completed successfully.

Runtime logs confirmed that the final GDN callable came from:

module=fla.ops.gated_delta_rule.chunk
source=fla/ops/gated_delta_rule/chunk.py

GPU/NPU Numerical Comparison

A comparison experiment was performed on a GPU server using the same model, dataset, and training hyperparameters.

The GPU and NPU loss values and convergence trends were generally aligned.

Under the validated configuration above, Qwen3.5-4B can complete packing training on Ascend NPUs using FLA’s native Triton-Ascend GDN implementation. The forward pass, backward pass, and model-saving pipeline all run successfully, while maintaining broadly consistent loss values and convergence trends with GPU training.

[gemini-code-assist]

Code Review

This pull request introduces helper functions to dynamically import flash linear attention kernels and patches the transformers' flash linear attention availability check for NPU environments. The review feedback suggests retrieving patched module-level attributes via sys.modules since they are not directly available on the model instance, and adding fallbacks to the standalone causal_conv1d package when fla is not installed to improve NPU compatibility.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since mod is an instance of Qwen3_5GatedDeltaNet, module-level patched variables like chunk_gated_delta_rule (patched via setattr(module, ...) in patch_qwen3_5_chunk_gated_delta_rule_with_mindspeed) are not attributes of the instance mod. Therefore, getattr(mod, 'chunk_gated_delta_rule', None) will return None. By looking up the module of mod via sys.modules, we can correctly retrieve the patched mindspeed version of chunk_gated_delta_rule and causal_conv1d, preventing an ImportError when fla is not installed.

Suggested change

	_try_import_flash_linear_attention_kernels()
	mod._swift_fla_causal_conv1d_fn = getattr(mod, 'causal_conv1d_fn', None) or causal_conv1d
	mod.chunk_gated_delta_rule = getattr(mod, 'chunk_gated_delta_rule', None) or chunk_gated_delta_rule
	import sys
	_try_import_flash_linear_attention_kernels()
	module = sys.modules.get(mod.__module__) if mod else None
	mod._swift_fla_causal_conv1d_fn = (
	getattr(mod, 'causal_conv1d_fn', None)
	or (getattr(module, 'causal_conv1d_fn', None) if module else None)
	or (getattr(module, 'causal_conv1d', None) if module else None)
	or causal_conv1d
	)
	mod.chunk_gated_delta_rule = (
	getattr(mod, 'chunk_gated_delta_rule', None)
	or (getattr(module, 'chunk_gated_delta_rule', None) if module else None)
	or chunk_gated_delta_rule
	)

[gemini-code-assist]

If fla is not installed, we can fallback to importing causal_conv1d_fn from the standalone causal_conv1d package, which is commonly installed on NPU environments.

Suggested change

	if causal_conv1d is None:
	try:
	from fla.modules.convolution import causal_conv1d as _causal_conv1d
	causal_conv1d = _causal_conv1d
	except Exception:
	pass
	if causal_conv1d is None:
	try:
	from fla.modules.convolution import causal_conv1d as _causal_conv1d
	causal_conv1d = _causal_conv1d
	except Exception:
	try:
	from causal_conv1d import causal_conv1d_fn as _causal_conv1d
	causal_conv1d = _causal_conv1d
	except Exception:
	pass

[gemini-code-assist]

If fla is not installed, but the standalone causal_conv1d package is installed, we can still support flash linear attention on NPU because we have the embedded mindspeed version of chunk_gated_delta_rule. Adding this fallback allows the NPU patch to be successfully enabled without requiring fla to be installed.

Suggested change

	def _is_flash_linear_attention_importable_on_npu() -> bool:
	try:
	from fla.modules.convolution import causal_conv1d # noqa: F401
	from fla.ops.gated_delta_rule import chunk_gated_delta_rule # noqa: F401
	return True
	except Exception:
	return False
	def _is_flash_linear_attention_importable_on_npu() -> bool:
	try:
	from fla.modules.convolution import causal_conv1d # noqa: F401
	from fla.ops.gated_delta_rule import chunk_gated_delta_rule # noqa: F401
	return True
	except Exception:
	try:
	from causal_conv1d import causal_conv1d_fn # noqa: F401
	return True
	except Exception:
	return False

[Copilot]

Pull request overview

This PR updates the Ascend NPU Qwen3.5 linear-attention “packing/padding-free” path to rely on upstream flash-linear-attention (FLA) public operators (causal-conv + GDN), adds NPU-specific FLA availability probing for Transformers, and introduces restoration logic to keep FLA’s native Triton-Ascend GDN callable from being replaced by MindSpeed patching/caching. It also removes ms-swift’s embedded MindSpeed GDN copy and updates related NPU/Qwen3.5 documentation.

Changes:

• Add an NPU-only supplement to Transformers’ is_flash_linear_attention_available() that treats FLA as available when its public causal-conv and GDN entry points are importable.
• Prefer upstream FLA GDN in Megatron/MindSpeed flows by restoring fla.ops.gated_delta_rule.chunk_gated_delta_rule (and refreshing mcore-bridge’s cached reference when present).
• Remove swift.model.chunk_gated_delta_rule and update Qwen3.5 docs/NPU docs to reflect the new runtime dispatch and recommended FLA installation.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
swift/model/npu_patcher.py	Re-export the new MindSpeed→FLA GDN restoration helper for NPU flows.
swift/model/npu_patch/init.py	Export patch_mindspeed_fla_gdn_implementation from the NPU patch package.
swift/model/npu_patch/model.py	Add NPU-only FLA import-based availability probing and avoid CUDA-init fused norm on NPU.
swift/model/npu_patch/mindspeed.py	Add best-effort restoration logic for upstream FLA GDN and refresh mcore-bridge cache.
swift/model/models/qwen.py	Avoid permanently caching missing FLA kernels; retry imports at runtime for packing/sequence-parallel.
swift/model/chunk_gated_delta_rule.py	Remove embedded MindSpeed-derived GDN implementation (duplicate maintenance).
swift/megatron/pipelines/train/sft.py	Restore upstream FLA GDN after MindSpeed repatch() in the Megatron SFT pipeline.
swift/megatron/init.py	Restore upstream FLA GDN during Megatron environment initialization on NPU.
docs/source/BestPractices/Qwen3_5-Best-Practice.md	Update NPU guidance and call-path notes for FLA GDN usage.
docs/source/BestPractices/NPU-support.md	Update NPU Qwen3.5 patch notes to reflect upstream FLA Triton-Ascend GDN usage.
docs/source_en/BestPractices/Qwen3_5-Best-Practice.md	English equivalent documentation updates for FLA GDN on NPU.
docs/source_en/BestPractices/NPU-support.md	English equivalent NPU Qwen3.5 patch notes updates.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.