Add Vortex NVFP4 scale kernel

src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp

@@ -0,0 +1,14 @@
+#include <torch/extension.h>
+
+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");
+}

src/h3_blackwell_runtime/csrc/nvfp4_scale.cu

@@ -0,0 +1,78 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+
+#include <cmath>
+#include <limits>
+
+namespace {
+
+constexpr int kThreads = 256;
+
+template <typename scalar_t>
+__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<int64_t>(blockDim.x) * gridDim.x;
+  int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + tid;
+  float local_max = 0.0f;
+  while (index < numel) {
+    const float value = static_cast<float>(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<int>(std::min<int64_t>((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<scalar_t><<<blocks, kThreads, 0, stream>>>(
+        input.data_ptr<scalar_t>(), partials.data_ptr<float>(), numel);
+  });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  final_scale_kernel<<<1, kThreads, 0, stream>>>(partials.data_ptr<float>(), output.data_ptr<float>(), blocks, static_cast<float>(divisor));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}

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(