Conversation
* feat: paged_store_kv support varlen. * fix: return value && perf.
…nhance the error message.
* 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.
* [ilu/ixformer] support moe using gdr * [ilu/ixformer] Use bfloat16 inv_smooth_scale in quant moe when input.dtype=torch.bfloat16. --------- Co-authored-by: xiaomei.wang <xiaomei.wang@iluvatar.com>
Co-authored-by: song.liu <song.liu@iluvatar.com>
* ci: add claude code automated PR review workflow Runs on NPU self-hosted runner for each PR, posts review comments via the internal Anthropic-compatible endpoint. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: pre-cleanup workspace before checkout for claude review NPU runner workspace can hold root-owned files left by prior container jobs; checkout fails with EACCES otherwise. Wipe via docker before checkout, matching the iluvatar workflow's approach. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: harden pre-cleanup step Run as root explicitly, use rm -rf with dotglob to catch hidden files, list dir at end to verify, and drop "|| true" so failures surface instead of being silently swallowed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: override npu image entrypoint in pre-cleanup The npu CI image's ENTRYPOINT runs sshd and swallows our `bash -c`, so the cleanup never executed. Use --entrypoint bash to bypass it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: provide GitHub auth for claude-code-action Add id-token: write permission and pass github_token explicitly so the action can authenticate to the GitHub API without going through the OIDC + GitHub App flow (which fails on this self-hosted runner). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: replace claude-code-action with custom review script The official action installs claude-code from claude.com which is geo-blocked from our network (proxy egress IP triggers region check). Self-roll a small Python step instead: compute the PR diff, POST it to the internal Anthropic-compatible endpoint, post the response as a PR issue comment via the GitHub API. Requires a new repo secret ANTHROPIC_AUTH_TOKEN (Bearer token for the internal endpoint). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: address review feedback for claude review - Add concurrency group to cancel superseded reviews on rapid pushes - Skip job for fork PRs (no secret access) - Drop redundant git fetch (fetch-depth: 0 already pulls full history) - Retry LLM call up to 3x with exponential backoff - Validate response shape and surface non-text/error payloads instead of throwing KeyError Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: harden claude review against second-round feedback - Run review script from BASE_SHA, not PR head, so a PR cannot modify the script to exfiltrate ANTHROPIC_AUTH_TOKEN / GITHUB_TOKEN. Falls back to PR head only when the script does not yet exist at base (first-time bootstrap, e.g. this PR). - Truncate comment body to GitHub's 65536-char hard limit. - Bail out of retry loop on 4xx (except 408/429) so auth/format errors do not waste attempts. - Add jitter to exponential backoff. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: byte-accurate comment limit and GitHub POST error handling - GitHub's 65KB comment limit is bytes (not chars); slice the encoded utf-8 buffer and decode with errors='ignore' to avoid corrupting a multi-byte boundary. - Catch HTTPError / URLError on the comment POST and surface the response body so failures don't lose the review output silently. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * ci: tighten claude review prompt format Enforce a strict output structure (Verdict / Summary / Must fix / Suggestions / Nits / Notes) with severity tags and required path:LINE citations. Bans emoji, requires a 2-3 sentence PR summary, and folds non-blocker items behind <details>. Adds debug-residue, layering, and silent-fallback to the review checklist. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* [ilu/ixformer] update static_quant for 2d scale * [ilu/ixformer] support IxformerStorePagedKVCache chunk_metadata pass * misc: refine yml. * [ilu/ixformer] add log --------- Co-authored-by: fan.jiang <fan.jiang@iluvatar.com> Co-authored-by: xudong.zhao <xudong.zhao@iluvatar.com> Co-authored-by: wens <zhaowenshuo.oo@bytedance.com> Co-authored-by: xiaomei.wang <xiaomei.wang@iluvatar.com>
XPU-Forces#326) * [ilu/ixformer] add IxformerPagedDecodeGQAWithKVDequant/IxformerPagedDecodeSWAWithKVDequant * fix: add _get_dequant_buffers * misc: update version. * fix: add max_total_seq_len for TTXPagedDecodeSWAWithKVDequant --------- Co-authored-by: 江帆 <fan.jiang@iluvatar.com> Co-authored-by: wens <zhaowenshuo.oo@bytedance.com>
* fix: [ttx/mlu][swa] move get_aux_mask on device * reduce extra casting kernels on the device
* feat: support MojoPerfillSageGQA * fix: modify the input shape form [T, Hq] to [Hq, T] * fix IxformerPagedPrefillSageGQA forward.
…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>
![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 the MojoHcPre operator along with its NPU backend implementation (TTXHcPre) and associated Triton kernels. It also adds comprehensive accuracy and performance tests. The review feedback highlights several critical issues in the Triton kernels that need to be addressed, including compilation failures due to non-power-of-two block dimensions and tl.arange sizes, a lack of assertions enforcing supported layouts (hc_mult == 4), and missing bounds capping using tl.minimum which can lead to redundant loop iterations and compilation issues.
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.
lyujheng
force-pushed
the
deepseekv4/hc_pre
branch
from
July 2, 2026 09:39
27135f0 to
805e0b8
Compare
lyujheng
force-pushed
the
deepseekv4/hc_pre
branch
from
July 2, 2026 09:54
805e0b8 to
d1c9adc
Compare
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.
Summary
Add the
hc_preoperator implementation for DeepSeekV4 on Ascend NPU using Triton.Changes
mojo_opset/backends/ttx/kernels/npu/hc_pre.py— Triton kernel for fused head combination pre-processing (sigmoid activation, Sinkhorn normalization, pre-weighted matmul), with tle and non-tle variantsmojo_opset/core/operators/hc_pre.py— MojoHcPre operator with torch reference implementationmojo_opset/backends/ttx/operators/hc_pre.py— TTX backend binding with runtime dispatch based on tle availability and batch sizePerformance (Device Latency, 910B, HC=4, dtype=bf16)
Accuracy (bf16, HC=4, atol=5e-3, rtol=5e-3)
Note
triton_v3.2.xbranch.