[PR #26] [MERGED] Add MXFP8 Support and Consolidated stable scaled mm API #87

New issue

Closed

opened 2026-10-26 22:20:15 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:15 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/26
Author: @contentis
Created: 1/23/2026
Status: ✅ Merged
Merged: 1/24/2026
Merged by: @comfyanonymous

Base: main ← Head: scaled_mm_v2

📝 Commits (10+)

85d3023 Add torch functional scaled_mm
3492c5b dlopen cublas making it optional at runtime
2b7ffa1 add MXFP8 and consolidate v2 MM api
186a862 add MXFP8 and consolidate v2 MM api
9918bdb update readme
182ea13 fix linear dispatch issue for MX
afc2386 Add turing+ptx, housekeeping
d8fc5f9 some perf tuning
9d610c2 add arm64 runner
39e1b83 Merge branch 'arm64_runner' into scaled_mm_v2

📊 Changes

21 files changed (+2051 additions, -31 deletions)

View changed files

📝 .github/workflows/build-wheels.yml (+86 -1)
📝 README.md (+4 -0)
📝 comfy_kitchen/__init__.py (+73 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+46 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+60 -0)
➕ comfy_kitchen/backends/cuda/ops/quantize_mxfp8.cu (+211 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+48 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+290 -13)
📝 comfy_kitchen/backends/triton/__init__.py (+11 -0)
📝 comfy_kitchen/backends/triton/quantization.py (+181 -0)
📝 comfy_kitchen/float_utils.py (+28 -0)
➕ comfy_kitchen/scaled_mm_v2.py (+99 -0)
📝 comfy_kitchen/tensor/__init__.py (+3 -0)
📝 comfy_kitchen/tensor/fp8.py (+3 -5)
➕ comfy_kitchen/tensor/mxfp8.py (+236 -0)
📝 pyproject.toml (+1 -0)
➕ samples/mxfp8_model_patcher.py (+370 -0)
📝 samples/nvfp4_linear.py (+4 -6)
📝 tests/test_qdq.py (+62 -0)

...and 1 more files

📄 Description

Updated scaled mm API

With torch 2.10 scaled_mm was introduced as stable API
Added wrapper to pipe through the stble API if the torch version is sufficient, fallback to experimental _scaled_mm otherwise (current behaviour)

MXFP8

Add MXFP8 (Microscaling FP8) block quantization support with E8M0 (power-of-2) scales and 32-element blocks, targeting SM >= 10.0 (Blackwell)
Introduce TensorCoreMXFP8Layout with full operator dispatch for mm, addmm, and linear

_{🔄 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/Comfy-Org/comfy-kitchen/pull/26 **Author:** [@contentis](https://github.com/contentis) **Created:** 1/23/2026 **Status:** ✅ Merged **Merged:** 1/24/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `scaled_mm_v2` --- ### 📝 Commits (10+) - [`85d3023`](https://github.com/Comfy-Org/comfy-kitchen/commit/85d3023e54e7e784585d95cee15b60c525ae1351) Add torch functional scaled_mm - [`3492c5b`](https://github.com/Comfy-Org/comfy-kitchen/commit/3492c5b98023e229fa03ad366244ceb2d2ef0fa0) dlopen cublas making it optional at runtime - [`2b7ffa1`](https://github.com/Comfy-Org/comfy-kitchen/commit/2b7ffa105df0e4a67f76eafb4f87a5fce2304663) add MXFP8 and consolidate v2 MM api - [`186a862`](https://github.com/Comfy-Org/comfy-kitchen/commit/186a8622bd49915f6a7f40a1d5febf9fe34d6d09) add MXFP8 and consolidate v2 MM api - [`9918bdb`](https://github.com/Comfy-Org/comfy-kitchen/commit/9918bdb0196c312b94f8625018a3db77173fe78a) update readme - [`182ea13`](https://github.com/Comfy-Org/comfy-kitchen/commit/182ea139b880be13ccc0d01ecdb87e37080ad48a) fix linear dispatch issue for MX - [`afc2386`](https://github.com/Comfy-Org/comfy-kitchen/commit/afc2386803d694df58f675ac0cf10a2d07c16b5a) Add turing+ptx, housekeeping - [`d8fc5f9`](https://github.com/Comfy-Org/comfy-kitchen/commit/d8fc5f9dc0b243a4bf5d3659d9f1410ebbb24ba0) some perf tuning - [`9d610c2`](https://github.com/Comfy-Org/comfy-kitchen/commit/9d610c2cebd6c7adc3dde521aaacec71d97f92e1) add arm64 runner - [`39e1b83`](https://github.com/Comfy-Org/comfy-kitchen/commit/39e1b830ff878313f9ae279e53d29b0cf0d67b7d) Merge branch 'arm64_runner' into scaled_mm_v2 ### 📊 Changes **21 files changed** (+2051 additions, -31 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/build-wheels.yml` (+86 -1) 📝 `README.md` (+4 -0) 📝 `comfy_kitchen/__init__.py` (+73 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+46 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+60 -0) ➕ `comfy_kitchen/backends/cuda/ops/quantize_mxfp8.cu` (+211 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+48 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+290 -13) 📝 `comfy_kitchen/backends/triton/__init__.py` (+11 -0) 📝 `comfy_kitchen/backends/triton/quantization.py` (+181 -0) 📝 `comfy_kitchen/float_utils.py` (+28 -0) ➕ `comfy_kitchen/scaled_mm_v2.py` (+99 -0) 📝 `comfy_kitchen/tensor/__init__.py` (+3 -0) 📝 `comfy_kitchen/tensor/fp8.py` (+3 -5) ➕ `comfy_kitchen/tensor/mxfp8.py` (+236 -0) 📝 `pyproject.toml` (+1 -0) ➕ `samples/mxfp8_model_patcher.py` (+370 -0) 📝 `samples/nvfp4_linear.py` (+4 -6) 📝 `tests/test_qdq.py` (+62 -0) _...and 1 more files_ </details> ### 📄 Description Updated scaled mm API - With torch 2.10 scaled_mm was introduced as stable [API](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_mm.html) - Added wrapper to pipe through the stble API if the torch version is sufficient, fallback to experimental _scaled_mm otherwise (current behaviour) **MXFP8** - Add MXFP8 (Microscaling FP8) block quantization support with E8M0 (power-of-2) scales and 32-element blocks, targeting SM >= 10.0 (Blackwell) - Introduce TensorCoreMXFP8Layout with full operator dispatch for mm, addmm, and linear --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:15 +00:00

closed this issue
added the
pull-request
label

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

Comfy-Org/comfy-kitchen#87

No description provided.

Rows
Columns