[PR #233] Optimize NVFP4 and MXFP8 quantization in Triton backend #243

New issue

Open

opened 2026-10-26 22:21:56 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:21:56 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/233
Author: @nukuihayu
Created: 10/8/2026
Status: 🔄 Open

Base: main ← Head: optimize-triton-nvfp4-mxfp8

📝 Commits (1)

5752a5d Optimize NVFP4 and MXFP8 quantization in Triton backend

📊 Changes

1 file changed (+155 additions, -192 deletions)

View changed files

📝 comfy_kitchen/backends/triton/quantization.py (+155 -192)

📄 Description

Changes

Process 32 quantization blocks in parallel per program instead of looping over 1–4 blocks, increasing useful work per program.
Replace NVFP4's one-hot reductions with reshape/split for packing, preserving nibble order and signed-zero behavior.
Remove MXFP8's unintended FP64 scale calculation. Encode E8M0 directly with CUDA-equivalent saturating upward rounding, using bit operations supported on SM89 and newer.
Use FP32 division and handle the minimum E8M0 scale explicitly. Remove the tiny-scale zeroing threshold to retain small finite inputs.

Performance

RTX 5090, BF16 input:

Operation	Shape	Before	After
NVFP4	256 × 1024	8.31 µs	1.89 µs
NVFP4	4096 × 4096	393.33 µs	17.72 µs
MXFP8	256 × 1024	47.95 µs	1.96 µs
MXFP8	4096 × 4096	2760.63 µs	12.98 µs

About 22× faster for NVFP4 and 213× for MXFP8 at 4096², with no material small-input regression across 12 tested operation/shape combinations.

CUDA Graph GPU timings include scale clearing and quantization, excluding Python and compilation. Median of seven alternating rounds, 20 calls per replay. PyTorch 2.13.0+cu130, Triton 3.7.1; baseline be003b7c.

Testing

103 tests passed; six CUDA Graph replays passed; Compute Sanitizer reported zero errors.
E8M0 conversion matches CUDA for all 2³¹ nonnegative FP32 bit patterns. This checks the conversion step, not whole-kernel equivalence.
Runtime tested on SM120; offline compilation passed for NVFP4 on SM100 and MXFP8 on SM89/SM100.
Ruff and git diff --check passed.

_{🔄 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/233 **Author:** [@nukuihayu](https://github.com/nukuihayu) **Created:** 10/8/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `optimize-triton-nvfp4-mxfp8` --- ### 📝 Commits (1) - [`5752a5d`](https://github.com/Comfy-Org/comfy-kitchen/commit/5752a5d497d461f40638d072e5fee591735490e1) Optimize NVFP4 and MXFP8 quantization in Triton backend ### 📊 Changes **1 file changed** (+155 additions, -192 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/triton/quantization.py` (+155 -192) </details> ### 📄 Description ## Changes - Process 32 quantization blocks in parallel per program instead of looping over 1–4 blocks, increasing useful work per program. - Replace NVFP4's one-hot reductions with reshape/split for packing, preserving nibble order and signed-zero behavior. - Remove MXFP8's unintended FP64 scale calculation. Encode E8M0 directly with CUDA-equivalent saturating upward rounding, using bit operations supported on SM89 and newer. - Use FP32 division and handle the minimum E8M0 scale explicitly. Remove the tiny-scale zeroing threshold to retain small finite inputs. ## Performance RTX 5090, BF16 input: | Operation | Shape | Before | After | | --- | --- | ---: | ---: | | NVFP4 | 256 × 1024 | 8.31 µs | 1.89 µs | | NVFP4 | 4096 × 4096 | 393.33 µs | 17.72 µs | | MXFP8 | 256 × 1024 | 47.95 µs | 1.96 µs | | MXFP8 | 4096 × 4096 | 2760.63 µs | 12.98 µs | About 22× faster for NVFP4 and 213× for MXFP8 at 4096², with no material small-input regression across 12 tested operation/shape combinations. CUDA Graph GPU timings include scale clearing and quantization, excluding Python and compilation. Median of seven alternating rounds, 20 calls per replay. PyTorch 2.13.0+cu130, Triton 3.7.1; baseline `be003b7c`. ## Testing - 103 tests passed; six CUDA Graph replays passed; Compute Sanitizer reported zero errors. - E8M0 conversion matches CUDA for all 2³¹ nonnegative FP32 bit patterns. This checks the conversion step, not whole-kernel equivalence. - Runtime tested on SM120; offline compilation passed for NVFP4 on SM100 and MXFP8 on SM89/SM100. - Ruff and `git diff --check` passed. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 22:21:56 +00:00

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#243

No description provided.

Rows
Columns