feat: add FlashComm1 and MMRS fusion support for Qwen3.5 NPU.

[Cooofish]

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 kernel-launch overhead in prefill.

The feature is disabled by default:

--enable_flashcomm1=false
--enable_mmrs_fusion=false

To enable the full FC1 path:

--enable_flashcomm1=true \
--enable_mmrs_fusion=true

More details about the design, usage, suitable workloads, dependencies, and validation guidance are documented in:

• https://github.com/xLLM-AI/xllm-docs/blob/main/src/content/docs/zh/features/flashcomm.md

Main Changes

• Add FC1 runtime options, global flags, model context fields, and worker/master option propagation.
• Add FlashComm1Context and helper functions for prefill sequence padding, sequence sharding, reduce, gather, and row-parallel layout handling.
• Add the NPU matmul_reduce_scatter API backed by torch_npu npu_mm_reduce_scatter_base.
• Integrate fused MMRS into supported row-parallel linear layers, with fallback to the original matmul + reduce path when MMRS is not applicable.
• Wire FC1 into Qwen3 and Qwen3-Next hybrid Dense model paths.
• Keep model output semantics unchanged by gathering back to normal hidden states before returning from the FC1 prefill path.
• Gate FC1 to prefill-only long-context cases by default: dp=1, cp=1, tp>=8, and prefill_tokens>=8192.
• Keep FC1 disabled during decode. TPOT and decode throughput are not expected to directly benefit from this feature.
• Disable FC1/MMRS when routed experts are present. The current implementation supports Dense models only; the routed-expert MoE path is outside this PR's scope.

MMRS Integration

The MMRS path intentionally uses the torch_npu fused operator instead of implementing a local fused kernel in xLLM:

• The xLLM wrapper validates inputs, resolves the HCCL process group, selects comm_mode, and calls npu_mm_reduce_scatter_base.
• Supported mmrs_comm_mode values are aiv, ai_cpu, and none.
• The default mode remains aiv.
• The implementation does not contain hard-coded risky-shape checks or magic-number-based switching between aiv and ai_cpu.
• ai_cpu can be selected explicitly through --mmrs_comm_mode=ai_cpu for environment-level diagnosis or fallback.
• Unsupported MMRS inputs return an undefined tensor to the caller, which then uses the existing row-parallel matmul + reduce path.

ops-transformers dependency

The torch_npu MMRS path relies on the operator implementation provided by ops-transformers. The Ascend A3 validation environment uses CANN 9.0.0 together with the MMRS-related operator built from ops-transformers 9.1.0.

When mmrs_comm_mode=aiv is used, an ops-transformers 9.1.0 or newer operator library containing the MMRS AIV fix must be built, installed, and loaded at runtime. This resolves the earlier FIXP configuration illegal / 0x3000084 failure observed for some shapes with the older AIV operator.

Only the required ops-transformers operator may be upgraded while retaining CANN 9.0.0, but the compile-time headers and runtime shared libraries must match. The loaded ops-transformers, torch_npu, ATB, and CANN components must remain ABI-compatible. The operator library path should be fixed explicitly in the deployment environment so that the tested build is loaded consistently.

The old shape-specific aiv to ai_cpu switching logic has been removed because it encoded incomplete risky-shape knowledge and could not guarantee safety for arbitrary production shapes. ai_cpu remains available as an explicit configuration option, not an automatic per-shape policy.

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:

• NPU backend.
• Prefill phase.
• Long input, currently prefill_tokens >= 8192.
• tp >= 8.
• dp = 1.
• cp = 1.
• Dense models using BF16/FP16 non-quantized row-parallel paths.

Cases where FC1 is disabled or expected to have limited benefit:

• Decode phase.
• Short prompts such as 2K input.
• Long-output workloads where decode dominates total latency.
• tp < 8, dp > 1, or cp > 1.
• Quantized row-parallel paths, which currently skip fused MMRS.
• MoE models with routed experts, which currently disable FC1/MMRS.

Validation

Build validation performed locally:

NPU_HOME_PATH=/usr/local/Ascend/cann-9.0.0 \
PYTHON_INCLUDE_PATH=/usr/local/python3.11.15/include/python3.11 \
PYTORCH_INSTALL_PATH=/usr/local/python3.11.15/lib/python3.11/site-packages/torch \
PYTORCH_NPU_INSTALL_PATH=/usr/local/libtorch_npu \
LIBTORCH_ROOT=/usr/local/python3.11.15/lib/python3.11/site-packages/torch \
ninja -C build_flash_npu3 xllm

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/ATB packages 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_93 as a 910B-family device and deploy the matching libatb_customize.so and libcustomize_ops.so together.

Correctness checks performed locally:

• TP=4, 2K input / 2K output, FC1 off vs FC1 on token IDs matched.
• Chunked prefill active validation passed for long and uneven prompt lengths.
• MMRS consistency validation passed with token mismatch = 0 and sampled logprob max diff < 0.03.

Performance and profiling checks performed locally:

• The FC1-off/on matrix was run with Qwen3.5-27B, graph=true, 1K output, 8K/16K/32K input, parallel=2/4/8/16, and two warmup requests.
• The selected primary matrix contains 72 individual measurements, covering 36 complete FC1-off/on comparisons collected from repeated TP=8 measurements. The July 19 single-round rerun is reported separately and is not used to replace the repeated 32K multi-concurrency conclusion.
• The prior MMRS AICore failures were resolved for the reported matrix after loading the compatible ops-transformers 9.1.0 operator build.
• TP=2 and TP=4 were retained as flag/no-op controls. The table below reports TP=8 because the default runtime policy requires tp >= 8 before FC1 becomes active.
• Positive percentages consistently mean that FC1 on is better: lower TTFT/TPOT/latency or higher throughput. TTFT uses the stable post-warmup request average and remains sensitive to scheduler and batch composition.

Input	Parallel	TTFT gain	TPOT gain	Prompt throughput gain	Decode throughput gain	Request throughput gain	Latency gain
8K	2	+3.95%	-0.01%	+0.19%	+0.19%	+0.19%	+0.19%
8K	4	+5.35%	-2.08%	-1.58%	-1.58%	-1.58%	-1.61%
8K	8	+5.57%	-1.25%	-0.68%	-0.68%	-0.68%	-0.70%
8K	16	+4.06%	-2.10%	-1.34%	-1.34%	-1.34%	-1.37%
16K	2	+3.26%	+0.92%	+1.12%	+1.12%	+1.12%	+1.10%
16K	4	+1.54%	+4.52%	+6.07%	+6.07%	+6.07%	+4.25%
16K	8	+1.27%	+3.08%	-0.49%	-0.49%	-0.49%	+2.84%
16K	16	+1.38%	+0.06%	+0.29%	+0.29%	+0.29%	+0.30%
32K	2	+6.10%	-0.60%	+0.40%	+0.40%	+0.40%	+0.39%
32K	4	+3.95%	-0.13%	+2.10%	+2.10%	+2.10%	+1.96%
32K	8	+3.63%	+1.15%	+2.65%	+2.65%	+2.65%	+1.71%
32K	16	+1.94%	+1.97%	+2.60%	+2.60%	+2.60%	+1.97%

Prefill-focused profiling confirmed that FC1+MMRS is active:

• MatmulReduceScatterV2 appears when FC1/MMRS is enabled.
• Part of the original row-parallel communication path is replaced by MMRS.
• allReduce time/count decreases in FC1-on traces.

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 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); the remaining 16K points are more sensitive to run-to-run variation and should not be treated as stable headline gains. These results reinforce that FC1 primarily targets prefill latency: a TTFT improvement does not guarantee an end-to-end throughput improvement when decode, layout conversion, and communication overhead dominate.

Current profiling limitations:

• Extra gather may be needed after ReduceScatter to restore the expected full hidden-state layout.
• MMRS has fixed overhead and is not beneficial for small decode-like shapes.
• Layout transforms and pad/unpad reduce the net gain.
• End-to-end TTFT is sensitive to scheduler batching variance.
• Quantized row-parallel paths currently skip fused MMRS.

For this reason, FC1 remains disabled by default and should be enabled only for suitable long-prefill deployments.

Reviewer Focus

Please focus review on:

• Default-off behavior for both FC1 and MMRS.
• Runtime activation gating for long-prefill only.
• MMRS input validation, comm_mode selection, and caller-side fallback behavior.
• Sequence sharding, local-tail padding, gather, and output-layout correctness.
• Qwen3 and Qwen3-Next Dense integration points.
• Runtime/model-output semantics after FC1 gathers hidden states back.
• ops-transformers 9.1.0 operator build and runtime library compatibility.

Recent Updates

MMRS wrapper simplification

• Simplified the FC1 MMRS C++ wrapper so it only validates inputs, resolves the HCCL group, selects torch_npu comm_mode, and calls npu_mm_reduce_scatter_base.
• Removed the duplicate local torch::matmul fallback. Unsupported MMRS cases return an undefined tensor and let the existing row-parallel caller use its fallback path.
• Added defensive defined() checks before reading MMRS tensor dimensions or metadata.
• Removed stale MMRS output validation and unused MMRS buffers/parameters.

Sequence-layout and model-scope review fixes

• Changed sequence sharding to allocate and copy only the local shard. Only the tail rank is zero-padded when necessary, avoiding a full padded-sequence temporary tensor.
• Consolidated gather behavior and removed the redundant gather-and-unpad helper.
• Simplified residual-layout validation and removed unreachable return paths.
• Replaced an NPU compile-time branch in the shared DenseMLP path with a runtime platform check to preserve other backends.
• Explicitly disable FC1/MMRS for models with routed experts; the current supported scope is Dense models.
• Restored the original mlp_(x) ModuleHolder invocation after FC1 context propagation moved to thread-local scope.

Configuration and diff-scope cleanup

• Removed hard-coded risky-shape magic numbers and automatic aiv to ai_cpu switching after validating the fixed ops-transformers 9.1.0 operator.
• Kept FlashComm1Context as a plain data struct and moved sequence-sharding helpers to free functions.
• Removed duplicate row-padding logic from linear.cpp and reused the shared FC1 helper.
• Restored CMake helper files, top-level CMakeLists.txt, and third-party submodule pointers to match upstream main.
• Removed noisy success-path FC1 logs and kept local build-environment inference out of the PR.

[gemini-code-assist]

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.

[Cooofish]

最新commit 971dd79已本地build test通过
100% tests passed, 0 tests failed out of 845
Total Test time (real) = 106.23 sec
All tests passed!

[Cooofish]

commit 6fa9994 本地build test已通过

I20260723 15:10:44.881584 2590209 setup.py:797]
================================================================================

I20260723 15:10:44.881612 2590209 setup.py:798]
All tests passed!

I20260723 15:10:44.881636 2590209 setup.py:799]
================================================================================

[Cooofish]

最新commit本地build test已通过
Commit: bf0a955
xllm_ops: fa7859e66694a2ece7e41492c3a8e95a415e15b5
Tests discovered: 880
Exit code: 0
All tests passed!