[KMCompiler] [ttx/npu] add DeepSeek-V4 operator rms_norm_dynamic_quant - #416
Closed
YangLong114514 wants to merge 330 commits into
Closed
YangLong114514 wants to merge 330 commits into
YangLong114514 wants to merge 330 commits into
Conversation
* add npu backend * add npu backend * add npu backend * add npu backend * add npu backend * add npu backend * add npu_gelu --------- Co-authored-by: mc-zhang <zhangbolun6@huawei.com>
…_device_guard feat(ttx/kernels/utils.py): add tensor device guard for triton kernels.
…e_unnecessary_code chore(ttx/kernels/npu): remove unnecessary code.
* feat: optimize conv1d fwd&bwd. * fix: test shape.
…ule_like feat(normalization): make normalization nn.Module like.
…-Forces#120) This reverts commit 94561d4.
…e." (XPU-Forces#120)" (XPU-Forces#122) This reverts commit ca98706.
* feat: refactor rope. * fix: test base. * fix: graph test.
Co-authored-by: lizichong <756066299@qq.com>
…t_tensor_factory_args feat(mojo_opset/core/operator.py): support torch.empty factory args.
…Forces#129) * fix: wrong logic in platform.py, add register fallback warning. * fix: change warning to debug.
* Refine ttx activation kernels and tests * Refactor normalization api and support patching rmsnorm * chore: specify backend for ttx tests
添加forward_diff_with的百分比对比功能 . 框架报错人性化提示 添加forward_diff_with的百分比对比功能，并提取函数优化代码 格式化代码 Update operator.py 修改精度校验方法 . .
…orces#113) * Optimize causal_conv1d_update_kernel_bdt_fwd ttx kernel * Remove useless static assert * Fix accuracy problem * Fix ut test error
* refine some core api * Update README * Add MoEGating torch impl
* Add MoE interface * fix some
…computed tokens in kv_cache (XPU-Forces#130) * feat(paged_prefill_attention): refactor paged_prefill_gqa to support computed tokens in kv_cache * with computed tokens in kvcache, kv should be longer than q * refactor reference impl and tests to support inequal q_len and kv_len * refactor ttx impl of paged_prefill_gqa into persistent-kernel style * fix(typo): fix typo in attention.py * fix: fix gqa_layout for paged_attention_prefill * fix: lift dtype of qk before scale * feat: make seqlens_kv optional for paged_attention_prefill * feat: add optional mask for paged_prefill_gqa
…XPU-Forces#335) * refactor: align naming convention, change max_xxx_lens to max_xxx_len * Update mojo_opset/tests/accuracy/operators/test_attention.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update mojo_opset/tests/accuracy/operators/test_attention_cudagraph.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: xiaomei.wang <xiaomei.wang@iluvatar.com>
* feat: add attention_gate core impl. * misc: add tests. * feat: add MojoFusedAttnOutputGate. * [ilu/ixformer] support MojoFusedAttnOutputGate --------- Co-authored-by: xiaomei.wang <xiaomei.wang@iluvatar.com>
…rces#330) * Enable torch deterministic config for mojo deterministic mode * Update curl command to use proxy for downloading * revert last Removed proxy option from curl command for downloading ixformer wheel. --------- Co-authored-by: wwens7 <zhaowenshuo.oo@bytedance.com>
* WIP: add shmem_manager & support compute overlap comm ops
* feat(ttx): add triton-dist fused comm+compute operators for Ascend NPU
Add TTXAllGatherGemm, TTXGemmAllReduce, and TTXGemmReduceScatter backends
that fuse GEMM with collective communication via aclshmem on Ascend NPU.
Ported from triton-dist-package reference kernels and verified on 2-card setup.
Also fixes libentry import issue across all NPU triton kernels by adding a
try/except fallback in npu/utils.py, and rewrites multi-card comm tests to
use torchrun instead of mp.spawn.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* tmp
* refactor(runtime): unify TTX shmem management via MojoSymmetricMemoryManager
- Add backend="ttx" support to MojoSymmetricMemoryManager with unique_id
bootstrap mode (eliminates ASH_MASTER_ADDR/PORT env var requirement)
- Move shmem preload (.so conflict workaround) into centralized _init_ttx_backend
- Delete _ensure_ash_init from kernel layer (was runtime concern, not kernel)
- TTX operators now use runtime.get_backend_manager() + runtime.allocate_peer_mem()
- Add proper finalize path via close() for TTX backend
- Add bfloat16 support to GemmAllReduce and GemmReduceScatter kernels
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: address code review feedback
- Use torch.npu.current_device() instead of self.rank for device_id
(fixes multi-node environments where rank != local device index)
- Fix typos: aclshmem_finialize → aclshmem_finalize, destory → destroy
- Fix README typo: mojo_opsetutils → mojo_opset.utils
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: address Claude Code review feedback
- Simplify libentry fallback to direct definition (clearer intent)
- Remove debug test scripts from repo root (test_ag_gemm_torchrun.py,
test_ttx_allgather_gemm.py)
- Consolidate try/except ImportError for triton-dist kernels with
explicit None defaults to avoid confusing downstream errors
- Use output.add_(self.bias) instead of allocating new tensor
- Add comments explaining why torch.zeros (not empty) is required
for atomic_add kernels
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(test): skip quant comm tests when backend lacks implementation
Add @bypass_not_implemented to single-rank quant comm tests so they
are skipped (not failed) when MOJO_BACKEND=ttx on MLU/ILU CI where
no TTX implementation exists for MojoQuantGemmAll2All/MojoAll2AllQuantGemm.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* refactor(test): remove single-rank and gloo tests from test_compute_with_comm
This file is dedicated to multi-card comm+compute fusion testing.
Remove single-rank parametrized tests (no real comm exercised, redundant
with other test files) and skip all tests in CI (require triton-dist +
multi-NPU environment with HCCL).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: address second round Claude Code review
- Narrow ImportError catch: gate on `import triton_dist` specifically,
log at debug level when unavailable
- Replace pytest.mark.skip with pytest.importorskip("triton_dist") so
tests auto-enable when the dependency is available
- Read MOJO_TTX_SHMEM_SIZE_MB at call time (not module import) so tests
can override via env
- Remove stale TODO comments (replaced by importorskip mechanism)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(ttx): lazy-import npu utils to avoid loading npu kernels on ILU
Move `from mojo_opset.backends.ttx.kernels.npu.utils import get_num_cores`
from module-level to inside _ensure_shmem methods. This prevents ILU CI
from triggering npu kernel module load (which fails on `tl.gather`).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(ci): tolerate pytest exit code 5 when triton-dist is unavailable
test_compute_with_comm.py uses pytest.importorskip("triton_dist") which
skips all tests when the package is not installed. Pytest returns exit
code 5 for "no tests collected" which fails CI. Allow exit code 5 as
a valid outcome (all-skipped is expected until triton-dist is in CI).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: address third round code review
- Fix CI exit code handling: `rc=$?; [ $rc -eq 0 ] || [ $rc -eq 5 ]`
correctly distinguishes "all skipped" (5) from real failures (1)
- Fix _ensure_shmem: always delegate to runtime.allocate_peer_mem()
(handles grow-on-demand) instead of early-returning on stale _peer_mem.
This prevents undersized buffer if operator is reused with larger K.
- Store self._runtime for consistent access across methods
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* Skip distributed test for temp
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: song.liu <song.liu@iluvatar.com>
…mask, simplify kernel structure (XPU-Forces#321) - Replace external aux_mask (pre-allocated triangular mask tensor) with inline causal mask computation directly in kernel - Remove redundant inner KV block loop for full pages (PAGE_SIZE == BLOCK_N), loading entire page in single iteration - Remove boundary_check for full-page loads to enable async SME copy path - Add num_stages=2 hint to full-page loop for pipeline pass trigger - Update autotune configs: larger BLOCK_M (256) and more warps (8/16) - Remove unused parameters: aux_mask_ptr, mask strides, AUX_MASK_SIZE, USE_AUX_MASK
…ces#338) Use pinned non-blocking H2D for group_offsets to avoid draining the GPU queue each call, and add small-M autotune tiles (BLOCK_M=16) so large-N MoE shapes pick a bandwidth-friendly config. Drop MAX_M from the autotune key to avoid re-tuning on varying per-group row counts.
Co-authored-by: song.liu <song.liu@iluvatar.com>
…#346) * refactor: distinguish between fused moe and non-fused moe * Update mojo_opset/core/operators/moe.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * refactor: update ixformer moe / quant_moe * refactor: support ep for fused moe (XPU-Forces#350) * refactor: distinguish between fused moe and non-fused moe * Update mojo_opset/core/operators/moe.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * refactor: update ixformer moe / quant_moe * refactor: support EP for fused MojoMoE / MojoQuantMoE * fix: fix group gemm calls * fix: unify cudagraph settings for groupgemm * ci: add tests for fused ep moe --------- Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [ilu/ixformer] refine: EP of fused moe. * fix: support moe ep with dp inputs * fix: revert ixformer calls to adapt old versions * ci: fix ci * fix: update max init of swa * Update flash_attention.py --------- Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: xiaomei.wang <xiaomei.wang@iluvatar.com>
…#345) * ci: fix tests for experts and quantexperts * Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [ilu/ixformer] fix moe ops. * [ilu/ixformer] update ixformer version. * test(moe_ep): validate small-op composed nn.Module matches MojoMoE(torch) Adds _SmallOpMoEModule — a plain nn.Module assembled from MojoMoEGating / MojoMoEDispatch / MojoExperts / MojoMoECombine — and asserts its forward matches MojoMoE(backend='torch') under both EP=1 (plain pytest) and EP=2 (torchrun --nproc-per-node=2). In the EP path, expert_outputs is padded back to [num_tokens * top_k, hidden] with zeros at non-local positions because ixformer's moe_combine kernel rejects sliced tensors; the per-rank reduce yields the same total as the torch reference. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * doc: add docs for MoEDispatch --------- Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: xudongzhao1006 <690335895@qq.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* [GENESIS-6582]feat(triton):swa supprt 4 dim tensor q * [GENESIS-6582]feat(triton):swa review comand * Update mojo_opset/backends/ttx/kernels/mlu/swa.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [GENESIS-6582]feat(triton):swa review comand -1 * Update mojo_opset/backends/ttx/kernels/mlu/swa.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update mojo_opset/backends/ttx/kernels/mlu/swa.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix: raise for not supported platforms * fix: fix decode-N step mask for swa * fix: separate decode-n swa * Revert "fix: raise for not supported platforms" This reverts commit dbea3c8. * fix: keep core ops clean * suport casul ask * clean code * fix ci failed * fix review comand --------- Co-authored-by: zhouronghai <zhouronghai@cambricon.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Neuromancer42 <chenyifan_1997@hotmail.com>
tl.cast do not support overflow_mode parameter
… extraction (XPU-Forces#318) Replace full argsort in _topk_stage1_kernel and _topk_merge_kernel with iterative max/min extraction, only outputting k elements per chunk.Remove _compact_sorted_blocks as it is no longer needed.
* feat: [MLU] add mlu experimental ops; fix post norm bug * ci: fix ci input gen * fix: fix torch_npu fused residual_add_rms_norm (norm_pos = 'post') * Update mojo_opset/experimental/operators/normalization.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com> Co-authored-by: Yifan Chen <chenyifan_1997@hotmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
…XPU-Forces#367) * [KMCompiler] Optimize FuseAddRmsNorm performance * [KMCompiler] Format FuseAddRmsNorm implementation * [KMCompiler] Address FuseAddRmsNorm review feedback
Co-authored-by: chenyifan.42 <chenyifan.42@bytedance.com>
* [lvzheng/ttx] Optimize _layernorm_fwd_kernel performance Refactor kernel into two stages: Stage 1 computes mean and variance in one pass; Stage 2 performs normalization and writes output. * [KMCompiler] Optimize layernorm kernels via single pass and nomask * [KMCompiler] Optimize Layernorm by tune NPU LayerNorm forward tiling and skip infer stats stores * [KMCompiler] Address LayerNorm review feedback --------- Co-authored-by: lvzheng <lyujheng@gmail.com>
…es#365) * [KMCompiler] Optimize silu with rowwise nomask kernels * [KMCompiler] Address SiLU review feedback
* [KMCompiler]opt for rms_norm * review update
…U-Forces#385) `embedding_nf4_dequant_impl` used torch.empty to allocate the output tensor. When an input id was out of vocabulary (>= vocab_size or < 0), the kernel's token_mask went False and `tl.store` was masked out — so the corresponding output row was never written. The row silently kept whatever bytes torch.empty had handed us from the allocator. On some runners those bytes happened to be near zero (test passed), on others they weren't (test failed with Max absolute difference of several thousand at exactly one row). This made `test_embedding_nf4_dequant_impl` flake in a way that only reproduced on specific CI hosts. Switch to torch.zeros so OOB rows deterministically read back as zero — matching the contract the reference implementation exercises via `expected = torch.zeros(...); expected[valid_mask] = ...`. Non-OOB behavior is unchanged: the kernel overwrites the zeros with the real dequantized values.
…es#364) * [KMCompiler] Optimize gelu with rowwise nomask kernels * [KMCompiler] Address GeLU review feedback * [KMCompiler] Fix GeLU libdevice import fallback
* fix bug of a: tl.constexpr=tl.cdiv(b, c) * Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: zhouronghai <zhouronghai@cambricon.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
* modify compiler_hint path to tl.extra.cann.extension; tl.max add propagate_nan=tl.PropagateNan.ALL; * swa add enable_ubuf_saving to solve ub overflow * normalization use casting_mode gemma to avoid F.rms_norm fp32 cast precision * remove print * flash attention use need_mask to do if else * Revert "flash attention use need_mask to do if else" This reverts commit d705131. * modify test_over_encoding to switch to triton-ascend * modify n_gram mask to solve random precision problem * switch to triton-ascend 3.2.1, use wget temprorarily for checking CI is right * switch to triton-ascend 3.2.1, use wget temprorarily for checking CI is right * switch to triton-ascend 3.2.1, use wget temprorarily for checking CI is right * add print to debug CI * modify ccec path * add CI debug * add CI debug * add CI debug * rollback ci and switch to cann8.5.0 image * add triton-ascend on CI * fix inder quant para error * fix perf test case * modify extract_slice _compute_vision_rope to adapter triton-ascend * add sync_solver=False to avoid groupgemm perf descend on triton-ascend * switch to byted-triton-x 3.2.1 * add --index-url to switch to byted-triton-x 3.2.1 * add --index-url to switch to byted-triton-x 3.2.1 * Update mojo_opset/backends/ttx/kernels/npu/over_encoding/n_gram.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * rollback to rmsnorm_fwd llama mode, triton-ascend 3.2.1 precision is OK * functions rms backward precision is not OK, change to gemma * no need to add propagate_nan to all * groupgemm: delete dot_pad_only_k hint and tl.multibuffer * modify ci and pyproject * use pyproject.toml to pip install --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces MojoRMSNormDynamicQuant, a fused RMSNorm and dynamic per-token quantization operator, along with its TTX backend implementation and corresponding Triton kernels for NPU. The feedback highlights critical issues in the Triton kernels, specifically potential out-of-bounds memory writes in unmasked kernels when the hidden dimension is not aligned to BLOCK_N_ALIGN. Additionally, the reviewer recommends adding division-by-zero guards in the prefill kernels and utilizing propagate_nan=tl.PropagateNan.ALL in tl.max operations to prevent performance overhead and align with NPU hardware behavior.
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.
YangLong114514
changed the title
mazixuan-byted
force-pushed
the
master
branch
from
September 14, 2026 15:54
7defc82 to
df2f3da
Compare
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
19 participants
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
Add TLE Triton implementation of MojoRMSNormDynamicQuant for DeepSeek-V4 on Ascend 910B.
Changes
Performance
Test environment: Ascend 910B, TLE Triton backend.(cann 9.0)
The above perf data was obtained from testing using the perf-rms_norm_dynamic_quant.py.py script.
Accuracy Test
All accuracy test cases passed.