8d02eb38fa
fix: transpose SF to MN-major layout before TMA stride checks
...
transform_sf_into_required_layout expects MN-major input (stride(-2)=1).
Our packed int32 SF is K-major (stride(-1)=1). Transpose the last two
dims, make contiguous, then transpose back so data is in MN-major order.
2026-05-11 07:32:10 +00:00
7154500f22
fix: reshape SF to 2D before transform_sf_into_required_layout
...
The C++ check_sf_layout stride assertion fails on 3D (experts, mn, K//64)
tensors. Reshape to 2D (experts*mn, K//64) before calling the transform
function, matching the expected stride layout.
2026-05-11 07:30:54 +00:00
388fd8dcfd
fix: pack UE4M3 into int32 before transform_sf_into_required_layout
...
The C++ transform function expects int32 (for kInt type) with 4 UE4M3
bytes packed per int32. We pack first, then transform for TMA alignment
and UTCCP transpose with recipe (1, 16).
2026-05-11 07:05:11 +00:00
acae75e109
fix: use transform_sf_into_required_layout for proper TMA-aligned SF
...
Instead of custom _pack_nvfp4_sf_for_utccp, use DeepGEMM's C++
transform_sf_into_required_layout with recipe (1, 1, 16) for NVFP4.
This handles TMA alignment and UTCCP transpose correctly.
2026-05-11 06:54:34 +00:00
5cb4fcaef3
fix: cast uint8 weights to int8 (kPackedFP4) for DeepGEMM compatibility
2026-05-11 06:36:32 +00:00
bbf9a5f46a
feat: fold weight_scale_2 into block scales in NVFP4 transform
...
- transform_nvfp4_weights_for_mega_moe now accepts weight_scale_2
- Folds global scale into block scales: UE4M3 * FP32 -> UE4M3
- Dequantize to f32, multiply by global scale, clamp [0,448], re-quantize
- This is needed because the kernel only applies one level of block scaling
2026-05-11 05:42:16 +00:00
36b439ee26
feat: NVFP4 mega MoE kernel (scale_vec::4X, UE4M3 block scales)
...
- New CUDA kernel: sm100_fp8_nvfp4_mega_moe_impl
- kGranK=16 (NVFP4 group_size=16, vs MXFP4's 32)
- kind::mxf4nvf4.block_scale.scale_vec::4X PTX instruction
- float_ue4m3_t scale factor type in instruction descriptor
- SF layout: scale_vec::4X (4 TMEM sub-columns per UMMA atom)
- UTCCP column stride: i*8 (vs MXFP4's i*4) for 4X layout
- L1 epilogue: UE4M3 activation scales (float→cutlass::float_e4m3_t)
- SF loading: kNumSFUint32 = kHidden/64 (4 UE4M3 per int32)
- New PTX wrappers: SM100_MMA_MXF4NVF4_2x1SM_SS, SM100_MMA_MXF4NVF4_SS
- Python API:
- fp8_nvfp4_mega_moe() with recipe=(1,1,16)
- transform_nvfp4_weights_for_mega_moe() for UE4M3→int32 UTCCP packing
- _pack_nvfp4_sf_for_utccp() helper
- C++ bindings:
- mega_nvfp4.hpp with NVFP4-specific SymmBuffer (SF stride K/16)
- JIT kernel header with kGranK=16 TMA descriptors
- Registered in python_api.cpp
NOTE: Both SFA and SFB must use UE4M3 (scale_format_ is 1-bit, shared).
The L1 epilogue converts float→UE4M3 for activation scales.
2026-05-11 05:41:08 +00:00
Zhean Xu
891d57b4db
Add various optimizations and Mega MoE benchmarks ( #316 )
...
* Merge with private repo
* Add Mega MoE Benchmark
* Minor fix
* Update
---------
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com >
2026-04-24 18:41:37 +08:00
Chenggang Zhao
7f2a703ed5
[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes ( #304 )
...
* Merge with private repo
* Update README
* Update README
* Update README
* Add PyTorch requirements
* Fix sync scopes for MQA logits (#256 )
* Update README
2026-04-17 09:45:14 +08:00
Ray Wang
d30fc36c8f
Fix sync issue of TMEM alloc/dealloc ( #292 )
2026-03-22 16:41:28 +08:00
Xin Qiu
35c4bc8771
fix: k_grouped_fp8_gemm_nt_contiguous crashes with n = 768 on H100 ( #238 )
2026-02-25 10:13:54 +08:00
Ray Wang
477618cd51
Fix a sync issue in SM100 MQA logits ( #285 )
2026-02-03 17:29:49 +08:00
Zhean Xu
0f5f266202
Multiple updates and refactorings ( #280 )
2026-01-16 17:06:52 +08:00
Ray Wang
38f8ef73a4
Multiple updates and refactorings ( #231 )
2025-11-21 17:49:47 +08:00
Zhean Xu
bb4424aad4
Fix sum_k * shape_m overflow
2025-11-19 11:51:36 +08:00
Ray Wang
ec5e9ed0b8
Fix SM90 MQA logits ( #229 )
2025-11-19 10:50:36 +08:00
Ray Wang
2f9d87877e
Use larger MMA shape ( #227 )
2025-11-14 11:38:15 +08:00
oliver könig
9f196058ae
chore: Build and store bdist wheels ( #181 )
...
* build: Minor tweeks for wheel build
Signed-off-by: oliver könig <okoenig@nvidia.com >
* ci: Workflows for wheel build
Signed-off-by: oliver könig <okoenig@nvidia.com >
* fix
Signed-off-by: oliver könig <okoenig@nvidia.com >
* fix
Signed-off-by: oliver könig <okoenig@nvidia.com >
* build: Add CachedWheel
Signed-off-by: oliver könig <okoenig@nvidia.com >
* add version to init
Signed-off-by: oliver könig <okoenig@nvidia.com >
* revert
Signed-off-by: oliver könig <okoenig@nvidia.com >
* revert
Signed-off-by: oliver könig <okoenig@nvidia.com >
* revert
Signed-off-by: oliver könig <okoenig@nvidia.com >
* v2
Signed-off-by: oliver könig <okoenig@nvidia.com >
* update
Signed-off-by: oliver könig <okoenig@nvidia.com >
* test
Signed-off-by: oliver könig <okoenig@nvidia.com >
* from packaging.version import parse
Signed-off-by: oliver könig <okoenig@nvidia.com >
* local version
Signed-off-by: oliver könig <okoenig@nvidia.com >
* remove file
Signed-off-by: oliver könig <okoenig@nvidia.com >
* revert
Signed-off-by: oliver könig <okoenig@nvidia.com >
* Updates and lint
* revert missing cudaextension args
Signed-off-by: oliver könig <okoenig@nvidia.com >
* Add timeout
* fix DG settings
Signed-off-by: oliver könig <okoenig@nvidia.com >
* DG_USE_LOCAL_VERSION
Signed-off-by: oliver könig <okoenig@nvidia.com >
* Update version
* Detect local changes
* Minor fix
* Revert CUTLASS
* Unify options
---------
Signed-off-by: oliver könig <okoenig@nvidia.com >
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com >
2025-10-10 18:23:40 +08:00
Chenggang Zhao
c1bf4cae4b
Fix version
2025-10-01 20:31:27 +08:00
Chenggang Zhao
07b82fb8cd
Fix old CUDA compatibility
2025-10-01 20:29:15 +08:00
Simon Mo
59f2c07cf2
Add SM100 kernels ( #201 )
...
Signed-off-by: simon-mo <simon.mo@hey.com >
2025-09-29 17:07:28 +08:00
Chenggang Zhao
80ceeb2c76
Add SM90 kernels ( #200 )
2025-09-29 17:00:23 +08:00
Ray Wang
3f71de7aa9
Make various updates and fixes ( #198 )
2025-09-25 16:19:07 +08:00
zhonghui-J
2da871e304
Fix grouped gemms performance issue. ( #168 )
2025-08-22 17:35:43 +08:00
Chenggang Zhao
e38c2e3103
Remove comments
2025-08-22 17:32:04 +08:00
Chenggang Zhao
f20256fd50
Compatible with CUDA 13
2025-08-22 17:30:47 +08:00
xiweny
affdb1cd90
Add sm_100f support and make nvcc 13 happy ( #157 )
...
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com >
2025-08-22 17:19:32 +08:00
Ray Wang
f85ec649d7
Make various updates and fixes: ( #164 )
...
- Add BF16 support for SM90 and SM100
- Refactor Python APIs
- Other fixes and code refactoring
2025-08-15 18:32:35 +08:00
Ray Wang
d9c363f86f
Make various updates and fixes:
...
- Add support for legacy CUDA versions; now compatible with CUDA 12.3 and newer
- Add support for NVRTC compilation
- Other fixes and code refactoring
2025-08-02 19:52:22 -07:00
yukuai26
aff9da0aba
Fix SM90 GEMM ( #149 )
...
* Fix sm90 GEMM
* Fix typo
---------
Co-authored-by: Kuai Yu <yukuai@deepseek.com >
2025-08-01 10:36:49 +08:00
Ray Wang
9da4a23561
Add more GPU architectures support ( #112 )
...
* Add more GPU architectures support
* Update layout.py
* Optimize performance, Add SM90 support, Add 1D2D SM100 support
* Add fmtlib submodule at commit 553ec11
---------
Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com >
2025-07-18 11:32:22 +08:00
Chenggang Zhao
03d0be3d2d
Simplify expression
2025-07-02 14:07:05 +08:00
fy1214
3fc6728dee
[add] fix smem_barrier size in wgrad way ( #122 )
2025-07-02 14:05:36 +08:00
yukuai
e82c4139da
Revert "Fixed the bug in get_swizzle_mode function related to elem_size setting. ( #115 )"
...
This reverts commit ac428e25e0 .
This PR causes wgrad to hang during testing. Revert it until we resolve the issue
2025-06-23 17:13:36 +08:00
TherLF
ac428e25e0
Fixed the bug in get_swizzle_mode function related to elem_size setting. ( #115 )
2025-06-23 09:37:10 +08:00
shixianc
0c88cd0139
Fix illegal memory address when skipping -1 m indices ( #113 )
...
Co-authored-by: Shixian Cui <shixian@amazon.com >
2025-06-16 10:44:31 +08:00
yukuai26
8dfa329827
Grouped GEMM skip useless computation for unaligned Ms ( #103 )
...
* Grouped GEMM skip useless computation for unaligned Ms
* Update readme.md
* small typo
* Rename variables
* Restore previous indent
* Format
* Refactor tests
* Add `SkipComputation` types
* Bug fixed
* Format
* Fix tests
* Add assertions
* Minor fix
---------
Co-authored-by: yukuai <yukuai@deepseek.com >
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com >
2025-05-27 13:43:38 +08:00
Chenggang Zhao
391755ada0
Fix JIT tests
2025-05-16 14:39:58 +08:00
Chenggang Zhao
78d8362e7a
Add a missing #pragma once
2025-05-15 18:10:05 +08:00
Chenggang Zhao
104a6ec109
Add __assertfail
2025-05-15 17:04:21 +08:00
Chenggang Zhao
3b412f458a
Unify kwargs usages
2025-05-15 16:53:52 +08:00
Chenggang Zhao
350989eef3
Unify ceil_divs
2025-05-15 16:48:32 +08:00
Chenggang Zhao
4373af2e82
Add DG_PRINT_CONFIGS
2025-05-15 16:36:40 +08:00
Chenggang Zhao
816b39053a
Refactor launch-related structures
2025-05-15 16:14:21 +08:00
Chenggang Zhao
e2d6a107ef
Cleanup some useless staffs
2025-05-14 15:46:45 +08:00
Zhean Xu
04278f6dee
Weight gradient kernels for dense and MoE models ( #95 )
...
* Init weight gradient kernels.
* Support unaligned n,k and gmem stride
* Update docs
* Several cleanups
* Remove restrictions on N
* Add stride(0) assertions
---------
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com >
2025-05-14 14:47:58 +08:00
Chenggang Zhao
8702f910e3
Fix 12.9 compatibility
2025-05-07 13:23:40 +08:00
Chenggang Zhao
085b4a1532
Add DG_PRINT_AUTOTUNE to README
2025-05-07 11:46:52 +08:00
Chenggang Zhao
daec8fd2fc
Fix pipeline stage edge cases
2025-05-07 11:40:34 +08:00
Gabriel Wu
bfe983c4c2
Refactor JIT compilation (+NVRTC support) ( #94 )
...
* [wip] refactor: compile to .cubin
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* refactor: compile to .cubin and add NVRTC option
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* fix: compiler version
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* feat: compat for old drivers
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* feat: save kernel name to file
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* feat: fix win compat
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* fix: windows compat
Signed-off-by: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com >
* feat: make API more general
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* feat: drop support for CUDA<12.3
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* doc: update README
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
* Some lints and refactor
* Refactor runtime
* Several fixes
* Refactor environment variables
* Code format
* Add a TODO
* Compatible with CUDA 12.3
* Fix indent
* Fix typing
* Drop support for Windows
* Add a TODO
---------
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com >
Signed-off-by: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com >
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com >
2025-05-07 11:38:14 +08:00