feta: kimi k3 support ulysess cp-MindSpeed-MM-AtomGit

已合并

feta: kimi k3 support ulysess cp #2933

suis-moi创建于 18 天前从suis-moi/MindSpeed-MM： master 合入到 Ascend/MindSpeed-MM： master

feta: kimi k3 support ulysess cp #2933

已合并

suis-moi创建于 18 天前

从suis-moi/MindSpeed-MM master合入到Ascend/MindSpeed-MM master

共 6 个文件变更+377-32

Mexamples/kimi_k3/README.md+37-2

	@@ -14,6 +14,7 @@
14	- [1. 环境准备](#1-环境准备)	14	- [1. 环境准备](#1-环境准备)
15	- [2. 环境搭建](#2-环境搭建)	15	- [2. 环境搭建](#2-环境搭建)
16	- [3. 安装配套版本的Triton-Ascend](#3-安装配套版本的triton-ascend)	16	- [3. 安装配套版本的Triton-Ascend](#3-安装配套版本的triton-ascend)
		17	+ - [4. 安装fla-npu以适配AscendC](#4-安装fla-npu以适配ascendc)
17	- [数据集准备及处理](#数据集准备及处理)	18	- [数据集准备及处理](#数据集准备及处理)
18	- [训练](#训练)	19	- [训练](#训练)
19	- [1. 准备工作](#1-准备工作)	20	- [1. 准备工作](#1-准备工作)
	@@ -99,6 +100,39 @@ cp -f ${MM_PATH}/mindspeed_mm/fsdp/ops/kda/triton_ascend/chunk.py \
99	pip install -e .	100	pip install -e .
100	```	101	```
101		102
		103	+### 4. 安装fla-npu以适配AscendC
		104	+
		105	+Kimi-K3 的 KDA 短卷积算子（`causal_conv1d_implementation: ascendc`）基于 fla-npu 的 AscendC 融合算子实现，需要安装 fla-npu。
		106	+
		107	+拉取flash-linear-attention-npu代码仓，并进入代码仓根目录，切到对应commitID
		108	+
		109	+```bash
		110	+git clone https://github.com/flashserve/flash-linear-attention-npu
		111	+cd flash-linear-attention-npu
		112	+git checkout c2e3d83f
		113	+```
		114	+
		115	+安装步骤：可参考fla-npu仓README：[flash-linear-attention-npu](https://github.com/flashserve/flash-linear-attention-npu/blob/release/v26.1.0/README.md)
		116	+
		117	+推荐使用以下安装命令
		118	+
		119	+```shell
		120	+# source 实际的cann路径
		121	+source /usr/local/Ascend/cann/set_env.sh
		122	+
		123	+# 编译算子 run 包，--soc 需指定为当前机器芯片类型 {ascend910b/ascend910_93/ascend950}
		124	+bash build.sh --soc=ascend910b --pkg --vendor_name=fla_npu
		125	+bash build_out/fla-npu-*.run
		126	+cd torch_custom/fla_npu/
		127	+bash build.sh
		128	+```
		129	+
		130	+检验fla_npu是否安装成功
		131	+
		132	+```bash
		133	+pip list \| grep fla_npu
		134	+```
		135	+
102	---	136	---
103		137
104	<a id="jump2"></a>	138	<a id="jump2"></a>
	@@ -162,7 +196,7 @@ NODE_RANK: 当前节点序号
162		196
163	\| 配置项 \| 配置路径 \| 参数说明 \| 调整说明 \|	197	\| 配置项 \| 配置路径 \| 参数说明 \| 调整说明 \|
164	\|--------\|----------\|----------\|----------\|	198	\|--------\|----------\|----------\|----------\|
165	-\| `ulysses_parallel_size` \| `parallel` \| ulysses-cp 并行度 \| 暂不支持，开发中 \|	199	+\| `ulysses_parallel_size` \| `parallel` \| ulysses-cp 并行度 \| 值为1时不开启，根据实际情况调整；\|
166	\| `expert_parallel_size` \| `parallel` \| EP专家并行度 \| 值为1时不开启，仅对MoE模型生效 \|	200	\| `expert_parallel_size` \| `parallel` \| EP专家并行度 \| 值为1时不开启，仅对MoE模型生效 \|
167	\| `ep_plan` \| `parallel` \| EP调度策略配置 \| 包含`dispatcher`、`use_npu_fused_ops`等子字段，`dispatcher`可选`alltoall` \|	201	\| `ep_plan` \| `parallel` \| EP调度策略配置 \| 包含`dispatcher`、`use_npu_fused_ops`等子字段，`dispatcher`可选`alltoall` \|
168	\| `num_to_forward_prefetch` \| `parallel->fsdp_plan` \| 前向计算时预取后续层参数 \| 减少通信等待开销 \|	202	\| `num_to_forward_prefetch` \| `parallel->fsdp_plan` \| 前向计算时预取后续层参数 \| 减少通信等待开销 \|
	@@ -170,6 +204,7 @@ NODE_RANK: 当前节点序号
170	\| `enable_preload` \| `data->dataloader_param` \| 数据预加载开关 \| 开启后数据加载与计算重叠，减少训练等待时间 \|	204	\| `enable_preload` \| `data->dataloader_param` \| 数据预加载开关 \| 开启后数据加载与计算重叠，减少训练等待时间 \|
171	\| `use_grouped_expert_matmul` \| `model` \| MoE专家分组矩阵乘融合算子开关 \| 开启后使用NPU融合算子加速MoE专家计算 \|	205	\| `use_grouped_expert_matmul` \| `model` \| MoE专家分组矩阵乘融合算子开关 \| 开启后使用NPU融合算子加速MoE专家计算 \|
172	\| `kda_implementation` \| `model` \| KDA算子实现选择 \| `fused`: triton-ascend-kernels融合大算子（默认）<br>`naive`: 仓内小算子实现，可用于功能对齐验证 \|	206	\| `kda_implementation` \| `model` \| KDA算子实现选择 \| `fused`: triton-ascend-kernels融合大算子（默认）<br>`naive`: 仓内小算子实现，可用于功能对齐验证 \|
		207	+\| `causal_conv1d_implementation` \| `model` \| KDA短卷积算子实现选择 \| `triton`: triton实现（默认）<br>`ascendc`: fla_npu AscendC融合算子（仅NPU），需安装fla-npu，参考[安装fla-npu以适配AscendC](#4-安装fla-npu以适配ascendc) \|
		yaoyaoxu12 天前使用ascendc融合算子，性能收益有多少？建议提供基线数据，以作参考 suis-moi 11 天前评论： 11 天前评论：
173	\| `skip_flash_attn_recompute` \| `model` \| 跳过full attention层flash attention重计算 \| 选择性重计算，需同时使能重计算和`enable_activation_offload` \|	208	\| `skip_flash_attn_recompute` \| `model` \| 跳过full attention层flash attention重计算 \| 选择性重计算，需同时使能重计算和`enable_activation_offload` \|
174	\| `skip_kda_recompute` \| `model` \| 跳过linear attention层KDA重计算 \| 选择性重计算，需同时使能重计算和`enable_activation_offload` \|	209	\| `skip_kda_recompute` \| `model` \| 跳过linear attention层KDA重计算 \| 选择性重计算，需同时使能重计算和`enable_activation_offload` \|
175	\| `recompute` \| `features` \| 重计算开关 \| 开启后可以节省显存占用 \|	210	\| `recompute` \| `features` \| 重计算开关 \| 开启后可以节省显存占用 \|
	@@ -260,7 +295,7 @@ NNODES: 一共几个节点
260	- 参考配置：当前 A3 单节点可配置 `num_hidden_layers=16`、`num_experts=32`；	295	- 参考配置：当前 A3 单节点可配置 `num_hidden_layers=16`、`num_experts=32`；
261	- 序列长度：mbs=1时支持6k序列长度以下；	296	- 序列长度：mbs=1时支持6k序列长度以下；
262	- 权重加载：当前采用随机初始化权重（加载预训练权重能力后续支持）；	297	- 权重加载：当前采用随机初始化权重（加载预训练权重能力后续支持）；
263	-- CP 长序列训练：暂不支持，开发中。	298	+- CP 长序列训练：支持 ulysses-cp 长序列训练，配置 `kimik3_config.yaml` 中 `parallel->ulysses_parallel_size` 调整并行度（值为1时不开启）。
264		299
265	<a id="jump3.4"></a>	300	<a id="jump3.4"></a>
266		301

Mexamples/kimi_k3/kimik3_config.yaml+2-0

	@@ -89,6 +89,8 @@ model:
89	use_grouped_expert_matmul: true	89	use_grouped_expert_matmul: true
90	# KDA算子实现：fused（triton-ascend-kernels融合大算子）/ naive（仓内小算子）	90	# KDA算子实现：fused（triton-ascend-kernels融合大算子）/ naive（仓内小算子）
91	kda_implementation: fused	91	kda_implementation: fused
		92	+ # KDA短卷积算子实现：triton / ascendc
		93	+ causal_conv1d_implementation: triton
92	# 选择性重计算：attention 相关融合算子（ViT/MLA 的 FA、KDA）不参与重算，	94	# 选择性重计算：attention 相关融合算子（ViT/MLA 的 FA、KDA）不参与重算，
93	skip_flash_attn_recompute: true	95	skip_flash_attn_recompute: true
94	skip_kda_recompute: true	96	skip_kda_recompute: true

Mmindspeed_mm/fsdp/models/kimi_k3/__init__.py+10-0

	@@ -84,6 +84,16 @@ class KimiK3ForConditionalGeneration(WeightInitMixin, _KimiK3ForConditionalGener
84	)	84	)
85	transformer_config.text_config.kda_implementation = kda_implementation	85	transformer_config.text_config.kda_implementation = kda_implementation
86		86
		87	+ # Causal conv1d kernel selection for KDA short convolutions: 'triton' (default)
		88	+ # or 'ascendc' (AscendC fused op from fla_npu, same as qwen3_5, NPU only).
		89	+ causal_conv1d_implementation = getattr(model_args, "causal_conv1d_implementation", "triton")
		90	+ if causal_conv1d_implementation not in ("triton", "ascendc"):
		91	+ raise ValueError(
		92	+ f"Unsupported causal_conv1d_implementation: {causal_conv1d_implementation}. "
		93	+ "Expected 'triton' or 'ascendc'."
		94	+ )
		95	+ transformer_config.text_config.causal_conv1d_implementation = causal_conv1d_implementation
		96	+
87	return transformer_config	97	return transformer_config
88		98
89	def tie_weights(self, args, *kwargs):	99	def tie_weights(self, args, *kwargs):

Mmindspeed_mm/fsdp/models/kimi_k3/modeling_kimi_k3.py+87-1

@@ -26,6 +26,7 @@ from typing import Optional

26 26

27 import numpy as np 27 import numpy as np

28 import torch 28 import torch

29 +import torch.distributed as dist

29 import torch.nn as nn 30 import torch.nn as nn

30 import torch.nn.functional as F 31 import torch.nn.functional as F

31 from transformers import activations 32 from transformers import activations

@@ -44,7 +45,7 @@ from transformers.models.llava.modeling_llava import \

44 from transformers.utils import is_flash_attn_2_available 45 from transformers.utils import is_flash_attn_2_available

45 46

46 from .configuration_kimi_k3 import KimiK3Config 47 from .configuration_kimi_k3 import KimiK3Config

47 -from .modeling_kimi_linear import KimiLinearForCausalLM 48 +from .modeling_kimi_linear import KimiLinearForCausalLM, set_seq_len, get_seq_len

48 49

49 # Flash attention imports 50 # Flash attention imports

50 if is_flash_attn_2_available(): 51 if is_flash_attn_2_available():

@@ -53,6 +54,13 @@ else:

53 flash_attn_varlen_func = None 54 flash_attn_varlen_func = None

54 from mindspeed_mm.fsdp.utils.device import IS_NPU_AVAILABLE 55 from mindspeed_mm.fsdp.utils.device import IS_NPU_AVAILABLE

55 from mindspeed_mm.fsdp.loss.loss_func import build_loss_func 56 from mindspeed_mm.fsdp.loss.loss_func import build_loss_func

57 +from mindspeed_mm.fsdp.distributed.parallel_state import get_parallel_state

58 +from mindspeed_mm.fsdp.distributed.context_parallel.communication import (

59 + all_to_all,

60 + gather_forward_split_backward,

61 + packed_data_split_forward_gather_backward_with_cp,

62 +)

63 +from mindspeed_mm.fsdp.distributed.context_parallel.utils import cal_split_sizes

56 64

57 if IS_NPU_AVAILABLE: 65 if IS_NPU_AVAILABLE:

58 import torch_npu 66 import torch_npu

@@ -106,6 +114,38 @@ def multihead_attention(

106 output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing, 114 output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,

107 where dim = num_heads * head_dim 115 where dim = num_heads * head_dim

108 """ 116 """

117 +

118 + # Modification start, ulysses cp

119 + ps = get_parallel_state() if dist.is_initialized() else None

120 + is_ulysses_enabled = ps is not None and ps.is_ulysses_enable()

121 + total_seq_len = get_seq_len("visual")

122 + head_num = q.shape[1]

123 + kv_head_num = k.shape[1]

124 +

125 + # ulysses validation

126 + if is_ulysses_enabled:

127 + ulysses_size = ps.get_ulysses_group_size()

128 + if head_num % ulysses_size != 0:

129 + raise ValueError(f"num_query_heads ({head_num}) must be divisible by ulysses_size ({ulysses_size})")

130 + if ulysses_size > kv_head_num:

131 + if ulysses_size % kv_head_num != 0:

132 + raise ValueError(

133 + f"ulysses_size ({ulysses_size}) must be divisible by num_key_value_heads ({kv_head_num})"

134 + )

135 + n_repeat = ulysses_size // kv_head_num

136 + # Shape before: (total_seq_len, kv_head_num, head_dim)

137 + # This repeats the K/V heads (dim 1) to match the ulysses_size (SP world size)

138 + # Shape after: (total_seq_len, kv_head_num * n_repeat, head_dim) where (kv_head_num * n_repeat) == ulysses_size

139 + k = torch.repeat_interleave(k, dim=1, repeats=n_repeat)

140 + v = torch.repeat_interleave(v, dim=1, repeats=n_repeat)

141 +

142 + if is_ulysses_enabled:

143 + q = all_to_all(q, ps.get_ulysses_group(), scatter_dim=1, gather_dim=0, gather_size=total_seq_len)

144 + k = all_to_all(k, ps.get_ulysses_group(), scatter_dim=1, gather_dim=0, gather_size=total_seq_len)

145 + v = all_to_all(v, ps.get_ulysses_group(), scatter_dim=1, gather_dim=0, gather_size=total_seq_len)

146 +

147 + # Modification end, ulysses cp

148 +

109 if IS_NPU_AVAILABLE: 149 if IS_NPU_AVAILABLE:

110 # Modification start 150 # Modification start

111 if skip_recompute: 151 if skip_recompute:

@@ -161,6 +201,9 @@ def multihead_attention(

161 if isinstance(attn_out, tuple): 201 if isinstance(attn_out, tuple):

162 attn_out = attn_out[0] 202 attn_out = attn_out[0]

163 203

204 + if is_ulysses_enabled:

205 + attn_out = all_to_all(attn_out, ps.get_ulysses_group(), scatter_dim=0, gather_dim=1)

206 +

164 attn_out = attn_out.flatten(start_dim=-2) 207 attn_out = attn_out.flatten(start_dim=-2)

165 208

166 return attn_out 209 return attn_out

@@ -692,6 +735,22 @@ class MoonViT3dEncoder(nn.Module):

692 if IS_NPU_AVAILABLE: 735 if IS_NPU_AVAILABLE:

693 cu_seqlens = tuple(cu_seqlens[1:].cpu().numpy().tolist()) 736 cu_seqlens = tuple(cu_seqlens[1:].cpu().numpy().tolist())

694 737

738 + # Modification start: ulysses cp

739 + seq_len, _ = hidden_states.size()

740 + sequence_lengths = torch.repeat_interleave(grid_thws[:, 1] * grid_thws[:, 2], grid_thws[:, 0]).cpu()

741 + set_seq_len("visual", seq_len)

742 +

743 + ps = get_parallel_state() if dist.is_initialized() else None

744 + # Split sequences across context parallel groups for distributed processing

745 + if ps is not None and ps.is_ulysses_enable():

746 + hidden_states = packed_data_split_forward_gather_backward_with_cp(

747 + hidden_states, dim=0, seq_lens=sequence_lengths

748 + )

749 + rope_freqs_cis = packed_data_split_forward_gather_backward_with_cp(

750 + rope_freqs_cis, dim=0, seq_lens=sequence_lengths

751 + )

752 + # Modification end: ulysses cp

753 +

695 cos = rope_freqs_cis.unsqueeze(-2).real.to(torch.float32).repeat_interleave(2, dim=-1).contiguous() 754 cos = rope_freqs_cis.unsqueeze(-2).real.to(torch.float32).repeat_interleave(2, dim=-1).contiguous()

696 sin = rope_freqs_cis.unsqueeze(-2).imag.to(torch.float32).repeat_interleave(2, dim=-1).contiguous() 755 sin = rope_freqs_cis.unsqueeze(-2).imag.to(torch.float32).repeat_interleave(2, dim=-1).contiguous()

697 set_global_param("cos", cos) 756 set_global_param("cos", cos)

@@ -703,6 +762,19 @@ class MoonViT3dEncoder(nn.Module):

703 max_seqlen, 762 max_seqlen,

704 rope_freqs_cis=rope_freqs_cis) 763 rope_freqs_cis=rope_freqs_cis)

705 764

765 + # Modification start: ulysses cp

766 + ps = get_parallel_state() if dist.is_initialized() else None

767 + if ps is not None and ps.is_ulysses_enable():

768 + gather_sizes = cal_split_sizes(get_seq_len("visual"), ps.get_ulysses_group_size())

769 + hidden_states = gather_forward_split_backward(

770 + hidden_states,

771 + ps.get_ulysses_group(),

772 + dim=0,

773 + grad_scale="up",

774 + gather_sizes=gather_sizes,

775 + )

776 + # Modification end: ulysses cp

777 +

706 hidden_states = self.final_layernorm(hidden_states) 778 hidden_states = self.final_layernorm(hidden_states)

707 return hidden_states 779 return hidden_states

708 780

@@ -1267,6 +1339,13 @@ class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel):

1267 1339

1268 # Chunk loss needs the full hidden-state sequence; force it on if enabled. 1340 # Chunk loss needs the full hidden-state sequence; force it on if enabled.

1269 enable_chunk_loss = getattr(self, "enable_chunk_loss", False) 1341 enable_chunk_loss = getattr(self, "enable_chunk_loss", False)

1342 + # The plain CE loss path below shifts labels within the local sequence

1343 + # shard, which is incorrect under ulysses cp, so chunk loss is required.

1344 + ps = get_parallel_state() if dist.is_initialized() else None

1345 + if ps is not None and ps.is_ulysses_enable() and not enable_chunk_loss:

1346 + raise ValueError(

1347 + "ulysses cp requires chunk loss; enable chunk loss or set ulysses_parallel_size to 1."

1348 + )

1270 if enable_chunk_loss: 1349 if enable_chunk_loss:

1271 output_hidden_states = True 1350 output_hidden_states = True

1272 1351

@@ -1377,6 +1456,13 @@ class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel):

1377 shift_labels.view(-1).to(shift_logits.device), 1456 shift_labels.view(-1).to(shift_logits.device),

1378 ) 1457 )

1379 1458

1459 + # Modification start: ulysses cp, gather the loss computed on the

1460 + # sequence shards back across the CP group before reducing.

1461 + if ps is not None and ps.is_cp_enable():

犀牛12 天前

上面有一个if labels is not None:分支, 如果不进这个分支,loss为空. 可能训练场景下不涉及该问题.

suis-moi

11 天前评论：

1462 + loss = gather_forward_split_backward(loss.unsqueeze(0), ps.get_cp_group(), dim=0)

cxiaolong12 天前

loss不用gather，从头切到尾就行，否则会多一些通信

suis-moi

11 天前评论：

1463 + loss = loss.sum()

atomgit-bot18 天前

🟡 Medium Priority

变更行（L1457-1460）：新增了 Ulysses CP loss 汇集逻辑：

受影响行为：当 is_cp_enable() 为 True 时，loss.unsqueeze(0) 无条件执行。但 loss 可能为 None：

建议：在 CP loss 汇集前增加 loss is not None 判断，避免在 eval/inference（无 labels）场景下对 None 调用 .unsqueeze(0) 崩溃。

改动建议

1463	+ if ps is not None and ps.is_cp_enable() and loss is not None:
1464	+ loss = gather_forward_split_backward(loss.unsqueeze(0), ps.get_cp_group(), dim=0)
1463 1465	loss = loss.sum()

应用建议

suis-moi

11 天前评论：

1464 + # Modification end: ulysses cp

1465 +

1380 if not return_dict: 1466 if not return_dict:

1381 output = (logits, ) + outputs[1:] 1467 output = (logits, ) + outputs[1:]

1382 return (loss, ) + output if loss is not None else output 1468 return (loss, ) + output if loss is not None else output

Mmindspeed_mm/fsdp/models/kimi_k3/modeling_kimi_linear.py+187-10

	@@ -23,9 +23,10 @@
23	# limitations under the License.	23	# limitations under the License.
24	import math	24	import math
25	from collections.abc import Callable	25	from collections.abc import Callable
26	-from typing import Any	26	+from typing import Any, List, Optional, Union
27		27
28	import torch	28	import torch
		29	+import torch.distributed as dist
29	import torch.nn.functional as F	30	import torch.nn.functional as F
30	import transformers	31	import transformers
31	from einops import rearrange	32	from einops import rearrange
	@@ -53,6 +54,12 @@ except ImportError:
53	return func	54	return func
54		55
55	from mindspeed_mm.fsdp.utils.device import IS_NPU_AVAILABLE	56	from mindspeed_mm.fsdp.utils.device import IS_NPU_AVAILABLE
		57	+from mindspeed_mm.fsdp.distributed.context_parallel.communication import (
		58	+ all_to_all,
		59	+ split_forward_gather_backward_with_cp,
		60	+)
		61	+from mindspeed_mm.fsdp.distributed.context_parallel.utils import generate_ulysses_cu_seqlen_params
		62	+from mindspeed_mm.fsdp.distributed.parallel_state import get_parallel_state
56		63
57	if IS_NPU_AVAILABLE:	64	if IS_NPU_AVAILABLE:
58	import torch_npu	65	import torch_npu
	@@ -71,6 +78,35 @@ if version.parse(transformers.__version__) < version.parse("4.56.0"):
71		78
72	logger = logging.get_logger(__name__)	79	logger = logging.get_logger(__name__)
73		80
		81	+_TOTAL_SEQ_LEN = None
		82	+_VISUAL_SEQ_LEN = None
		83	+_VISUAL_PER_SEQ_LEN = None
		84	+
		85	+
		86	+def set_seq_len(seq_type: str = None, seq_len: Optional[Union[int, List[int]]] = None) -> None:
		87	+ if seq_type == "total":
		88	+ global _TOTAL_SEQ_LEN
		89	+ _TOTAL_SEQ_LEN = seq_len
		90	+ elif seq_type == "visual":
		91	+ global _VISUAL_SEQ_LEN
		92	+ _VISUAL_SEQ_LEN = seq_len
		93	+ elif seq_type == "per_visual":
		94	+ global _VISUAL_PER_SEQ_LEN
		95	+ _VISUAL_PER_SEQ_LEN = seq_len
		96	+ else:
		97	+ raise ValueError(f"Invalid sequence type: '{seq_type}'. Expected 'total', 'visual' or 'per_visual'.")
		98	+
		99	+
		100	+def get_seq_len(seq_type: str = None) -> int:
		101	+ if seq_type == "total":
		102	+ return _TOTAL_SEQ_LEN
		103	+ elif seq_type == "visual":
		104	+ return _VISUAL_SEQ_LEN
		105	+ elif seq_type == "per_visual":
		106	+ return _VISUAL_PER_SEQ_LEN
		107	+ else:
		108	+ raise ValueError(f"Invalid sequence type: '{seq_type}'. Expected 'total', 'visual' or 'per_visual'.")
		109	+
74		110
75	class KimiK_3_MoeRMSNormGated(nn.Module):	111	class KimiK_3_MoeRMSNormGated(nn.Module):
76	def __init__(self, hidden_size, eps=1e-6, **kwargs):	112	def __init__(self, hidden_size, eps=1e-6, **kwargs):
	@@ -510,6 +546,17 @@ class KimiMLAAttention(nn.Module):
510	value_states = F.pad(	546	value_states = F.pad(
511	value_states, [0, self.q_head_dim - self.v_head_dim])	547	value_states, [0, self.q_head_dim - self.v_head_dim])
512		548
		549	+ # Modification start, Context Parallel
		550	+ # Pass total_seq_len so the patched flash attention performs the Ulysses
		551	+ # all-to-all (scatter heads, gather sequence) internally. The value pad
		552	+ # above happens before that all-to-all, same as the non-CP path.
		553	+ total_seq_len = get_seq_len("total")
		554	+ ps = get_parallel_state() if dist.is_initialized() else None
		555	+ if ps is not None and ps.get_ulysses_group_size() > self.num_key_value_heads:
		556	+ key_states = repeat_kv(key_states, self.num_key_value_groups)
		557	+ value_states = repeat_kv(value_states, self.num_key_value_groups)
		558	+ # Modification end, Context Parallel
		559	+
513	attention_interface: Callable = eager_attention_forward	560	attention_interface: Callable = eager_attention_forward
514	if self.config._attn_implementation != "eager":	561	if self.config._attn_implementation != "eager":
515	attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]	562	attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
	@@ -523,6 +570,7 @@ class KimiMLAAttention(nn.Module):
523	dropout=0.0 if not self.training else self.attention_dropout,	570	dropout=0.0 if not self.training else self.attention_dropout,
524	scaling=self.scaling,	571	scaling=self.scaling,
525	skip_flash_attn_recompute=getattr(self.config, "skip_flash_attn_recompute", False),	572	skip_flash_attn_recompute=getattr(self.config, "skip_flash_attn_recompute", False),
		573	+ total_seq_len=total_seq_len,
526	**kwargs,	574	**kwargs,
527	)	575	)
528		576
	@@ -565,20 +613,27 @@ class KimiDeltaAttention(nn.Module):
565	self.hidden_size, projection_k_size, bias=False)	613	self.hidden_size, projection_k_size, bias=False)
566	self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)	614	self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
567		615
		616	+ self.causal_conv1d_implementation = getattr(config, "causal_conv1d_implementation", "triton")
568	self.q_conv1d = ShortConvolution(	617	self.q_conv1d = ShortConvolution(
569	hidden_size=projection_k_size,	618	hidden_size=projection_k_size,
570	kernel_size=self.conv_size,	619	kernel_size=self.conv_size,
571	activation='silu',	620	activation='silu',
		621	+ implementation=self.causal_conv1d_implementation,
		622	+ head_num=self.num_k_heads,
572	)	623	)
573	self.k_conv1d = ShortConvolution(	624	self.k_conv1d = ShortConvolution(
574	hidden_size=projection_k_size,	625	hidden_size=projection_k_size,
575	kernel_size=self.conv_size,	626	kernel_size=self.conv_size,
576	activation='silu',	627	activation='silu',
		628	+ implementation=self.causal_conv1d_implementation,
		629	+ head_num=self.num_k_heads,
577	)	630	)
578	self.v_conv1d = ShortConvolution(	631	self.v_conv1d = ShortConvolution(
579	hidden_size=projection_size,	632	hidden_size=projection_size,
580	kernel_size=self.conv_size,	633	kernel_size=self.conv_size,
581	activation='silu',	634	activation='silu',
		635	+ implementation=self.causal_conv1d_implementation,
		636	+ head_num=self.num_heads,
582	)	637	)
583		638
584	self.A_log = torch.nn.Parameter(torch.log(torch.empty(	639	self.A_log = torch.nn.Parameter(torch.log(torch.empty(
	@@ -604,6 +659,19 @@ class KimiDeltaAttention(nn.Module):
604	self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')	659	self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
605	self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)	660	self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
606		661
		662	+ @staticmethod
		663	+ def _get_local_conv1d_weight(weight: torch.Tensor, ulysses_rank: int, local_dim: int) -> torch.Tensor:
		664	+ # Modification: shard depthwise conv1d weights to match head-sharded channels
		665	+ # under Ulysses SP. Each ShortConvolution here is per-projection (q/k/v), so a
		666	+ # single contiguous channel slice is enough.
		667	+ if weight.shape[0] < (ulysses_rank + 1) * local_dim:
		668	+ raise ValueError(
		669	+ f"conv1d weight dim ({weight.shape[0]}) is smaller than the requested "
		670	+ f"local slice end ({(ulysses_rank + 1) * local_dim})"
		671	+ )
		672	+ offset = ulysses_rank * local_dim
		673	+ return weight[offset: offset + local_dim]
		674	+
607	def forward(	675	def forward(
608	self,	676	self,
609	hidden_states: torch.Tensor,	677	hidden_states: torch.Tensor,
	@@ -626,12 +694,25 @@ class KimiDeltaAttention(nn.Module):
626	if self.training:	694	if self.training:
627	assert mode == 'chunk', "Only chunk mode is supported in training."	695	assert mode == 'chunk', "Only chunk mode is supported in training."
628		696
		697	+ # Modification start: Ulysses context parallel (LASP) for linear attention.
		698	+ ps = get_parallel_state() if dist.is_initialized() else None
		699	+ is_ulysses_enabled = ps is not None and ps.is_ulysses_enable()
		700	+ # Modification end
		701	+
629	cu_seqlens = kwargs.get('cu_seqlens')	702	cu_seqlens = kwargs.get('cu_seqlens')
630	indices = None	703	indices = None
631	if attention_mask is not None:	704	if attention_mask is not None:
632	- indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])	705	+ if is_ulysses_enabled:
633	- hidden_states = index_first_axis(	706	+ # Under Ulysses SP the sequence is sharded across CP ranks while the
634	- rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)	707	+ # mask spans the full sequence. Compute the unpad indices/cu_seqlens
		708	+ # from the full mask here, but defer the unpadding until after the
		709	+ # all-to-all has gathered the full sequence, so both match the
		710	+ # non-CP semantics exactly.
		711	+ indices, cu_seqlens, _ = get_unpad_data(attention_mask)
		712	+ else:
		713	+ indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
		714	+ hidden_states = index_first_axis(
		715	+ rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
635		716
636	conv_state_q, conv_state_k, conv_state_v = None, None, None	717	conv_state_q, conv_state_k, conv_state_v = None, None, None
637	recurrent_state = None	718	recurrent_state = None
	@@ -641,35 +722,108 @@ class KimiDeltaAttention(nn.Module):
641	self.layer_idx]	722	self.layer_idx]
642	recurrent_state = cache_params.recurrent_states[self.layer_idx]	723	recurrent_state = cache_params.recurrent_states[self.layer_idx]
643		724
		725	+ # Modification start: Ulysses context parallel (LASP) for linear attention.
		726	+ # The all-to-all after the projections gathers the full sequence and scatters
		727	+ # heads, so the causal convolutions and the chunk kernel run on the full
		728	+ # sequence with local heads only.
		729	+ if is_ulysses_enabled:
		730	+ ulysses_group = ps.get_ulysses_group()
		731	+ ulysses_size = ps.get_ulysses_group_size()
		732	+ ulysses_rank = ps.get_ulysses_rank()
		733	+ if self.num_heads % ulysses_size != 0:
		734	+ raise ValueError(
		735	+ f"SP size ({ulysses_size}) must divide num_heads ({self.num_heads}) "
		736	+ "for KimiDeltaAttention LASP"
		737	+ )
		738	+ local_num_heads = self.num_heads // ulysses_size
		739	+ local_key_dim = self.head_k_dim * local_num_heads
		740	+ local_value_dim = self.head_dim * local_num_heads
		741	+ total_seq_len = get_seq_len("total")
		742	+ q_conv_kwargs = {
		743	+ "weight": self._get_local_conv1d_weight(self.q_conv1d.weight, ulysses_rank, local_key_dim)}
		744	+ k_conv_kwargs = {
		745	+ "weight": self._get_local_conv1d_weight(self.k_conv1d.weight, ulysses_rank, local_key_dim)}
		746	+ v_conv_kwargs = {
		747	+ "weight": self._get_local_conv1d_weight(self.v_conv1d.weight, ulysses_rank, local_value_dim)}
		748	+ else:
		749	+ # Keep the default module weights (also keeps the non-NPU fla
		750	+ # ShortConvolution path, which has no `weight` kwarg, intact).
		751	+ q_conv_kwargs, k_conv_kwargs, v_conv_kwargs = {}, {}, {}
		752	+ # Modification end
		753	+
644	q_proj_states = self.q_proj(hidden_states)	754	q_proj_states = self.q_proj(hidden_states)
645	k_proj_states = self.k_proj(hidden_states)	755	k_proj_states = self.k_proj(hidden_states)
646	v_proj_states = self.v_proj(hidden_states)	756	v_proj_states = self.v_proj(hidden_states)
		757	+ # Modification start: LASP all-to-all (scatter heads, gather sequence).
		758	+ if is_ulysses_enabled:
		759	+ q_proj_states = all_to_all(
		760	+ q_proj_states, ulysses_group, scatter_dim=2, gather_dim=1, gather_size=total_seq_len)
		761	+ k_proj_states = all_to_all(
		762	+ k_proj_states, ulysses_group, scatter_dim=2, gather_dim=1, gather_size=total_seq_len)
		763	+ v_proj_states = all_to_all(
		764	+ v_proj_states, ulysses_group, scatter_dim=2, gather_dim=1, gather_size=total_seq_len)
		765	+ if indices is not None:
		766	+ # Unpad the gathered full sequences (deferred from the entry, see
		767	+ # the note above); matches the non-CP unpad semantics exactly.
		768	+ q_proj_states = index_first_axis(
		769	+ rearrange(q_proj_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
		770	+ k_proj_states = index_first_axis(
		771	+ rearrange(k_proj_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
		772	+ v_proj_states = index_first_axis(
		773	+ rearrange(v_proj_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
		774	+ # Modification end
647	q, conv_state_q = self.q_conv1d(	775	q, conv_state_q = self.q_conv1d(
648	x=q_proj_states,	776	x=q_proj_states,
649	cache=conv_state_q,	777	cache=conv_state_q,
650	output_final_state=use_cache,	778	output_final_state=use_cache,
651	cu_seqlens=cu_seqlens,	779	cu_seqlens=cu_seqlens,
		780	+ **q_conv_kwargs,
652	)	781	)
653	k, conv_state_k = self.k_conv1d(	782	k, conv_state_k = self.k_conv1d(
654	x=k_proj_states,	783	x=k_proj_states,
655	cache=conv_state_k,	784	cache=conv_state_k,
656	output_final_state=use_cache,	785	output_final_state=use_cache,
657	cu_seqlens=cu_seqlens,	786	cu_seqlens=cu_seqlens,
		787	+ **k_conv_kwargs,
658	)	788	)
659	v, conv_state_v = self.v_conv1d(	789	v, conv_state_v = self.v_conv1d(
660	x=v_proj_states,	790	x=v_proj_states,
661	cache=conv_state_v,	791	cache=conv_state_v,
662	output_final_state=use_cache,	792	output_final_state=use_cache,
663	cu_seqlens=cu_seqlens,	793	cu_seqlens=cu_seqlens,
		794	+ **v_conv_kwargs,
664	)	795	)
665	g = self.f_b_proj(self.f_a_proj(hidden_states))	796	g = self.f_b_proj(self.f_a_proj(hidden_states))
		797	+ beta = self.b_proj(hidden_states)
		798	+ # Modification start: LASP all-to-all for the decay gate and beta.
		799	+ if is_ulysses_enabled:
		800	+ g = all_to_all(g, ulysses_group, scatter_dim=2, gather_dim=1, gather_size=total_seq_len)
		801	+ beta = all_to_all(beta, ulysses_group, scatter_dim=2, gather_dim=1, gather_size=total_seq_len)
		802	+ if indices is not None:
		803	+ g = index_first_axis(
		804	+ rearrange(g, "b s ... -> (b s) ..."), indices).unsqueeze(0)
		805	+ beta = index_first_axis(
		806	+ rearrange(beta, "b s ... -> (b s) ..."), indices).unsqueeze(0)
		807	+ # Modification end
666	g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)	808	g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
667	- beta = self.b_proj(hidden_states).float()	809	+ beta = beta.float()
668		810
669	q, k = map(lambda x: rearrange(	811	q, k = map(lambda x: rearrange(
670	x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))	812	x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
671	v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)	813	v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
672		814
		815	+ # Modification start: slice per-head decay params to the local V-heads under LASP.
		816	+ if is_ulysses_enabled:
		817	+ v_head_offset = ulysses_rank * local_num_heads
		818	+ v_head_slice = slice(v_head_offset, v_head_offset + local_num_heads)
		819	+ A_log = self.A_log[v_head_slice]
		820	+ dt_bias_offset = ulysses_rank * local_value_dim
		821	+ dt_bias = self.dt_bias[dt_bias_offset: dt_bias_offset + local_value_dim]
		822	+ else:
		823	+ A_log = self.A_log
		824	+ dt_bias = self.dt_bias
		825	+ # Modification end
		826	+
673	if mode == 'chunk':	827	if mode == 'chunk':
674	kda_implementation = getattr(self.config, "kda_implementation", "fused")	828	kda_implementation = getattr(self.config, "kda_implementation", "fused")
675	if kda_implementation == "fused":	829	if kda_implementation == "fused":
	@@ -684,8 +838,8 @@ class KimiDeltaAttention(nn.Module):
684	v=v,	838	v=v,
685	g=g,	839	g=g,
686	beta=beta,	840	beta=beta,
687	- A_log=self.A_log,	841	+ A_log=A_log,
688	- dt_bias=self.dt_bias,	842	+ dt_bias=dt_bias,
689	initial_state=recurrent_state,	843	initial_state=recurrent_state,
690	output_final_state=True,	844	output_final_state=True,
691	use_qk_l2norm_in_kernel=True,	845	use_qk_l2norm_in_kernel=True,
	@@ -706,8 +860,8 @@ class KimiDeltaAttention(nn.Module):
706	v=v,	860	v=v,
707	g=g,	861	g=g,
708	beta=beta,	862	beta=beta,
709	- A_log=self.A_log,	863	+ A_log=A_log,
710	- dt_bias=self.dt_bias,	864	+ dt_bias=dt_bias,
711	initial_state=recurrent_state,	865	initial_state=recurrent_state,
712	output_final_state=True,	866	output_final_state=True,
713	use_qk_l2norm_in_kernel=True,	867	use_qk_l2norm_in_kernel=True,
	@@ -733,6 +887,17 @@ class KimiDeltaAttention(nn.Module):
733	cache_params.conv_states[self.layer_idx] = (	887	cache_params.conv_states[self.layer_idx] = (
734	conv_state_q, conv_state_k, conv_state_v)	888	conv_state_q, conv_state_k, conv_state_v)
735		889
		890	+ # Modification start: LASP all-to-all back (scatter sequence, gather heads),
		891	+ # so the gated norm and o_proj run on the sequence-sharded layout again.
		892	+ if is_ulysses_enabled:
		893	+ if indices is not None:
		894	+ # Restore the padded full-sequence layout before scattering the
		895	+ # sequence back, so pad positions are zero exactly as in the
		896	+ # non-CP path.
		897	+ o = pad_input(o.squeeze(0), indices, batch_size, total_seq_len)
		898	+ o = all_to_all(o, ulysses_group, scatter_dim=1, gather_dim=2)
		899	+ # Modification end
		900	+
736	if self.use_full_rank_gate:	901	if self.use_full_rank_gate:
737	g = self.g_proj(hidden_states)	902	g = self.g_proj(hidden_states)
738	else:	903	else:
	@@ -742,7 +907,7 @@ class KimiDeltaAttention(nn.Module):
742		907
743	o = rearrange(o, 'b t h d -> b t (h d)')	908	o = rearrange(o, 'b t h d -> b t (h d)')
744	o = self.o_proj(o)	909	o = self.o_proj(o)
745	- if attention_mask is not None:	910	+ if attention_mask is not None and not is_ulysses_enabled:
746	o = pad_input(o.squeeze(0), indices, batch_size, q_len)	911	o = pad_input(o.squeeze(0), indices, batch_size, q_len)
747		912
748	return o	913	return o
	@@ -1266,6 +1431,18 @@ class KimiLinearModel(KimiPreTrainedModel):
1266	linear_attn_mask = self._update_linear_attn_mask(	1431	linear_attn_mask = self._update_linear_attn_mask(
1267	attention_mask, cache_position)	1432	attention_mask, cache_position)
1268		1433
		1434	+ # Modification start: Ulysses context parallel patch.
		1435	+ # cu_seqlen params must be generated from the full-sequence position_ids
		1436	+ # before the sequence is split across CP ranks.
		1437	+ total_seq_len = inputs_embeds.shape[1]
		1438	+ set_seq_len("total", total_seq_len)
		1439	+ ps = get_parallel_state() if dist.is_initialized() else None
		1440	+ if ps is not None and ps.is_ulysses_enable():
		1441	+ kwargs.update(generate_ulysses_cu_seqlen_params(position_ids))
		1442	+ position_ids = split_forward_gather_backward_with_cp(position_ids, dim=1)
		1443	+ inputs_embeds = split_forward_gather_backward_with_cp(inputs_embeds, dim=1)
		1444	+ # Modification end
		1445	+
1269	hidden_states = inputs_embeds	1446	hidden_states = inputs_embeds
1270	if past_key_values is not None:	1447	if past_key_values is not None:
1271	assert isinstance(past_key_values, KimiDynamicCache)	1448	assert isinstance(past_key_values, KimiDynamicCache)

Mmindspeed_mm/fsdp/ops/kda/short_conv.py+54-19

	@@ -28,9 +28,13 @@ class ShortConvolution(nn.Conv1d):
28	kernel_size (int): Size of the convolution kernel	28	kernel_size (int): Size of the convolution kernel
29	bias (bool, optional): Whether to include learnable bias. Defaults to False.	29	bias (bool, optional): Whether to include learnable bias. Defaults to False.
30	activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'.	30	activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'.
31	- backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'.	31	+ backend (Optional[str], optional): Backend implementation ('triton' or 'cuda') for the decode `step` path. Defaults to 'triton'.
32	device (Optional[torch.device], optional): Device to place the layer on. Defaults to None.	32	device (Optional[torch.device], optional): Device to place the layer on. Defaults to None.
33	dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None.	33	dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None.
		34	+ implementation (str, optional): Implementation of the non-decode forward path,
		35	+ 'triton' (default) or 'ascendc' (AscendC fused op from fla_npu, NPU only).
		36	+ head_num (Optional[int], optional): Number of attention heads packed into the channel dim,
		37	+ required by the AscendC causal_conv1d op. Defaults to None (treated as 1).
34	**kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility)	38	**kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility)
35		39
36	Attributes:	40	Attributes:
	@@ -53,6 +57,8 @@ class ShortConvolution(nn.Conv1d):
53	backend: str \| None = 'triton',	57	backend: str \| None = 'triton',
54	device: torch.device \| None = None,	58	device: torch.device \| None = None,
55	dtype: torch.dtype \| None = None,	59	dtype: torch.dtype \| None = None,
		60	+ implementation: str = 'triton',
		61	+ head_num: int \| None = None,
56	**kwargs,	62	**kwargs,
57	):	63	):
58	super().__init__(	64	super().__init__(
	@@ -67,6 +73,15 @@ class ShortConvolution(nn.Conv1d):
67	)	73	)
68		74
69	self.hidden_size = hidden_size	75	self.hidden_size = hidden_size
		76	+ # Number of attention heads packed into the channel dim, required by the
		77	+ # AscendC causal_conv1d op. Defaults to 1 (whole channel dim as one head).
		78	+ self.head_num = head_num if head_num is not None else 1
		79	+ if implementation not in ('triton', 'ascendc'):
		80	+ raise ValueError(
		81	+ f"Unsupported causal conv1d implementation: {implementation}. "
		82	+ "Expected 'triton' or 'ascendc'."
		83	+ )
		84	+ self.implementation = implementation
70	self.activation = None	85	self.activation = None
71		86
72	if activation is not None:	87	if activation is not None:
	@@ -110,6 +125,7 @@ class ShortConvolution(nn.Conv1d):
110	output_final_state: bool = False,	125	output_final_state: bool = False,
111	cu_seqlens: torch.LongTensor \| None = None,	126	cu_seqlens: torch.LongTensor \| None = None,
112	chunk_indices: torch.LongTensor \| None = None,	127	chunk_indices: torch.LongTensor \| None = None,
		128	+ weight: torch.Tensor \| None = None,
113	**kwargs,	129	**kwargs,
114	) -> tuple[torch.Tensor, torch.Tensor \| None]:	130	) -> tuple[torch.Tensor, torch.Tensor \| None]:
115	"""	131	"""
	@@ -130,15 +146,13 @@ class ShortConvolution(nn.Conv1d):
130	Shape: [B+1]	146	Shape: [B+1]
131	chunk_indices (Optional[torch.LongTensor]):	147	chunk_indices (Optional[torch.LongTensor]):
132	Chunk indices for variable-length sequences. Default: `None`.	148	Chunk indices for variable-length sequences. Default: `None`.
		149	+ weight (`Optional[torch.Tensor]`):
		150	+ Optional weight override of shape `[D_local, 1, W]` (e.g. a head-sharded
		151	+ slice under Ulysses context parallel). Default: `None` (use `self.weight`).
133		152
134	Returns:	153	Returns:
135	Tensor of shape `[B, T, D]`.	154	Tensor of shape `[B, T, D]`.
136	"""	155	"""
137	- # Import here to avoid circular dependency
138	- # from fla.modules.conv.causal_conv1d import causal_conv1d
139	- from mindspeed_mm.fsdp.models.qwen3_5.causal_conv1d import causal_conv1d
140	-
141	-
142	B, T, *_ = x.shape	156	B, T, *_ = x.shape
143	N = B if cu_seqlens is None else len(cu_seqlens) - 1	157	N = B if cu_seqlens is None else len(cu_seqlens) - 1
144	if mask is not None:	158	if mask is not None:
	@@ -157,23 +171,44 @@ class ShortConvolution(nn.Conv1d):
157	)	171	)
158	return y, cache	172	return y, cache
159		173
160	- # cuda backend do not support:	174	+ if weight is None:
161	- # 1. both `cu_seqlens` and `cache` being provided	175	+ weight = self.weight
162	- # 2. both `cu_seqlens` and `output_final_state` being provided	176	+
163	- # and other small issues	177	+ if self.implementation == 'ascendc':
164	- # to simplify the implementation, we just switch to triton backend	178	+ # Import here to avoid circular dependency and to keep torch_npu/fla_npu
165	- if self.backend == 'cuda' and cache is not None:	179	+ # optional for the triton path.
166	- warnings.warn(	180	+ from mindspeed_mm.fsdp.ops.gdn.causal_conv1d_ascendc import causal_conv1d_ascendc
167	- "The CUDA backend does not support both `cu_seqlens` and `cache` being provided, "	181	+
168	- "or both `cu_seqlens` and `output_final_state` being provided. "	182	+ # The AscendC op needs the number of heads packed in the channel dim.
169	- "Switching to the Triton backend instead. ",	183	+ # Derive it from the (possibly head-sharded, e.g. under Ulysses CP) weight
170	- stacklevel=2,	184	+ # so a weight override scales H proportionally.
		185	+ channels_per_head = self.hidden_size // self.head_num
		186	+ if weight.shape[0] % channels_per_head != 0:
		187	+ raise ValueError(
		188	+ f"weight channels ({weight.shape[0]}) must be a multiple of "
		189	+ f"channels_per_head ({channels_per_head})"
		190	+ )
		191	+ H = weight.shape[0] // channels_per_head
		192	+ y, final_state = causal_conv1d_ascendc(
		193	+ x=x,
		194	+ weight=rearrange(weight, "d 1 w -> d w"), # AscendC op expects [D, W] weight
		195	+ H=H,
		196	+ bias=self.bias,
		197	+ residual=residual,
		198	+ initial_state=cache,
		199	+ activation=self.activation,
		200	+ cu_seqlens=cu_seqlens,
		201	+ output_final_state=output_final_state,
171	)	202	)
172	- self.backend = 'triton'	203	+ # The AscendC op returns the head layout [B, H, T, d]; convert back to [B, T, D]
		204	+ y = y.transpose(1, 2).reshape(B, T, -1)
		205	+ return y, final_state
		206	+
		207	+ from mindspeed_mm.fsdp.models.qwen3_5.causal_conv1d import causal_conv1d
173		208
174	return causal_conv1d(	209	return causal_conv1d(
175	x=x,	210	x=x,
176	- weight=rearrange(self.weight, "d 1 w -> w d"), # NOTE：修改了权重格式	211	+ weight=rearrange(weight, "d 1 w -> w d"),
177	bias=self.bias,	212	bias=self.bias,
178	residual=residual,	213	residual=residual,
179	initial_state=cache,	214	initial_state=cache,