add deepep and test, modify test_moe_quant and test_attention_quant - #349

Closed

qiushi13 wants to merge 324 commits into

heqiushi/moe_quant

This pull request is big! We're only showing the most recent 250 commits

Commits on Jan 26, 2026

refactor(ttx/kernels/utils.py): cover the autotuners.
LeoLau94
committed
73ad1ba
test(tests/base/test_tensor_device_guard): add more test cases.
LeoLau94
committed
b5e9e31
fix(ttx/kernels/utils): remove the debug code.
LeoLau94
committed
6362fa8
feat(kernels/utils): default to disable the tensor device bakayoky.
LeoLau94
committed
22b0c21
feat(mojo_opset/utils/patching.py): rewrite the python assertion to enhance the error message.
LeoLau94
committed
35e9691
refactor(mojo_opset/utils/patching.py): add control through environment varaible.
LeoLau94
committed
b900f67
test(tests/base): add unit test for assertion rewrite.
LeoLau94
committed
43a42d4
add npu backend && activation ops (#106)

old-steel
and
mc-zhang
authored
be6da0b
Merge pull request #101 from XPU-Forces/liuyuan/feat/tensor_device_guard
LeoLau94
authored
d4e3a2e
chore(ttx/kernels/npu): remove unnecessary code.
LeoLau94
committed
934d8c4
Merge pull request #112 from XPU-Forces/liuyuan/chore/remove_unnecessary_code
LeoLau94
authored
9f850cb

Commits on Jan 27, 2026

feat(normalization): make normalization nn.Module like.
LeoLau94
committed
94561d4
feat: optimize conv1d fwd&bwd. (#107)
wwens7
authored
e343107
Merge pull request #114 from XPU-Forces/liuyuan/feat/nn_module_like
LeoLau94
authored
2d2b389
feat: add top-k-sampling ttx and torch implement
YYYYimo
committed
ad2cb3a
Revert "feat(normalization): make normalization nn.Module like." (#120)
LeoLau94
authored
ca98706
Revert "Revert "feat(normalization): make normalization nn.Module like." (#120)" (#122)
LeoLau94
authored
6c2b6aa
refactor: rope. (#119)
wwens7
authored
7894a2e
feat(mojo_opset/core/operator.py): support torch.empty factory args.
LeoLau94
committed
fcd7234
Merge branch 'master' into yinzijun/topk
YYYYimo
authored
3370275
fix: change super().__init__(...) into non-param call version
YYYYimo
committed
9d83436
chore(normalization): add more annotations.
LeoLau94
committed
d0faa7b
fix: register top_k_sampling_impl and remove comparison with torch
YYYYimo
committed
5c55011

Commits on Jan 28, 2026

Commits on Jan 29, 2026

Commits on Jan 30, 2026

feat: optimize causal_conv1d_update_kernel_bdt_fwd ttx kernel. (#113)
zhanzy178
authored
754b9e2

Commits on Feb 2, 2026

Clean code (#135)
zhangjihang-BD
authored
bb1e20c

Commits on Feb 4, 2026

Add MoE interface (#140)
zhangjihang-BD
authored
8324c88

Commits on Feb 5, 2026

feat(paged_prefill_attention): refactor paged_prefill_gqa to support computed tokens in kv_cache (#130)
yfchen-byted
authored
bb4d948

Commits on Feb 6, 2026

Commits on Feb 8, 2026

fix: change default gqa layout of Paged Attention (#147)
yfchen-byted
authored
ad54d5f
fix: update paged_decode_attn to support different gqa_layout (#149)
yfchen-byted
authored
1834375
Update convolution tiling size T_CHK_SIZE to 256 to use double buffer (#146)
zhanzy178
authored
91c1e72
Refactor core linear api (#150)
zhangjihang-BD
authored
fee573c

Commits on Feb 11, 2026

TTX support import mlu kernel (#152)
zhangjihang-BD
authored
bbf941e
fix: ci network. (#154)

wwens7
and
zhangjihang-BD
authored
5b5a3ed

Commits on Feb 12, 2026

fix: store_lowrank op. (#148)

Leetrytry
and
Workman-Lee
authored
799f74b
chore(*): add bakayokes. (#153)
LeoLau94
authored
6046e40
Refine ttx backend (#155)
zhangjihang-BD
authored
c05b534
feat: refine kv_cache, moe impl, rope support varlen & rope_cut. (#134)
wwens7
authored
52be2a3
Refactor tests and support multi-process to speedup (#151)

zhangjihang-BD
and
wwens7
authored
77643c5

Commits on Feb 25, 2026

feat: add MOJO_DETERMINISTIC env var, refine RoPE. (#157)
wwens7
authored
1bc4581
feat: add mojo deepseek v3 model (#144)
Hailey-Zh
authored
7d6cc9e

Commits on Feb 26, 2026

add seed-oss model for mojo_opset (#142)
2052580
authored
9a162b9

Commits on Feb 27, 2026

Fix formatting in README.md table (#159)
zhangjihang-BD
authored
e77b1aa

Commits on Mar 3, 2026

misc: add modeling utils. (#161)
wwens7
authored
974d4c0
fix: correct config format after Triton version update. (#165)
wwens7
authored
4377800
fix: amd platform donot register torch implement (#166)
zhangjihang-BD
authored
c795c91

Commits on Mar 4, 2026

feat(ttx/kernels/npu): implement two-stage top-k sampling fused kernel with bitonic sort; further refinements needed [skip ci]
YYYYimo
committed
9ea91b9

Commits on Mar 5, 2026

feat(modeling/common.py): add common infra for modeling. (#164)
LeoLau94
authored
5696089
fix: remove redundant code. (#167)
wwens7
authored
3a34032

Commits on Mar 6, 2026

add wan2.2 model for mojo_opset (#136)
2052580
authored
089b902

Commits on Mar 10, 2026

fix: pass accuracy test with poor performance
YYYYimo
committed
8b76781
feat: optimize page attn, split Q chunks inside kernel
chenyifan.42
committed
969b980
Merge pull request #168 from XPU-Forces/chenyifan/feat/optimize_page_attn
yfchen-byted
authored
f1c5d45

Commits on Mar 13, 2026

Add MojoLinear back, remove some useless files (#171)
zhangjihang-BD
authored
a78a32c

Commits on Mar 16, 2026

Refactor some ops (#173)
zhangjihang-BD
authored
8ad9eb0

Commits on Mar 18, 2026

Commits on Mar 19, 2026

fix: fix paged decode attn ttx kernel (#178)
yfchen-byted
authored
8959d71

Commits on Mar 23, 2026

feat: add top-k-sampling ttx and torch implement
YYYYimo
committed
09e23db
fix: change super().__init__(...) into non-param call version
YYYYimo
committed
502eb06
fix: register top_k_sampling_impl and remove comparison with torch
YYYYimo
committed
e255b7e
fix：add mask
YYYYimo
committed
4b058f9
feat(ttx/kernels/npu): implement two-stage top-k sampling fused kernel with bitonic sort; further refinements needed [skip ci]
YYYYimo
committed
24ea1e3
fix: pass accuracy test with poor performance
YYYYimo
committed
b2b6838
bug: compile fail
YYYYimo
committed
2592b88
add test script
YYYYimo
committed
677a420
Merge branch 'yinzijun/topk' of github.com:XPU-Forces/mojo_opset into yinzijun/topk
YYYYimo
committed
9e7862e

Commits on Mar 24, 2026

feat: Add torch_npu ops to replace plugin (#127)
Florence-Li
authored
d5d1503

Commits on Mar 25, 2026

Add core ops (#182)

authored
64b56f7
still have some bug with accuracy
YYYYimo
committed
7313718

Commits on Mar 26, 2026

feat(ttx/kernels/npu): implement two-stage top-k sampling fused kernel with sort
YYYYimo
committed
74fe820

Commits on Mar 27, 2026

Adapt for iluvatar，already test accuracy ok for test_normalization.py (#184)

madengfei
and
yao.guo
authored
a51f1de
Merge remote-tracking branch 'origin/master' into yinzijun/topk
YYYYimo
committed
411af54
remove unused test files
YYYYimo
committed
0705973
misc: add ilu ci yml, fix some. (#188)

wwens7
and
HunterHantao
authored
d87da2c
feat: support npu graph. (#180)
wwens7
authored
d38884a
feat: add sliding windows support for attention fwd/bwd (#145)
yfchen-byted
authored
82fc2f0
fix test framework and register operator
YYYYimo
committed
7f3081e
feat: support mojo parallel (#186)

authored
adec54e

Commits on Mar 30, 2026

Commits on Mar 31, 2026

feat: add grouped norm ops (#193)
yfchen-byted
authored
eba1e64
[GemmDequant] impl int8 gemm for npu ttx & skills (#187)
TaochunChern
authored
6a38894
misc: add mlu ci. (#194)
wwens7
authored
0185593
fix: update topk status and remove duplicate topp in Readme
YYYYimo
committed
cd4efbf
Merge pull request #195 from XPU-Forces/yinzijun/fixreadme
LeoLau94
authored
5c157d9
feat: add MojoGeneratePosEmbs and MojoApplyRoPE
lipengfei1409
committed
2df8780

Commits on Apr 1, 2026

[GemmDequant] fix register (#197)

TaochunChern
and
chenyifan.42
authored
3d102db
feat(*): support quantized operators for torch_npu backend. (#191)

LeoLau94
and
wwens7
authored
d1fc8e4
refactor: separate kv_pos_emb from rope
chenyifan.42
committed
040bc1a
fix: add ttx impl of gen_rot_pos_embed
chenyifan.42
committed
4b41936
fix: fix torch_npu rope
chenyifan.42
committed
4e00969
fix: fix failed tests
chenyifan.42
committed
bae82bf
feat: support Indexer. (#196)
wwens7
authored
1b5a6c8
fix: update api of generate pos emb
chenyifan.42
committed
65b57d0
Merge branch 'master' into lipengfei/test_rope
chenyifan.42
committed
22fb2c5
fix: fix rope in indexer
chenyifan.42
committed
a85255f
misc: refactor test_moe. (#203)
wwens7
authored
7fe39dd
Merge branch 'master' into lipengfei/test_rope
chenyifan.42
committed
e13a4b4
fix: fix ttx_graph for mojo rope
chenyifan.42
committed
d787209

Commits on Apr 2, 2026

style: clean code
chenyifan.42
committed
aab51e4
style: rename rope to applyrope
chenyifan.42
committed
75927c9
fix: add support for large PAGE_SIZE of paged prefill/decode gqa/swa
chenyifan.42
committed
d04d76a
style: clean code
chenyifan.42
committed
718d038
refactor: separate RotaryEmbedding generation and ApplyRoPE
yfchen-byted
authored
9b1f851
ci: raise NotImplementError inside operator for torch_npu backend
chenyifan.42
committed
6b081c6
Merge pull request #208 from XPU-Forces/chenyifan/fix_bigger_page
yfchen-byted
authored
27b3a5a

Commits on Apr 3, 2026

fix: update ci for moe gating (#211)
yfchen-byted
authored
171a152
enhance：move runtime utils from modeling (#200)

Xu-pan
and
mc-zhang
authored
7e84a45

Commits on Apr 7, 2026

feat: support qkv & swiglu mlp parallel. (#192)
wwens7
authored
1e06cee
style: update comments and docstrings of moe and rope for clarification
chenyifan.42
committed
405713e
style: update according to review
chenyifan.42
committed
6c04c50
fix: fix support for apply-rope of padded_prefill scenario
chenyifan.42
committed
8c6edf2
fix: fix pos_emb dim for torch_npu apply_rope
chenyifan.42
committed
98dbb7a
fix: skip unsupported case of torch_npu apply_rope
chenyifan.42
committed
21d8a49

Commits on Apr 8, 2026

Merge pull request #214 from XPU-Forces/chenyifan/style/clean_rope
yfchen-byted
authored
ed80291
Add mlu backend (#156)

ranjola
and
lihuanran
authored
457d733
fix: fix nan for padded input
chenyifan.42
committed
db4fc47
feat: [NPU] add backend test of XopsMoEGating, XopsFusedSwiGLUMoEScaleDynamicQuantize, XopsMoEInitRoutingDynamicQuant

OpenDarrenlu
authored and
yfchen-byted
committed
9f7fbe7
style: remove unused import
yfchen-byted
committed
fc7edcc
feat: initial version of mlu swa (#210)

authored
8ed7aaa

Commits on Apr 9, 2026

[Ilu backend] Fix the implement of normalizations and add other ops. (#199)

authored
1928802
mis: add logger level description in readme. (#224)
wwens7
authored
f628cfc
[ilu] Add ixformer backend, layernorm and rmsnorm as examples. (#223)

authored
05fb4f5
ci: update ci styles and distinguish platform and device (#222)
yfchen-byted
authored
f270c05
feat: add initial version of mlu group_rmsnorm triton backend (#217)

authored
bd72ebb

Commits on Apr 10, 2026

Commits on Apr 13, 2026

add runtime.py: decouple runtime states with modeling (#228)
K13Byte
authored
813d66f

Commits on Apr 14, 2026

fix[ep]: fix get ep_rank error from ep_mesh (#231)
lipengfei1409
authored
76546d2
fix: clean code of rope functions (#230)
yfchen-byted
authored
d02a202

Commits on Apr 15, 2026

misc: update Docker image version for ilu. (#232)
wwens7
authored
8a8fd6d
fix qkv_partition (#233)
K13Byte
authored
c78ed58
optimize swa kernel for mlu (#227)

authored
51668ee

Commits on Apr 16, 2026

feat: [MLU] adapt xops backend; fix MojoPrefillGQA output error trans (#237)
jessicagao01
authored
1a7d146

Commits on Apr 17, 2026

fix: unify all index-type inputs to int32 except embedding and sampling. (#234)
wwens7
authored
dcd0ffb
ci: update attn tests with padding logic (#240)
yfchen-byted
authored
6f66730

Commits on Apr 19, 2026

Commits on Apr 22, 2026

添加paged_attention_decode triton kernel
Xing Hongbo
committed
6541d14
添加paged_attention_decode triton kernel
Xing Hongbo
committed
33af5e1
tile_k >= block_size时简化代码
Xing Hongbo
committed
5b1126d
tile_k == block_size时简化代码
Xing Hongbo
committed
b557eac
删除测试代码
Xing Hongbo
committed
ee6dcf7
删除测试代码
Xing Hongbo
committed
e6bdbf9
删除测试代码
Xing Hongbo
committed
ad2f44a
修改v_stride错误
Xing Hongbo
committed
2b1888f
修复block_size=1024时nram使用超限
Xing Hongbo
committed
bd6b3d8
变量命名规范
Xing Hongbo
committed
a2264e8
简化load block_table代码
Xing Hongbo
committed
92dd15e
简化load block_tables代码
Xing Hongbo
committed
205f54c
修复block_size>tile load block错误
Xing Hongbo
committed
3141896
调整kernel性能相关参数
Xing Hongbo
committed
b1cdfd2
修复load block table错误
Xing Hongbo
committed
596e612
feat: adapt oe. (#247)
wwens7
authored
9888bbd
删除性能自动调优相关参数
Xing Hongbo
committed
3fd4344

Commits on Apr 23, 2026

Expose backend registry accessors (#254)

yskhhh
and
byteyushikang
authored
717f90f
refactor: quantize core impl. (#251)
wwens7
authored
d894401

Commits on Apr 25, 2026

fix: revert some apis. (#260)
wwens7
authored
8507e9f
feat: add mlu fa paged prefill triton backend (#248)
Camzhanbin
authored
a4c7777
Merge branch 'master' into paged_attention_decode
yfchen-byted
authored
e627e40

Commits on Apr 27, 2026

add fa_paged_decode.py
Xing Hongbo
committed
2fb3f05
fix oe weight loading (#261)
K13Byte
authored
356362a
feat: store kv && prefill attn support padding. (#259)

wwens7
and
fei.han
authored
c566cc5
feat: support swiglu limit & attn sink. (#262)
wwens7
authored
ebe1143
Merge pull request #255 from hongbo-X/paged_attention_decode
Neuromancer42
authored
6fd307f
[ilu/ixformer] Add paged attention ops (squash) (#246)

madengfei
and
yao.guo
authored
a4fb822
Merge pull request #257 from XPU-Forces/rope
aaddai
authored
574dd82

Commits on Apr 28, 2026

[ilu/ttx] Optimize ttx ops: PA(prefill/decode)、SWA(prefill/decode)、GroupGemm. (#238)

madengfei
and
yao.guo
authored
25d099f
feat: add _op_post_init() func for OverEncodingNGram to apply meta_model init (#264)

lipengfei1409
and
chenyifan.42
authored
9e5e173
misc: refine input args and outputs of Quant & Dequant -like kernels. (#268)

wwens7
and
chenyifan.42
authored
a71ae86

Commits on Apr 29, 2026

feat: refine npu ttx skills. (#265)
wwens7
authored
7062628

Commits on Apr 30, 2026

misc: unify semantics of all qkvlen. (#267)
wwens7
authored
39ba5bc

Commits on May 6, 2026

Commits on May 7, 2026

fix: gemm related. (#275)
wwens7
authored
67cd273
feat: support quant attn (#271)
Neuromancer42
authored
c220acf

Commits on May 8, 2026

[ilu/ixformer] Add PA(prefill/decode) cudagraph implement， SWA(prefill/decode) both eager and cudagraph supported. (#269)

authored
573c644
feat: support ilu ixformer quant moe & group gemm. (#282)

Mayyyybe
and
王晓梅
authored
271eaed

Commits on May 11, 2026

feat: add mlu kv_cache triton backend (#250)
aaddai
authored
4588f0e
[ilu/ixformer] Add w8a8 moe support. (#286)

Mayyyybe
and
xiaomei.wang
authored
ce38c02
feat: add vision 2D RoPE ref, npu ttx impl and tests. (#283)

Xu-pan
and
wwens7
authored
ed36ab8

Commits on May 12, 2026

Fix ILU KV cache padded head dim masking

byteyushikang
authored and
madengfei
committed
6018353
[ilu/ttx] Improved test accurary speed by use the selecetd profile, and optim the SDPA kernel. (#290)

authored
cc461dd

Commits on May 13, 2026

add mrope ops 0506 (#279)

qiudepei
and
wangyuwei22
authored
7e3762f
Refactor paged KV store to chunk metadata (#287)
K13Byte
authored
937e8c7
[ilu/ttx] add ilu triton OverEncodingNGram and OverEncoding (#293)

authored
162a89a

Commits on May 14, 2026

feat: Ilu/ixformer support quant, rope, kv op. (#299)

authored
b9e8509
[ilu/ixformer] Support ixformer quant paged prefill via F16 flash att… (#300)

authored
98b46b8
refactor: oe input args name. (#291)
wwens7
authored
5353e99
[ilu/ttx] Add CUDAGraph implements for PA and SWA. (#284)

madengfei
and
yao.guo
authored
7a4ae23
[ilu/ttx] Add ILU dynamic quant kernel for TTX MoE quantization. (#295)

authored
81f62a8
misc: refine readme. (#280)
wwens7
authored
3885083

Commits on May 15, 2026

[ilu/ttx] Add iluvatar group_rmsnorm triton impl. (#289)

authored
bec95db
[ilu/ttx] add paged decode quant GQA kernel implementation.

yao.guo
authored and
madengfei
committed
87ecab9
[ilu/ttx] Optimize paged decode quant GQA kernel with online softmax.

AbeFei
authored and
madengfei
committed
0376dee

Commits on May 18, 2026

Ilu/ttx quant to main pr (#304)

authored
2635522
[ilu/ixformer] add over encoding (#305)

authored
3591a3a

Commits on May 19, 2026

Refine core API experimental operators (#308)
zhangjihang-BD
authored
80fcaa1
misc: add some test cases. (#310)

wwens7
and
xiaomei.wang
authored
7b34652
[ilu/ixformer] support bf16 moe. (#314)

Mayyyybe
and
xiaomei.wang
authored
759898b
feat: quant_moe support differenct quant config for up_proj and down_proj (#312)

authored
2f3f69e

Commits on May 20, 2026

[ilu/ixformer] ixformer-PrefillGQAWithKVDequant/SWA dequant part use ixformer kernel. paged_kv_cache_dequant_varlen/window api support b/f16 k/v scale. (#311)

Mayyyybe
and
xudongzhao1006
authored
3625129
[ilu/ttx] opt n gram prefill and decode impl

sky-fun
authored and
madengfei
committed
c453d64

Commits on May 21, 2026

optimize paged decode swa redundancy mask and boundary_check

sky-fun
authored and
madengfei
committed
2b6193b
feat: ixformer support 2d scale for static quant. (#316)
wwens7
authored
38aa578

Commits on May 22, 2026

[ilu/ixformer] support moe using gdr (#320)

Mayyyybe
and
xiaomei.wang
authored
1294aca

Commits on May 25, 2026

[ilu/ttx] Add iluvatar moe and quant_moe (#319)

songliucsu
and
song.liu
authored
0947308
ci: add claude code automated PR review workflow (#323)

zhangjihang-BD
and
claude
authored
f97a481

Commits on May 26, 2026

Ilu/cache int8 (#324)

authored
37cb522

Commits on May 27, 2026

[ilu/ixformer] add IxformerPagedDecodeGQAWithKVDequant/IxformerPagedD… (#326)

authored
55a5c44

Commits on May 28, 2026

fix: [ttx/mlu][swa] move get_aux_mask on device (#334)
chuang-tang
authored
86e5f1f
Ilu:sage attn (#331)
lipengfei1409
authored
3cf35fe

Commits on May 29, 2026

refactor: align naming convention, change max_xxx_lens to max_xxx_len (#335)

authored
bc40f28
[ilu/ixformer] use bf16 * bf16 gemm in MojoMoEGating. (#333)

Mayyyybe
and
xiaomei.wang
authored
6b36db7
feat: add attention_gate core impl. (#327)

wwens7
and
xiaomei.wang
authored
ec66676
Enable torch deterministic config for mojo deterministic mode (#330)

byteyushikang
and
wwens7
authored
6ea4592

Commits on Jun 1, 2026

Shmem dev (#339)

zhangjihang-BD
and
claude
authored
5d6e97d
[ilu/ttx] optimize iluvatar triton group_norm impl (#325)

songliucsu
and
song.liu
authored
26806b2
[ilu/ttx] optimize paged_prefill_gqa: remove aux_mask, inline causal mask, simplify kernel structure (#321)
AbeFei
authored
fe82961

Commits on Jun 5, 2026

add deepep and test, modify test_moe_quant and test_attention_quant
qiushi13
committed
e69da79
update test_deepep.py
qiushi13
committed
9a52847

Commits on Jun 11, 2026

fix: 入参增加pg, deepep.py从core移到experimental
qiushi13
committed
8121473
feat: add MojoRotaryEmbedding
MikuGhoul
committed
659292c
feat: add MojoGatherRopeStore
HJzhang-sjtu
committed
de0222b
feat: add MojoPagedAttentionStoreKvCache
HJzhang-sjtu
committed
bc5ed60
feat: add MojoPagedCacheDequant
HJzhang-sjtu
committed
d5b789e

Commits on Jun 12, 2026

feat: add MojoFusedAttnGateConcat
HJzhang-sjtu
committed
7abb604
fix: add assert and fix c8 quant in golden
HJzhang-sjtu
committed
af1b5fe
Merge pull request #355 from XPU-Forces/hjzhang/attn
HJzhang-sjtu
authored
fa4763c
feat: add fused ag scale quant and qk rmsnorm
jessicagao01
committed
53e83cf

Commits on Jun 15, 2026