👋 您好，欢迎向 MindSpeed MM 提交 Issue！
我们已收到您的反馈，感谢你对开源社区的支持。🎉
📅 处理时效： 维护团队将在 24 小时内 查看并回复您的问题（工作日）。
🚨 紧急联系： 如果您的问题非常紧急，可通过以下方式联系我们：
💬 微信：
请确保 Issue 描述清晰，包含复现步骤和日志，这将帮助我们更快定位问题。谢谢！
感谢使用。
当前代码暂不支持cuseqlens场景，建议bugfix通过后再使用，我们将尽快修改。
在那以前，建议优先使用cuseqlens=None的形式。
packing转换为非packing参考代码
def varlen_to_nonvarlen(cu_seqlens, *vars):
B = len(cu_seqlens) - 1 # batch size
max_len = max(cu_seqlens[i+1] - cu_seqlens[i] for i in range(B)) # 最长序列长度
# 创建全 0 的非 varlen 输出张量
nonvarlens = [torch.zeros(B, max_len, *v.shape[2:], device=v.device, dtype=v.dtype) for v in vars]
# 填充数据
for i in range(B):
start = cu_seqlens[i]
end = cu_seqlens[i+1]
seq_len = end - start
if seq_len > 0: # 避免除以 0
for j in range(len(nonvarlens)):
nonvarlens[j][i, :seq_len] = vars[j][0, start:end]
return nonvarlens
当前已知问题至少包含：
g = chunk_local_cumsum(g, cu_seqlens=cu_seqlens, head_first=True)
建议改为head_first=False，否则读取数据可能存在问题。
@triton.jit(do_not_specialize=['T'])
def recompute_w_u_fwd_kernel(
k,
v,
beta,
w,
u,
A,
g,
gk,
cu_seqlens,
chunk_indices,
T,
B,
H: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
NT: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_GK: tl.constexpr,
IS_VARLEN: tl.constexpr
):
core_id = tl.program_id(0)
total_cores = tl.num_programs(0)
T_max = T
其中入参T和循环体内变量重名，需要改名，某些版本的CANN存在bug，可能导致概率性输出nan。
chunk_scaled_dot_kkt_fwd 的输出也存在极大值，这个问题可以稳定复现。
实际复现条件和整网实际数据有关系，正在自我诊断，有后续进展会在 issue 上更新。
cu_seqlens = [0, 115, 116]
A = chunk_scaled_dot_kkt_fwd(
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_size=chunk_size,
output_dtype=torch.float32
)
A: torch.Size([1, 116, 4, 64])，按照 [116, 256] shape绘制
另外，代码中添加 debug，打印其中的部分数据
if USE_G:
p_g = tl.make_block_ptr(g + g_batch_off + bos + i_h * T_, (T_local,), (1,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,))
b_g_diff = b_g[:, None] - b_g[None, :]
b_g_diff = tl.minimum(tl.maximum(b_g_diff, -50.0), 50.0)
b_A *= tl.exp(b_g_diff)
#? DEBUG >>>>
p_A_ = tl.make_block_ptr(debug_A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
tl.store(p_A_, b_g_diff.to(p_A_.dtype.element_ty), boundary_check=(0, 1))
#? DEBUG <<<<
b_A *= b_beta[:, None]
p_A = tl.make_block_ptr(A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
也就是说我把 b_A *= tl.exp(b_g_diff)中的 b_g_diff拿出来了，这部分绘图结果如下：
这个地方的数据有点太大了，不正常
上述情况可能和模型侧实际数据相关，正在自我排查。
chunk_scaled_dot_kkt_fwd 的输出也存在极大值，这个问题可以稳定复现。
实际复现条件和整网实际数据有关系，正在自我诊断，有后续进展会在 issue 上更新。
cu_seqlens = [0, 115, 116]A = chunk_scaled_dot_kkt_fwd( k=k, g=g, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size, output_dtype=torch.float32 )
A: torch.Size([1, 116, 4, 64])，按照[116, 256]shape绘制
另外，代码中添加 debug，打印其中的部分数据
if USE_G: p_g = tl.make_block_ptr(g + g_batch_off + bos + i_h * T_, (T_local,), (1,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) b_g_diff = b_g[:, None] - b_g[None, :] b_g_diff = tl.minimum(tl.maximum(b_g_diff, -50.0), 50.0) b_A *= tl.exp(b_g_diff) #? DEBUG >>>> p_A_ = tl.make_block_ptr(debug_A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) tl.store(p_A_, b_g_diff.to(p_A_.dtype.element_ty), boundary_check=(0, 1)) #? DEBUG <<<< b_A *= b_beta[:, None] p_A = tl.make_block_ptr(A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))也就是说我把
b_A *= tl.exp(b_g_diff)中的b_g_diff拿出来了，这部分绘图结果如下：
这个地方的数据有点太大了，不正常
上述情况可能和模型侧实际数据相关，正在自我排查。
可否提供chunk_scaled_dot_kkt_fwd 的输入数据，我们此前未发现此算子存在问题，需要复现一下，谢谢
保存方法参考
A = chunk_scaled_dot_kkt_fwd(
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_size=chunk_size,
output_dtype=torch.float32
)
torch.save(dict(
A=A,
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_size=chunk_size,
output_dtype=torch.float32
), "chunk_scaled_dot_kkt_fwd.pt")
经过定位，发现存在两个问题引发了这个，已经修复
- 如您所述，使用
head_first稳定会造成数据异常，如下是走 chunk_local_cumsum 采用不同head_first时，A=chunk_scaled_dot_kkt_fwd()的输出 A 的数据绘图
head_first=True
head_first=False
w, u = recompute_w_u_fwd()存在已知 bug，
ref to https://gitcode.com/Ascend/MindSpeed/pull/3360/diffs
我会尝试更优雅的临时修复方式，后续提出 PR 进行修复。
另外建议暂时默认使用 head_first=False，如果认可，我会在 pr 中带上这一点
https://gitcode.com/Ascend/MindSpeed-MM/blob/master/mindspeed_mm/fsdp/models/qwen3_5/chunk_gated_delta_rule.py#L30
经过定位，发现存在两个问题引发了这个，已经修复
- 如您所述，使用
head_first稳定会造成数据异常，如下是走 chunk_local_cumsum 采用不同head_first时，A=chunk_scaled_dot_kkt_fwd()的输出 A 的数据绘图head_first=True
head_first=False
w, u = recompute_w_u_fwd()存在已知 bug，ref to https://gitcode.com/Ascend/MindSpeed/pull/3360/diffs
我会尝试更优雅的临时修复方式，后续提出 PR 进行修复。
另外建议暂时默认使用 head_first=False，如果认可，我会在 pr 中带上这一点
https://gitcode.com/Ascend/MindSpeed-MM/blob/master/mindspeed_mm/fsdp/models/qwen3_5/chunk_gated_delta_rule.py#L30
那看来都是已知问题，我上面的回复已经覆盖了。
recompute的问题，如我前面回复所述，实际为NPUIR编译器的bug，此问题当前主线已经修复，可以参考我在NPUIR提出的相关issue
因此我们在core上的代码实际也是临时规避，可能不需要太优雅
但是当前仍然不推荐使用NPUIR主线，3.13修复recompute问题，3.18又引入了新问题需要规避，未在issue中跟踪。
补充说明 head_first 的输出情况，方便日后问题跟踪
cu_seqlens = tensor([ 0, 113, 114, 115, 116]
原始输入的 g
head_first 为 False 时，输出 g 的数据：
head_first 为 True 时，输出 g 的数据：
可以看到底部的数据是有差异的
@LinMingZhe 你好 当前GDN 支持varlen场景了吗？项目急需支持。
@LinMingZhe 你好 当前GDN 支持varlen场景了吗？项目急需支持。
可以参考这个PR，前段时间卡CI了，预计最快今天合入，急需的话可以参考PR修改，或者直接克隆PR代码
PS1：varlen场景还需要注意下conv1d等计算的varlen支持度，此PR仅处理GDN
PS2：PR版本代码在seqlen>64K时，可能触发卡死问题，为编译器问题，内部问题单跟踪中
@LinMingZhe 你好 当前GDN 支持varlen场景了吗？项目急需支持。
可以参考这个PR，前段时间卡CI了，预计最快今天合入，急需的话可以参考PR修改，或者直接克隆PR代码
PS1：varlen场景还需要注意下conv1d等计算的varlen支持度，此PR仅处理GDN
PS2：PR版本代码在seqlen>64K时，可能触发卡死问题，为编译器问题，内部问题单跟踪中
好的，了解。另外，我们当前环境是cann8.5.0，这个需要升级到cann9.0.0吗？
@LinMingZhe 你好 当前GDN 支持varlen场景了吗？项目急需支持。
可以参考这个PR，前段时间卡CI了，预计最快今天合入，急需的话可以参考PR修改，或者直接克隆PR代码
PS1：varlen场景还需要注意下conv1d等计算的varlen支持度，此PR仅处理GDN
PS2：PR版本代码在seqlen>64K时，可能触发卡死问题，为编译器问题，内部问题单跟踪中好的，了解。另外，我们当前环境是cann8.5.0，这个需要升级到cann9.0.0吗？
当前上仓的版本是在8.5.1下验证的，理论上8.5.0区别不大，cann9.0.0的一些版本我们有发现别的问题在跟踪
已验证解决前向 nan 的问题，感谢支持
环境信息
🐛 问题描述
mindspeed_mm/fsdp/models/qwen3_5/chunk_gated_delta_rule.py使能 cu_seqlens 后必现 nan 问题
chunk_scaled_dot_kkt_fwd 和 recompute_w_u_fwd 都会先后出现 nan。
详细细节待定位...
欢迎加入社区，感谢您对社区的贡献 🎉!