diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp new file mode 100644 index 0000000..d5a8371 --- /dev/null +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp @@ -0,0 +1,14 @@ +#include + +torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor); + +torch::Tensor nvfp4_activation_scale(torch::Tensor input, double divisor) { + TORCH_CHECK(input.is_cuda(), "nvfp4_activation_scale expects a CUDA tensor"); + TORCH_CHECK(input.is_contiguous(), "nvfp4_activation_scale expects contiguous input"); + TORCH_CHECK(input.dim() == 2, "nvfp4_activation_scale expects a 2D tensor"); + return nvfp4_activation_scale_cuda(input, divisor); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("nvfp4_activation_scale", &nvfp4_activation_scale, "Vortex NVFP4 activation scale"); +} diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu new file mode 100644 index 0000000..9783f07 --- /dev/null +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu @@ -0,0 +1,78 @@ +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kThreads = 256; + +template +__global__ void partial_absmax_kernel(const scalar_t* __restrict__ input, float* __restrict__ partials, int64_t numel) { + __shared__ float shared[kThreads]; + const int tid = threadIdx.x; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + float local_max = 0.0f; + while (index < numel) { + const float value = static_cast(input[index]); + local_max = fmaxf(local_max, fabsf(value)); + index += stride; + } + shared[tid] = local_max; + __syncthreads(); + for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) { + if (tid < offset) { + shared[tid] = fmaxf(shared[tid], shared[tid + offset]); + } + __syncthreads(); + } + if (tid == 0) { + partials[blockIdx.x] = shared[0]; + } +} + +__global__ void final_scale_kernel(const float* __restrict__ partials, float* __restrict__ output, int64_t count, float divisor) { + __shared__ float shared[kThreads]; + const int tid = threadIdx.x; + float local_max = 0.0f; + for (int64_t index = tid; index < count; index += blockDim.x) { + local_max = fmaxf(local_max, partials[index]); + } + shared[tid] = local_max; + __syncthreads(); + for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) { + if (tid < offset) { + shared[tid] = fmaxf(shared[tid], shared[tid + offset]); + } + __syncthreads(); + } + if (tid == 0) { + output[0] = shared[0] / divisor; + } +} + +} // namespace + +torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor) { + c10::cuda::CUDAGuard device_guard(input.device()); + const auto numel = input.numel(); + TORCH_CHECK(numel > 0, "nvfp4_activation_scale input must be non-empty"); + TORCH_CHECK(divisor > 0.0, "nvfp4_activation_scale divisor must be positive"); + + const int blocks = static_cast(std::min((numel + kThreads - 1) / kThreads, 4096)); + auto partials = torch::empty({blocks}, input.options().dtype(torch::kFloat32)); + auto output = torch::empty({}, input.options().dtype(torch::kFloat32)); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "vortex_nvfp4_activation_scale", [&] { + partial_absmax_kernel<<>>( + input.data_ptr(), partials.data_ptr(), numel); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + final_scale_kernel<<<1, kThreads, 0, stream>>>(partials.data_ptr(), output.data_ptr(), blocks, static_cast(divisor)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} diff --git a/src/h3_blackwell_runtime/nvfp4_quant.py b/src/h3_blackwell_runtime/nvfp4_quant.py index 18b1ccb..d3e9bd4 100644 --- a/src/h3_blackwell_runtime/nvfp4_quant.py +++ b/src/h3_blackwell_runtime/nvfp4_quant.py @@ -2,7 +2,10 @@ from __future__ import annotations +import os import time +from functools import lru_cache +from pathlib import Path import torch @@ -23,12 +26,45 @@ def _record_timing(timings: dict[str, list[float]] | None, name: str, fn): return value +@lru_cache(maxsize=1) +def _vortex_scale_extension(): + from torch.utils.cpp_extension import load + + root = Path(__file__).resolve().parent + return load( + name="h3_vortex_nvfp4_scale", + sources=[str(root / "csrc" / "nvfp4_scale.cpp"), str(root / "csrc" / "nvfp4_scale.cu")], + extra_cflags=["-O3"], + extra_cuda_cflags=["-O3"], + verbose=os.getenv("H3_NVFP4_SCALE_VERBOSE", "").lower() in {"1", "true", "yes", "on"}, + ) + + def nvfp4_activation_scale(tensor: torch.Tensor, *, timings: dict[str, list[float]] | None = None) -> torch.Tensor: """Compute Comfy Kitchen's current per-tensor NVFP4 activation scale.""" from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX + divisor = float(F8_E4M3_MAX * F4_E2M1_MAX) + backend = os.getenv("H3_NVFP4_SCALE_BACKEND", "torch").lower() + if backend == "vortex": + if tensor.dim() != 2: + raise ValueError(f"Vortex NVFP4 scale backend requires 2D input, got {tensor.dim()}D") + if not tensor.is_cuda: + raise ValueError("Vortex NVFP4 scale backend requires CUDA input") + if not tensor.is_contiguous(): + raise ValueError("Vortex NVFP4 scale backend requires contiguous input") + try: + extension = _vortex_scale_extension() + return _record_timing(timings, "vortex_absmax_scale", lambda: extension.nvfp4_activation_scale(tensor, divisor)) + except Exception: + if os.getenv("H3_NVFP4_SCALE_STRICT", "").lower() in {"1", "true", "yes", "on"}: + raise + backend = "torch" + + if backend != "torch": + raise ValueError(f"Unsupported H3_NVFP4_SCALE_BACKEND={backend!r}") amax = _record_timing(timings, "scale_absmax", lambda: torch.amax(tensor.abs())) - return _record_timing(timings, "scale_finalize", lambda: amax / (F8_E4M3_MAX * F4_E2M1_MAX)) + return _record_timing(timings, "scale_finalize", lambda: amax / divisor) def vortex_quantize_nvfp4(