feat: use post4 libtorch_npu with early Python init and native ProcessGroupHCCL.#2006
Draft
yingxudeng wants to merge 2 commits into
Draft
feat: use post4 libtorch_npu with early Python init and native ProcessGroupHCCL.#2006yingxudeng wants to merge 2 commits into
yingxudeng wants to merge 2 commits into
Conversation
…sGroupHCCL. Replace the C++ custom-op TP workaround with a native Python ProcessGroup approach that mirrors the CUDA path: - xllm.cpp: early Py_Initialize + import torch_npu before workers start, so post4 libtorch_npu's PyGILState_Ensure in empty_with_format() has a valid interpreter. Dynamic device index, exception-safe monkeypatch, narrow RuntimeError catch, guarded PyEval_SaveThread, CHECK aclInit. - collectives.py: full Python implementation using ProcessGroupHCCL for NPU and ProcessGroupNCCL for CUDA, registered via @custom_op. - py_causal_lm.cpp: call Python init_tp_group() like the CUDA path instead of the C++ set_tp_process_group() hack. - npu_ops_library.cpp: remove all_reduce_/all_gather C++ ops (now Python), add update_decode_graph_metadata NPU impl (tensor copy). - npu_paged_attention.py: use PyTorch index_put for KV cache write (thread-safe vs ATB ReshapeAndCache singleton). Causal mask grows on demand for seq_len > 2048. - device.cpp: conditionalize init_device_context — Python path uses SetDevice only, C++ path retains full torch_npu::init_npu. - master.cpp: force npu_kernel_backend=TORCH for python model_impl (ATB path skips tp_group creation). Restore in-process multi-device TP for Python executor (new multi-process arch takes single device per process, but Python model needs all visible devices). - executor.py: treat "eager" as graph-off to avoid inductor fallthrough. Verified: Qwen3-4B TP2 --model_impl=python --python_graph_backend=eager on dyx-xllm-cann9-main with post4 libtorch_npu.so (71MB). Inference output correct across Chinese/English/code tests.
yingxudeng
requested review from
Clement-Wang26,
DongheJin,
DragonFive,
JimHsiung,
Kang-Meng,
RobbieLeung,
XuZhang99,
liujinguang0125,
liutongxuan,
walsonyang,
xiao-yu-chen,
yq33victor and
zhang-minchao
as code owners
July 22, 2026 10:36
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
yingxudeng
marked this pull request as draft
July 22, 2026 12:50
Add scripts/deps/torch_npu_install.py mirroring the TileLang pattern: detect installed torch_npu version via importlib.metadata, download and install the post4 wheel from JD storage if missing or mismatched. Handles pip 26+ strict PEP 440 filename validation by downloading to a temp dir with a compliant filename before pip install. Integrated into setup.py as _ensure_torch_npu_ready(), called before _ensure_tilelang_ascend_ready() in the NPU build path.
yingxudeng
commented
Jul 22, 2026
| TORCH_NPU_WHEELS: dict[tuple[str, str], str] = { | ||
| # (cann_version, arch) → wheel URL | ||
| ("9.0", "arm"): ( | ||
| "http://storage.jd.local/bmp-packages/NPU_CICD/compile-base/" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Related Issues
Change Type
Pull Request Checklist
Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.
PR Title and Commit Messages
<type>: <subject>.Pre-commit Checks
pre-commitby runningpip install pre-commitor an equivalent command.pre-commit install.pre-commit run --all-filesand fixed any reported issues.Self Review
.agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.mainbranch.Build and Test Coverage
python setup.py build testhas passed on a CUDA machine.python setup.py build testhas passed on an NPU machine.python setup.py build testhas passed on an MLU machine.Reviewer Notes