feat: add FlashComm1 and MMRS fusion support for Qwen3.5 NPU.#1890
feat: add FlashComm1 and MMRS fusion support for Qwen3.5 NPU.#1890Cooofish wants to merge 23 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the Flash Communication 1 (FC1) sequence-parallel optimization for tensor parallel inference on NPU, integrating it into layers like RowParallelLinear and FusedMoE, and adding a fused matmul_reduce_scatter NPU kernel. The review feedback identifies several style guide violations, including the use of member functions in the FlashComm1Context struct and direct access to FLAGS_ global variables. It also highlights critical issues such as potential undefined behavior from missing tensor definition checks, duplicate helper functions, a performance regression in MoE reduction when ep_size == 1, and a type mismatch where int64_t is used instead of int32_t for CANN/ATB host data.
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.
a2db15a to
dd852e4
Compare
Cooofish
changed the title
Cooofish
changed the title
9bb1d73 to
2c1850e
Compare
Cooofish
changed the title
- Infer NPU Python, PyTorch, and torch_npu paths during CMake configure when the environment is unset. - Propagate CARGO_HOME into Rust custom commands and align the workspace Cargo mirror name with the root Cargo config. - Remove stale selected-gather and Device::type_str() remnants after the FC1 context refactor. - Keep FC1 default-off and raise the activation threshold to long-prefill TP=8 cases.
- Keep the wrapper focused on validating inputs, resolving the HCCL group, selecting comm_mode, and invoking torch_npu npu_mm_reduce_scatter_base. - Return an empty tensor on unsupported inputs so the existing row-parallel matmul plus reduce_scatter fallback remains owned by the caller. - Remove the duplicate local torch::matmul fallback and reject-reason string builder to make the FC1 MMRS path easier to review and maintain.
Apply clang-format 20.1.6 to the FC1/MMRS diff and fix include ordering reported by the format check.
yingxudeng
changed the title
本地python setup.py build test已通过100% tests passed, 0 tests failed out of 833 Total Test time (real) = 125.67 sec The following tests did not run: commit b56d8e5也测试通过Triton NPU tests: CTest: I20260722 10:23:49.394614 setup.py:797] ================================================================================ |
Preserve the latest main quantization refactor while restoring the FC1/MMRS row-parallel overload, including its guarded MMRS fallback and tensor-defined checks. Formatting was validated with the repository clang-format configuration.
| #include "common/global_flags.h" | ||
| #include "framework/parallel_state/parallel_state.h" | ||
|
|
||
| DEFINE_bool(enable_flashcomm1, |
There was a problem hiding this comment.
确实不该放这里，我把这些配置定义移到了KernelConfig，改成通过 from_flags() 解析，并在create_options()中传递给FlashComm1Options
| : std::nullopt; | ||
| CHECK(weight_scale.has_value() && weight_scale.value().defined()) | ||
| << "weight_scale is required for w8a8_dynamic quant matmul."; | ||
| #if defined(USE_DCU) |
There was a problem hiding this comment.
这里的 DCU 分支是复制原RowParallel forward量化逻辑时带入的，不该在FC1/MMRS用。已将该 overload 和调用路径限制到 USE_NPU。
| .rank_tablefile(eplb_config.rank_tablefile()) | ||
| .expert_parallel_degree(eplb_config.expert_parallel_degree()) | ||
| .enable_chunked_prefill(scheduler_config.enable_chunked_prefill()) | ||
| .enable_flashcomm1(FLAGS_enable_flashcomm1) |
There was a problem hiding this comment.
和第一处是同一问题，已经改为把这些配置定义移到KernelConfig，通过 from_flags() 解析，并在create_options()中传递给FlashComm1Options
Move FC1 and MMRS flags into KernelConfig, remove direct gflag access from the runtime context, and limit the reduce-mode FC1 path to NPU. Preserve the existing defaults and Options propagation.
| auto output_shape = mmrs_input.sizes().vec(); | ||
| output_shape[0] = fc1_ctx->padded_local_num_tokens; | ||
| output_shape[1] = weight_.size(0); | ||
| torch::Tensor mmrs_output = torch::empty(output_shape, input.options()); |
There was a problem hiding this comment.
torch::empty 了，但是实际没有用？是不是不太合理？只需要shape吧
There was a problem hiding this comment.
确实不合理。现已删除 torch::empty 以及整个封装链路中无效的 output 参数。
| x = moe_mlp_(x, input_params); | ||
| } else { | ||
| x = mlp_(x); | ||
| x = mlp_->forward(x); |
There was a problem hiding this comment.
Why are we calling both mlp_(x) and mlp_->forward(x) here?
There was a problem hiding this comment.
这里可能是对 diff 的误解。x = mlp_(x) 是被删除的代码，当前实现中只保留了 x = mlp_->forward(x)，没有同时调用mlp_(x) 和 mlp_->forward(x)
There was a problem hiding this comment.
sorry. 为什么改成了 x = mlp_->forward(x);
| input = xllm::parallel_state::scatter(input, process_group_); | ||
| } | ||
|
|
||
| xllm::kernel::ScaledQuantizeParams quantize_params; |
There was a problem hiding this comment.
Does torch::Tensor RowParallelLinearImpl::forward heavily duplicate the existing code?
There was a problem hiding this comment.
原本考虑避免本次改动影响其他后端原有的量化、scatter 和 reduction 行为，保留了原 forward 路径。
在commit f189579中已经改为：
保留原有 forward(input) 和 FC1 使用的 forward(input, reduce_mode) 两个入口，但两个入口现在统一调用同一个 forward_impl()。量化、scatter、matmul、MMRS、fallback 和结果通信逻辑只保留一份。
原接口会根据 enable_result_reduction_ 映射到 NONE 或 ALL_REDUCE，维持原有行为；只有显式传入 REDUCE_SCATTER 或 MATMUL_REDUCE_SCATTER 时才会启用 FC1 的 sequence-sharded 处理，避免影响其他 RowParallelLinear 调用路径。
已通过 clang-format 和 NPU 单文件编译检查。
| return gathered; | ||
| } | ||
|
|
||
| torch::Tensor maybe_pad_for_reduce(const torch::Tensor& input, |
There was a problem hiding this comment.
Is maybe_pad_for_reduce unused?
There was a problem hiding this comment.
是的。已删除 maybe_pad_for_reduce 的声明和实现。
Remove the unused preallocated MMRS output parameter because torch_npu returns the result tensor directly. Also delete the unreferenced maybe_pad_for_reduce helper.
Route the existing and FC1-aware RowParallelLinear forward overloads through one implementation. Preserve the default reduction behavior while applying sequence-sharded handling only to explicit FC1 reduce modes.
|
最新commit 971dd79已本地build test通过 |
|
|
||
| // MLP forward | ||
| if (moe_mlp_) { | ||
| x = moe_mlp_(x, input_params); |
There was a problem hiding this comment.
对应的 moe 不用改嘛？moe 版本精度是否有问题
| mmrs_weight_t_.size(0) == weight_.size(1) && | ||
| mmrs_weight_t_.size(1) == weight_.size(0); | ||
| if (!valid) { | ||
| mmrs_weight_t_ = weight_.transpose(0, 1).contiguous(); |
There was a problem hiding this comment.
这个操作不能在模型初始化或者 load_state_dict 末尾做吗？
| const int64_t world_size = process_group->world_size(); | ||
| const int64_t k = a.size(1); | ||
| if (world_size == 2) { | ||
| return k == 3072 || k == 8704; |
| if (residual.size(0) == ctx.original_num_tokens) { | ||
| return shard_sequence(residual, ctx); | ||
| } | ||
| CHECK_EQ(residual.size(0), ctx.padded_local_num_tokens) |
There was a problem hiding this comment.
走到这里的时候，是否是必定失败的。下面的 return residual; 是不是死代码
| const FlashComm1Context& ctx) { | ||
| auto gathered = gather_sequence(input, ctx); | ||
| if (ctx.pad_size > 0) { | ||
| return gathered.slice(0, 0, ctx.original_num_tokens); |
There was a problem hiding this comment.
gather_sequence 里面也有 gathered.slice(0, 0, ctx.original_num_tokens); 和这里的操作区别目的是什么
| const int64_t shard_start = | ||
| static_cast<int64_t>(ctx.tp_rank) * ctx.padded_local_num_tokens; | ||
| const int64_t shard_end = shard_start + ctx.padded_local_num_tokens; | ||
| return padded_input.slice(0, shard_start, shard_end).contiguous(); |
There was a problem hiding this comment.
padded_input = pad_rows_by_copy(input, ctx.padded_num_tokens); 和 padded_input.slice(0, shard_start, shard_end).contiguous(); 是不是多次copy了，能否优化？
| } | ||
|
|
||
| torch::Tensor output; | ||
| #if !defined(USE_NPU) |
3 participants
Summary
This PR adds FlashComm1 (FC1) support for NPU tensor-parallel inference.
FC1 is a prefill-side communication optimization for long-context workloads. It shards hidden states along the sequence dimension across TP ranks, lets supported row-parallel linear layers consume the sharded layout, and uses Matmul + ReduceScatter fusion to reduce part of the communication and launch overhead in prefill.
The feature is disabled by default:
To enable the full FC1 path:
More details about design, usage, suitable workloads, and validation guidance are documented in:
Main Changes
FlashComm1Contextand helper functions for prefill sequence padding, sequence sharding, reduce, gather, and row-parallel layout handling.matmul_reduce_scatterkernel API backed by torch_npunpu_mm_reduce_scatter_base.dp=1,cp=1,tp>=8, andprefill_tokens>=8192.MMRS Integration
The MMRS path intentionally uses the torch_npu fused operator instead of implementing a local fused kernel in xLLM:
comm_mode, and callsnpu_mm_reduce_scatter_base.mmrs_comm_modevalues areaiv,ai_cpu, andnone.aiv.ai_cputhrough the existing mode selection logic.Default Activation Policy
FC1 and MMRS are both opt-in. Even when enabled by flags, FC1 only becomes active when runtime conditions are suitable.
Default suitable case:
prefill_tokens >= 8192.tp >= 8.dp = 1.cp = 1.Cases where FC1 is disabled or expected to have limited benefit:
tp < 8,dp > 1, orcp > 1.Validation
Build validation performed locally:
Result: passed, producing
build_flash_npu3/xllm/xllm.Ascend A3 and ACL Graph compatibility
On Ascend A3 (
Ascend910_93), some CANN 9.0-based xllm_ops/ATBpackages may incorrectly select the 910A CustomPagedAttention runner
when ACL Graph is enabled. This can cause CustomPagedAttention setup
to fail with status 27.
This issue is independent of FlashComm1/MMRS. Use an xllm_ops build
that correctly classifies
Ascend910_93as a 910B-family device anddeploy the matching
libatb_customize.soandlibcustomize_ops.sotogether.
Correctness checks performed locally:
Performance and profiling checks performed locally:
graph=true, 1K output, 8K/16K/32K input, parallel=2/4/8/16, and two warmup requests.tp >= 8before FC1 becomes active.MatmulReduceScatterV2appears when FC1/MMRS is enabled.Performance Notes
FC1 is not a general-purpose speedup switch. In the selected multi-round TP=8 matrix, the most balanced results are 32K input at parallel=4/8/16: TTFT improves by 1.94%-3.95%, prompt/decode/request throughput improves by 2.10%-2.65%, and average latency improves by 1.71%-1.97%. The 32K, parallel=2 case also improves TTFT by 6.10%, with throughput and latency essentially unchanged to slightly better.
The 8K cases improve TTFT by 3.95%-5.57%, but throughput and average latency are flat to about 1.6% worse. At 16K, parallel=2 is balanced (+3.26% TTFT and about +1.1% throughput), while parallel=8/16 improves TTFT by about 2.6% but regresses end-to-end throughput substantially. The 16K, parallel=4 result was remeasured three times; the final repeat is effectively neutral (+0.40% TTFT, -0.33% throughput), showing why a single TTFT average must not be interpreted as a stable FC1 gain. These results reinforce that FC1 primarily targets prefill latency; TTFT improvement does not guarantee an end-to-end throughput improvement when decode and layout/communication overhead dominate.
Current profiling limitations:
For this reason, FC1 remains disabled by default and should be enabled only for suitable long-prefill deployments.
Reviewer Focus
Please focus review on:
comm_modeselection, and fallback behavior.Recent Updates
MMRS wrapper simplification
comm_mode, and callsnpu_mm_reduce_scatter_base.torch::matmulfallback from the wrapper. Unsupported MMRS cases now return an empty tensor and let the existing row-parallel fallback remain in the caller.Gemini review fixes
FlashComm1Contextas a plain data struct and moved sequence-sharding helpers to free functions.linear.cppand reused the shared FC1 helper.defined()checks before inspecting MMRS input tensor metadata.ModelConfig::get_instance().backend(),EPLBConfig::get_instance().expert_parallel_degree(), preserve theep_size == 1single-reduce MoE optimization, and no longer contain the flagged Qwen3 ATBint64_thostData vectors.Diff-scope cleanup
main.CMakeLists.txtto match upstreammain.