feat: add DeepSeek V4 NPU kernel Python bindings. - #2241
Conversation
yingxudeng
changed the title
| ) | ||
|
|
||
|
|
||
| def grouped_moe_with_selected_experts( |
There was a problem hiding this comment.
这里GPU没有就不需要加。
kernels_cuda和kernels_npu暴露的接口本来就不需要一样
There was a problem hiding this comment.
已修改：删除 CUDA 侧 grouped_moe_with_selected_experts 占位实现及其导出。该接口仅在当前 NPU 实现中暴露，不再要求 kernels_cuda 与 kernels_npu 保持相同接口。
There was a problem hiding this comment.
张老师，提交pr跑流水线有一个case没过，好像是说需要cuda和npu的接口保持一致，能帮忙看看这个后续是否需要修改测试用例还是对齐断言的接口？
| del clamp_limit, glu_alpha, glu_bias | ||
| # Output is half of input's last dim (SwiGLU splits gate/up). | ||
| out_dim = x.size(-1) // 2 | ||
| act_quantized = x.new_empty(x.size(0), out_dim, dtype=torch.int8) |
There was a problem hiding this comment.
硬编码 2D 输出形状？ 算子是必定输出2D形状吗
There was a problem hiding this comment.
已修改：fake 实现不再硬编码二维输出，输出 shape 改为 (*x.shape[:-1], x.size(-1) // 2)，保留输入的前导维度；同时补充了 3D 输入测试。
| bias=None, | ||
| quant_scale=None, | ||
| quant_offset=None, | ||
| group_index=group_list.to(torch.int64), |
There was a problem hiding this comment.
已修改：移除了重复的 group_list.to(torch.int64)，只在进入算子调用前执行一次必要的 dtype 转换。
| return_softmax_lse, | ||
| ) | ||
| out = q.new_empty(q.shape, dtype=q.dtype) | ||
| lse = q.new_empty((0,), dtype=q.dtype) |
There was a problem hiding this comment.
为何不用 return_softmax_lse ，而是永远传 0 回去
There was a problem hiding this comment.
已修改：fake 实现现在根据 return_softmax_lse 返回结果。关闭时返回 shape 为 (0,) 的 float32 空 tensor；开启时返回 shape 为 q.shape[:-1] + (1,) 的 float32 tensor，并补充了两种分支的测试。
| TORCH_LIBRARY_IMPL(xllm_ops, CompositeExplicitAutograd, m) { | ||
| m.impl("build_cp_context", TORCH_FN(xllm::build_cp_context_npu)); | ||
| // ---- DeepSeek-V4 DSA kernels ---- | ||
| m.impl("moe_gating_top_k_hash", |
There was a problem hiding this comment.
CompositeExplicitAutograd 是通用的和设备无关的吧，你的下面调用都是 aclnn 的吧和设备有关吧
There was a problem hiding this comment.
已修改：所有带必需 Tensor 输入且调用 ACLNN 的 DSV4 算子已从 CompositeExplicitAutograd 移到 PrivateUse1，包括 gating、SwiGLU、HC、compressor、sparse attention 和 lightning indexer。仅保留两个允许所有 Tensor 参数为空的 metadata factory 在 CompositeExplicitAutograd，因为这种调用无法从输入推导 NPU dispatch key，且实现会显式选择输出设备。另补充了 dispatch key 注册测试。修复见 5bc5ea0。
0moyi0-2024
left a comment
0moyi0-2024
left a comment
There was a problem hiding this comment.
补充两点 grouped MoE 的 EP 契约问题，建议在本 PR 内修正并增加单测。
| return torch_npu.npu_moe_token_unpermute( | ||
| permuted_tokens=output, | ||
| sorted_indices=expanded_row_idx.abs(), | ||
| probs=topk_weights.to(output.dtype), |
There was a problem hiding this comment.
这里直接把原始 topk_weights 传给 token_unpermute，但 active_expert_range 只路由当前 rank 的专家。原生 C++ FusedMoEImpl::select_experts 会先把非本地 expert 的权重置零；否则 EP all-reduce 时当前 rank 可能贡献非本地路由权重。建议在 helper 内按 [start_expert_id, start_expert_id + local_expert_count) 做 mask，并在测试里校验传给 probs 的值。
There was a problem hiding this comment.
已在 5bc5ea0 修复：按当前 rank 的 active expert range 对 topk_weights 做 mask，再把本地权重传给 npu_moe_token_unpermute；同时增加混合本地/非本地 expert ID 的测试，校验非本地概率为 0。
| sorted_hidden_i8, pertoken_scale = _kernels.dynamic_quant(expanded_hidden) | ||
| if pertoken_scale is None: | ||
| raise RuntimeError("dynamic_quant did not return a per-token scale") | ||
| group_list = expert_tokens.to(torch.int64) |
There was a problem hiding this comment.
建议显式约束 expert_tokens 的长度为 local_expert_count（至少切片并校验）。这个函数的权重只有本地专家，但当前直接把完整 expert_tokens 作为 group_list；如果算子返回全局 expert 计数，group 数会与 w13/w2 的专家维不一致。原生 C++ 路径对此有 CHECK_EQ(group_list.size(0), local_physical_experts_num())。
There was a problem hiding this comment.
已在 5bc5ea0 修复：校验 expert_tokens 至少覆盖本地专家数量，并只使用 expert_tokens[:local_expert_count] 作为 group list；同时校验本地专家数与 w13/w2 的第一维一致，并补充测试。
3 participants
Summary
This is PR 1 of the split DeepSeek V4 Python integration. Runtime/metadata and model integration will follow separately.
Validation
ruff check: passedruff format --check: passedgit diff --check: passed6 passed363/363, passed withMAX_JOBS=326 passed(grouped GEMM x2, partial RoPE, compressor, QLI, sparse attention)