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 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 \
--flashcomm1_min_prefill_tokens=8192 \
--mmrs_comm_mode=aiv

More details about design, usage, suitable workloads, 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 NPU matmul_reduce_scatter kernel 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 model paths.
• Use scoped FC1 context for Qwen3-Next layers instead of threading an extra fc1_ctx argument through all forward signatures.
• 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 in decode. TPOT/decode throughput are not expected to directly benefit from this feature.

MMRS Integration

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

• 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.
• If a specific shape hits an AIV AICore issue, the path can fall back to ai_cpu through the existing mode selection logic.
• Unsupported MMRS cases return to the caller-side fallback path instead of hiding fallback behavior inside the wrapper.

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.
• 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.

Scope Notes

This PR intentionally keeps FC1 opt-in and avoids unrelated cleanup.

Excluded from this PR:

• RFC/design documents.
• Temporary local test files.
• CMake helper changes unrelated to FC1.
• Third-party submodule pointer changes.
• Generated third-party build directories.

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 \

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 matrix now contains all 72 individual measurements, covering 36 complete FC1-off/on comparisons. The prior AICore failures were resolved after deploying compatible custom-operator builds.
• 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/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 latest 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:

• 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 such as concat/transpose 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 shape checks, comm_mode selection, and fallback behavior.
• Sequence sharding, padding, gather, and output-layout correctness.
• Qwen3 and Qwen3-Next integration points.
• Runtime/model-output semantics after FC1 gathers hidden states back.

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 duplicate local torch::matmul fallback from the wrapper. Unsupported MMRS cases now return an empty tensor and let the existing row-parallel fallback remain in the caller.
• This keeps the integration closer to the upstream torch_npu operator and reduces review surface.

Gemini review fixes

• Kept FlashComm1Context as a plain data struct and moved sequence-sharding helpers to free functions.
• Removed duplicate row-padding helper from linear.cpp and reused the shared FC1 helper.
• Added defensive defined() checks before inspecting MMRS input tensor metadata.
• Verified that the previously flagged production paths already use ModelConfig::get_instance().backend(), EPLBConfig::get_instance().expert_parallel_degree(), preserve the ep_size == 1 single-reduce MoE optimization, and no longer contain the flagged Qwen3 ATB int64_t hostData vectors.

Diff-scope cleanup

• Restored CMake helper files and third-party submodule pointers to match upstream main.
• Restored top-level CMakeLists.txt to match upstream main.
• Removed unused MMRS parameters and noisy success-path FC1 logs.
• The PR now keeps CMake helpers, third-party pointers, and local build-environment inference out of the FlashComm1 change set.

[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.