[PR #157] [CLOSED] Add stochastic FP8 rounding to Triton backend #183

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/157
Author: @tangzzycc
Created: 9/5/2026
Status: ❌ Closed

Base: main ← Head: feat/triton-stochastic-fp8

📝 Commits (2)

d2dd0cf Add Triton stochastic FP8 rounding
73ed759 Skip Triton stochastic FP8 tests below SM89

📊 Changes

5 files changed (+261 additions, -1 deletions)

View changed files

📝 README.md (+1 -1)
📝 comfy_kitchen/backends/triton/__init__.py (+13 -0)
📝 comfy_kitchen/backends/triton/quantization.py (+123 -0)
📝 tests/test_constraints.py (+8 -0)
📝 tests/test_qdq.py (+116 -0)

📄 Description

Summary

Add a Triton implementation of stochastic_rounding_fp8 for NVIDIA GPUs with native FP8 support (compute capability 8.9+). This avoids the slow, memory-heavy eager fallback reported in Comfy-Org/ComfyUI#15311 when the native CUDA backend is unavailable.

Changes

Match the CUDA path: FP16 input rounding followed by FP32 stochastic-rounding arithmetic.
Support FP32, FP16, and BF16 input with E4M3FN or E5M2 output.
Reuse the contiguous uint8 RNG buffer as output, avoiding an additional output allocation.
Use adaptive block sizes and 64-bit offsets for large tensors.
Register the operation for CUDA devices with compute capability >= 8.9 and add tests.

Correctness and compatibility

Exhaustive finite-value check: all 65,536 FP16 bit patterns x 256 RNG values, with zero Triton/CUDA mismatches for both FP8 formats.
Offline compilation: 96 combinations across SM89, SM90, SM100, SM120, all input/output types, and all block sizes.
Every compiled variant uses 4 warps, 0 bytes shared memory, and 0 bytes global scratch memory.

RTX 5090 performance

BF16 to E4M3FN, warmed kernels, median latency, rotating working set >= 256 MiB:

Elements	CUDA	Triton	Eager	Triton vs CUDA	Triton vs Eager
1,048,576	0.0136 ms	0.0266 ms	0.3538 ms	1.95x slower	13.3x faster
8,388,608	0.0329 ms	0.0263 ms	0.3696 ms	1.25x faster	14.1x faster
16,777,216	0.0628 ms	0.0415 ms	0.8942 ms	1.51x faster	21.5x faster

For 16,777,216 contiguous BF16 elements, both CUDA and Triton add 0 MiB of peak PyTorch-allocated memory because the output aliases the RNG buffer. Eager adds 240 MiB.

Validation

118 passed, 4 skipped
ruff: All checks passed
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/157 **Author:** [@tangzzycc](https://github.com/tangzzycc) **Created:** 9/5/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `feat/triton-stochastic-fp8` --- ### 📝 Commits (2) - [`d2dd0cf`](https://github.com/Comfy-Org/comfy-kitchen/commit/d2dd0cfe3a3b51cf9cfe2d68a33e53a7dcdd71d0) Add Triton stochastic FP8 rounding - [`73ed759`](https://github.com/Comfy-Org/comfy-kitchen/commit/73ed759ac02d07ed2a091d834e86f8e8aa8eb89d) Skip Triton stochastic FP8 tests below SM89 ### 📊 Changes **5 files changed** (+261 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -1) 📝 `comfy_kitchen/backends/triton/__init__.py` (+13 -0) 📝 `comfy_kitchen/backends/triton/quantization.py` (+123 -0) 📝 `tests/test_constraints.py` (+8 -0) 📝 `tests/test_qdq.py` (+116 -0) </details> ### 📄 Description ## Summary Add a Triton implementation of `stochastic_rounding_fp8` for NVIDIA GPUs with native FP8 support (compute capability 8.9+). This avoids the slow, memory-heavy eager fallback reported in Comfy-Org/ComfyUI#15311 when the native CUDA backend is unavailable. ## Changes - Match the CUDA path: FP16 input rounding followed by FP32 stochastic-rounding arithmetic. - Support FP32, FP16, and BF16 input with E4M3FN or E5M2 output. - Reuse the contiguous `uint8` RNG buffer as output, avoiding an additional output allocation. - Use adaptive block sizes and 64-bit offsets for large tensors. - Register the operation for CUDA devices with compute capability >= 8.9 and add tests. ## Correctness and compatibility - Exhaustive finite-value check: all 65,536 FP16 bit patterns x 256 RNG values, with zero Triton/CUDA mismatches for both FP8 formats. - Offline compilation: 96 combinations across SM89, SM90, SM100, SM120, all input/output types, and all block sizes. - Every compiled variant uses 4 warps, 0 bytes shared memory, and 0 bytes global scratch memory. ## RTX 5090 performance BF16 to E4M3FN, warmed kernels, median latency, rotating working set >= 256 MiB: | Elements | CUDA | Triton | Eager | Triton vs CUDA | Triton vs Eager | |---:|---:|---:|---:|---:|---:| | 1,048,576 | 0.0136 ms | 0.0266 ms | 0.3538 ms | 1.95x slower | 13.3x faster | | 8,388,608 | 0.0329 ms | 0.0263 ms | 0.3696 ms | 1.25x faster | 14.1x faster | | 16,777,216 | 0.0628 ms | 0.0415 ms | 0.8942 ms | 1.51x faster | 21.5x faster | For 16,777,216 contiguous BF16 elements, both CUDA and Triton add 0 MiB of peak PyTorch-allocated memory because the output aliases the RNG buffer. Eager adds 240 MiB. ## Validation ```text 118 passed, 4 skipped ruff: All checks passed git diff --check: passed ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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