[PR #348] feat: support mxfp4 × mxfp4 gemm on sm100 #382

Open
opened 2026-10-26 21:17:05 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/348
Author: @z52527
Created: 6/1/2026
Status: 🔄 Open

Base: nv_dev ← Head: fea-fp4-synced


📝 Commits (10+)

  • 7d27b8a FP4: port MXF4 kernel + wrappers onto main's refactored APIs
  • 7c2d685 FP4: progress port towards runtime — IMA fixed, smem cap, multicast clamp
  • 6539975 FP4: layout override to fea-fp4 tested config + test int8 packing
  • 24eb9c2 FP4: tests/test_fp4.py — use FP32 output (kernel's only working epilogue path)
  • ae98142 FP4: tests/test_fp4.py — full suite passes on synced branch (10/10)
  • dfdc5ae FP4: port fea-fp4's full get_best_fp4_config heuristic — perf parity
  • 31aef98 FP4: rewrite tests/test_fp4.py to mirror test_fp8_fp4.py style
  • 8d841f6 FP4: align .cuh comments with sm100_fp8_fp4_gemm_1d1d.cuh style
  • ae4e77d FP4: address codex review (PDL sync, strict FP32 D, prune unused fields)
  • 959e308 FP4: remove .item() host sync in m-grouped contiguous swap_ab gating

📊 Changes

4 files changed (+1386 additions, -9 deletions)

View changed files

📝 csrc/apis/gemm.hpp (+55 -9)
➕ csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp (+462 -0)
➕ deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh (+735 -0)
➕ tests/test_fp4.py (+134 -0)

📄 Description

Adds an FP4×FP4 (MXF4) GEMM path on SM100/B200 using SM100_MMA_MXF4_SS directly, separate from main's unified MXF8F6F4 kernel. Dispatched from the existing fp8_fp4 APIs (dense / m-grouped contiguous / m-grouped masked) when both operands are FP4. Tested with tests/test_fp4.py on B200 (28/28 pass). Bench: median ~1.07× vs cublasTest 13.3 direct, 10/13 shapes ≥ cuBLAS.

Dense GEMM perf (B200, CUDA 12.9, vs cublasTest 13.3 direct binary)

M×N×K DG TFLOPS cBLT TFLOPS DG/cBLT
128×2048×7168 328 338 0.97×
128×24576×1536 980 745 1.32×
512×2048×7168 1178 1422 0.83×
512×7168×2048 1439 1190 1.21×
1024×7168×2048 2056 1399 1.47×
4096×4096×7168 5332 3939 1.35×
8192×8192×8192 5702 5239 1.09×

TODO

  • BF16 D epilogue — kernel currently hardcodes fp32 output
  • Contract-rejection tests — wrapper has DG_HOST_ASSERTs for aliasing / K-major / k % 8 but no test coverage (repo-wide gap, not FP4-only)
  • Sparse-MoE swap_ab gating uses m / num_groups as a host estimate accurate for dense, but misses swap_ab on -1-padded sparse layouts (padding inflates the estimate above block_m). Avoids a device sync on m_indices; revisit if sparse MoE perf is a bottleneck.
  • ~5–15% perf gap vs cuBLAS on mid-M shapes (M ∈ {256, 512}, N ∈ {2k, 4k}, K=7168)

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/348 **Author:** [@z52527](https://github.com/z52527) **Created:** 6/1/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `fea-fp4-synced` --- ### 📝 Commits (10+) - [`7d27b8a`](https://github.com/deepseek-ai/DeepGEMM/commit/7d27b8a071b5ffd4ba6c3763a27b9ce25ac9d0cf) FP4: port MXF4 kernel + wrappers onto main's refactored APIs - [`7c2d685`](https://github.com/deepseek-ai/DeepGEMM/commit/7c2d685dbd91eb80992c02456dc5520b6120cba4) FP4: progress port towards runtime — IMA fixed, smem cap, multicast clamp - [`6539975`](https://github.com/deepseek-ai/DeepGEMM/commit/65399759a54837a2db3112f986c2fbbdb707af74) FP4: layout override to fea-fp4 tested config + test int8 packing - [`24eb9c2`](https://github.com/deepseek-ai/DeepGEMM/commit/24eb9c280abf4897b20dc2bd1946d84c0b4d788b) FP4: tests/test_fp4.py — use FP32 output (kernel's only working epilogue path) - [`ae98142`](https://github.com/deepseek-ai/DeepGEMM/commit/ae98142d4d82608343768e8cc9350be1f79b6bcb) FP4: tests/test_fp4.py — full suite passes on synced branch (10/10) - [`dfdc5ae`](https://github.com/deepseek-ai/DeepGEMM/commit/dfdc5aeb3b3f3f2071690410332828512e95c888) FP4: port fea-fp4's full get_best_fp4_config heuristic — perf parity - [`31aef98`](https://github.com/deepseek-ai/DeepGEMM/commit/31aef98fe7d293455ea807888f08e3a2e37f19cd) FP4: rewrite tests/test_fp4.py to mirror test_fp8_fp4.py style - [`8d841f6`](https://github.com/deepseek-ai/DeepGEMM/commit/8d841f664fa83a6802136653adb90afb13caea18) FP4: align .cuh comments with sm100_fp8_fp4_gemm_1d1d.cuh style - [`ae4e77d`](https://github.com/deepseek-ai/DeepGEMM/commit/ae4e77d8c7316101241bb64f7b219d278e5ca22a) FP4: address codex review (PDL sync, strict FP32 D, prune unused fields) - [`959e308`](https://github.com/deepseek-ai/DeepGEMM/commit/959e308af8078e59d8eabccd7142915d55e6ebba) FP4: remove .item() host sync in m-grouped contiguous swap_ab gating ### 📊 Changes **4 files changed** (+1386 additions, -9 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/gemm.hpp` (+55 -9) ➕ `csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp` (+462 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh` (+735 -0) ➕ `tests/test_fp4.py` (+134 -0) </details> ### 📄 Description Adds an FP4×FP4 (MXF4) GEMM path on SM100/B200 using `SM100_MMA_MXF4_SS` directly, separate from main's unified MXF8F6F4 kernel. Dispatched from the existing fp8_fp4 APIs (dense / m-grouped contiguous / m-grouped masked) when both operands are FP4. Tested with `tests/test_fp4.py` on B200 (28/28 pass). Bench: median ~1.07× vs `cublasTest 13.3` direct, 10/13 shapes ≥ cuBLAS. ## Dense GEMM perf (B200, CUDA 12.9, vs cublasTest 13.3 direct binary) | M×N×K | DG TFLOPS | cBLT TFLOPS | DG/cBLT | |------------------|----------:|------------:|--------:| | 128×2048×7168 | 328 | 338 | 0.97× | | 128×24576×1536 | 980 | 745 | 1.32× | | 512×2048×7168 | 1178 | 1422 | 0.83× | | 512×7168×2048 | 1439 | 1190 | 1.21× | | 1024×7168×2048 | 2056 | 1399 | 1.47× | | 4096×4096×7168 | 5332 | 3939 | 1.35× | | 8192×8192×8192 | 5702 | 5239 | 1.09× | ## TODO - BF16 D epilogue — kernel currently hardcodes fp32 output - Contract-rejection tests — wrapper has `DG_HOST_ASSERT`s for aliasing / K-major / `k % 8` but no test coverage (repo-wide gap, not FP4-only) - Sparse-MoE swap_ab gating uses `m / num_groups` as a host estimate accurate for dense, but misses swap_ab on `-1`-padded sparse layouts (padding inflates the estimate above `block_m`). Avoids a device sync on `m_indices`; revisit if sparse MoE perf is a bottleneck. - ~5–15% perf gap vs cuBLAS on mid-M shapes (M ∈ {256, 512}, N ∈ {2k, 4k}, K=7168) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#382
No description provided.