Repository navigation
Conversation
Use tensor index generation for NPU-resident counts with at least 64 source/expert groups. Keep the CPU reference for smaller and host count groups, and move indices to non-CPU token devices before reordering. Ascend 950PR standalone reference and gradient checks: 21 passed, 26 CUDA-only skips. Changed-file pre-commit passed. No end-to-end speedup claim. Signed-off-by: hazelduan <duanxueqi@outlook.com>
6 tasks done
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
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.
AutoEP's NPU permutation path needs a device implementation instead of invoking CUDA Triton or copying large source/expert count groups to the host. This change builds NPU permutation indices with PyTorch tensor operations when NPU-resident counts contain at least 64 source/expert groups. Smaller groups and host counts use the CPU reference path; indices and aligned counts are transferred to the token device before reordering.
It preserves ordering, padding and token gradients, including empty experts and strided inputs. The threshold is internal; no configuration or optional package is added. Other dispatch steps can still synchronize with the host.
Split from #8784; no Qwen or HiFloat8 dependency. Related RFC: #8777. Base: 3b80932; validated source: b7bee89. An open-PR search for NPU permutation returned no match on 2026-10-14. PR #8784 retains an NPU host correctness fallback, so dispatch selection overlaps: when combining, retain this threshold selection and the shared non-CPU device-transfer guard.
Validation:
PYTHONPATH=<isolated-permutation-repo> python -m pytest tests/unit/v1/moe/test_autoep_permute_rows.py -q -rs.The private combined-stack validation also completed Qwen3.5-35B-A3B text-only EP4/ZeRO2 for 30 matched steps. All 40 layers execute; only routed/shared experts in layers 0/39 train (12 tensors/rank), batch 1/GAS1, sequence length 256, eager attention/experts, no activation checkpointing or checkpoint save, seed 42, AdamW 1e-5/cosine, same model weights and 400-row data. Every rank has exact input hashes, loss, gradient norms, native operator counts and all final trainable tensors. Maximum-rank steady-step median (excluding 5 warm-up steps) is 4.1886s baseline vs 4.1906s candidate (+0.046%); training NPU peak is 33,487,014,400 bytes and initialization peak is 86,776,612,352 bytes in both. This combined-stack result does not attribute an independent speedup or numerical improvement to this split.
Native commands and logs are retained at
/data/dxq/hif8_a5_20260914/experiments/quantize_convergence_20261014_01:bash native_gates.sh,bash paired_ep4.sh,qwen_native_stack.log,permutation_native.log,comparison_ep4_baseline_r1_ep4_candidate_r1.json. The runner bind-mounts the original Conda environment (dxq_hif8_torchao_py311) and CANN/driver; it changes no global packages or device permissions. Standalone runtime file hashes match each immutable source revision below. The model run uses a private composition of main PR and split runtime hunks, with the overlapping dispatch hunks reconciled; this composition is not the main PR contents.Both model comparison variants include permutation optimization, so the paired run does not measure its independent benefit. The standalone tests establish index generation and local reordering. CUDA/other accelerator execution, distributed/offload combinations, activation checkpointing and isolated end-to-end performance/memory effects remain untested. No speedup is claimed.
AI assistance was used to implement and review this change; the human submitter owns review and validation.