From 0e855b8e55e5402c2331cf874caeb3bb30a30bcf Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Sat, 15 Aug 2026 02:22:37 +0700 Subject: [PATCH] Prototype Vortex native NVFP4 quantizer --- src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp | 15 +++ src/h3_blackwell_runtime/csrc/nvfp4_scale.cu | 119 ++++++++++++++++++ src/h3_blackwell_runtime/nvfp4_quant.py | 47 +++++++ tools/profile_nvfp4_linear.py | 6 +- 4 files changed, 185 insertions(+), 2 deletions(-) diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp index f67f1e2..f600d84 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp @@ -1,7 +1,10 @@ #include +#include + torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor); torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divisor, torch::Tensor partials, torch::Tensor output, int64_t blocks, int64_t threads); +std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x); torch::Tensor nvfp4_activation_scale(torch::Tensor input, double divisor) { TORCH_CHECK(input.is_cuda(), "nvfp4_activation_scale expects a CUDA tensor"); @@ -26,7 +29,19 @@ torch::Tensor nvfp4_activation_scale_into(torch::Tensor input, double divisor, t return nvfp4_activation_scale_into_cuda(input, divisor, partials, output, blocks, threads); } +std::vector quantize_nvfp4_bf16(torch::Tensor input, torch::Tensor scale, bool pad_16x) { + TORCH_CHECK(input.is_cuda(), "quantize_nvfp4_bf16 expects a CUDA tensor"); + TORCH_CHECK(input.is_contiguous(), "quantize_nvfp4_bf16 expects contiguous input"); + TORCH_CHECK(input.dim() == 2, "quantize_nvfp4_bf16 expects a 2D tensor"); + TORCH_CHECK(input.scalar_type() == torch::kBFloat16, "quantize_nvfp4_bf16 expects bfloat16 input"); + TORCH_CHECK(scale.is_cuda(), "quantize_nvfp4_bf16 expects a CUDA scale tensor"); + TORCH_CHECK(scale.device() == input.device(), "quantize_nvfp4_bf16 scale must be on the input device"); + TORCH_CHECK(scale.numel() == 1, "quantize_nvfp4_bf16 scale must be scalar"); + return quantize_nvfp4_bf16_cuda(input, scale, pad_16x); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_activation_scale", &nvfp4_activation_scale, "Vortex NVFP4 activation scale"); m.def("nvfp4_activation_scale_into", &nvfp4_activation_scale_into, "Vortex NVFP4 activation scale with caller workspace"); + m.def("quantize_nvfp4_bf16", &quantize_nvfp4_bf16, "Vortex BF16 to TensorCore NVFP4 quantizer"); } diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu index 77184bd..40c874c 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu @@ -2,14 +2,21 @@ #include #include +#include + #include #include #include +#include namespace { constexpr int kThreads = 256; +int64_t roundup(int64_t value, int64_t multiple) { + return ((value + multiple - 1) / multiple) * multiple; +} + __inline__ __device__ float warp_reduce_max(float value) { for (int offset = 16; offset > 0; offset >>= 1) { value = fmaxf(value, __shfl_down_sync(0xffffffff, value, offset)); @@ -43,6 +50,90 @@ __inline__ __device__ float bf16_abs_bits_to_float(uint32_t bits) { return __uint_as_float(bits << 16); } +__inline__ __device__ float bf16_bits_to_float(uint16_t bits) { + return __uint_as_float(static_cast(bits) << 16); +} + +__inline__ __device__ uint8_t encode_fp4_e2m1(float value) { + const bool negative = signbit(value); + float abs_value = fabsf(value); + uint8_t code = 0; + if (abs_value >= 5.0f) { + code = 7; + } else if (abs_value >= 3.5f) { + code = 6; + } else if (abs_value >= 2.5f) { + code = 5; + } else if (abs_value >= 1.75f) { + code = 4; + } else if (abs_value >= 1.25f) { + code = 3; + } else if (abs_value >= 0.75f) { + code = 2; + } else if (abs_value >= 0.25f) { + code = 1; + } + return negative && code != 0 ? static_cast(code | 0x8u) : code; +} + +__inline__ __device__ uint8_t encode_fp8_e4m3(float value) { + __nv_fp8_e4m3 encoded(value); + return *reinterpret_cast(&encoded); +} + +__inline__ __device__ float decode_fp8_e4m3(uint8_t value) { + __nv_fp8_e4m3 encoded; + *reinterpret_cast(&encoded) = value; + return static_cast(encoded); +} + +__global__ void quantize_nvfp4_bf16_kernel( + const uint16_t* __restrict__ input, + const float* __restrict__ scale, + uint8_t* __restrict__ qdata, + uint8_t* __restrict__ block_scale, + int64_t rows, + int64_t cols, + int64_t q_rows, + int64_t q_cols, + int64_t scale_rows, + int64_t scale_cols) { + const int64_t row = blockIdx.x; + const float tensor_scale = scale[0]; + for (int64_t block_col = threadIdx.x; block_col < scale_cols; block_col += blockDim.x) { + float local_max = 0.0f; + float values[16]; + #pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t col = block_col * 16 + i; + float value = 0.0f; + if (row < rows && col < cols) { + value = bf16_bits_to_float(input[row * cols + col]); + } + values[i] = value; + local_max = fmaxf(local_max, fabsf(value)); + } + uint8_t scale_byte = 0; + float block_scale_value = 0.0f; + if (local_max > 0.0f && tensor_scale > 0.0f) { + scale_byte = encode_fp8_e4m3(local_max / (tensor_scale * 6.0f)); + block_scale_value = decode_fp8_e4m3(scale_byte); + } + if (row < scale_rows) { + block_scale[row * scale_cols + block_col] = scale_byte; + } + if (row < q_rows) { + #pragma unroll + for (int pair = 0; pair < 8; ++pair) { + const float denom = tensor_scale * block_scale_value; + const uint8_t even = denom > 0.0f ? encode_fp4_e2m1(values[pair * 2] / denom) : 0; + const uint8_t odd = denom > 0.0f ? encode_fp4_e2m1(values[pair * 2 + 1] / denom) : 0; + qdata[row * q_cols + block_col * 8 + pair] = static_cast((even << 4) | odd); + } + } + } +} + template __global__ void partial_absmax_kernel(const scalar_t* __restrict__ input, float* __restrict__ partials, int64_t numel) { __shared__ float shared[kThreads]; @@ -171,3 +262,31 @@ torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divis C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } + +std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x) { + c10::cuda::CUDAGuard device_guard(input.device()); + const int64_t rows = input.size(0); + const int64_t cols = input.size(1); + TORCH_CHECK(cols % 2 == 0, "quantize_nvfp4_bf16 expects an even feature dimension"); + const int64_t q_rows = pad_16x ? roundup(rows, 16) : rows; + const int64_t q_cols = pad_16x ? roundup(cols, 16) / 2 : cols / 2; + const int64_t scale_rows = roundup(q_rows, 128); + const int64_t scale_cols = q_cols / 8; + auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8)); + auto block_scale = torch::empty({scale_rows, scale_cols}, input.options().dtype(torch::kFloat8_e4m3fn)); + auto stream = at::cuda::getCurrentCUDAStream(); + constexpr int threads = 256; + quantize_nvfp4_bf16_kernel<<>>( + reinterpret_cast(input.data_ptr()), + scale.data_ptr(), + qdata.data_ptr(), + reinterpret_cast(block_scale.data_ptr()), + rows, + cols, + q_rows, + q_cols, + scale_rows, + scale_cols); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {qdata, block_scale}; +} diff --git a/src/h3_blackwell_runtime/nvfp4_quant.py b/src/h3_blackwell_runtime/nvfp4_quant.py index ab92d3b..fb1987a 100644 --- a/src/h3_blackwell_runtime/nvfp4_quant.py +++ b/src/h3_blackwell_runtime/nvfp4_quant.py @@ -166,3 +166,50 @@ def vortex_quantize_nvfp4( ), ), ) + + +def vortex_native_quantize_nvfp4( + tensor: torch.Tensor, + *, + scale: torch.Tensor | float | None = None, + timings: dict[str, list[float]] | None = None, +): + """Prototype native BF16 activation packer for TensorCoreNVFP4Layout.""" + if tensor.dim() != 2: + raise ValueError(f"NVFP4 activation quantization requires a 2D tensor, got {tensor.dim()}D") + if tensor.dtype != torch.bfloat16: + raise ValueError("vortex_native_quantize_nvfp4 currently expects BF16 input") + if not tensor.is_cuda or not tensor.is_contiguous(): + raise ValueError("vortex_native_quantize_nvfp4 requires contiguous CUDA input") + + from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout + + orig_shape = tuple(tensor.shape) + orig_dtype = tensor.dtype + if scale is None: + scale = nvfp4_activation_scale(tensor, timings=timings) + scale = _record_timing(timings, "scale_to_device", lambda: torch.as_tensor(scale, device=tensor.device, dtype=torch.float32)) + extension = _vortex_scale_extension() + qdata, block_scale = _record_timing( + timings, + "vortex_quantize_nvfp4", + lambda: extension.quantize_nvfp4_bf16( + tensor, + scale, + TensorCoreNVFP4Layout.get_padded_shape(orig_shape) != orig_shape, + ), + ) + return _record_timing( + timings, + "params_wrap", + lambda: QuantizedTensor( + qdata, + "TensorCoreNVFP4Layout", + TensorCoreNVFP4Layout.Params( + scale=scale, + orig_dtype=orig_dtype, + orig_shape=orig_shape, + block_scale=block_scale, + ), + ), + ) diff --git a/tools/profile_nvfp4_linear.py b/tools/profile_nvfp4_linear.py index 5f7d005..94a6a8f 100644 --- a/tools/profile_nvfp4_linear.py +++ b/tools/profile_nvfp4_linear.py @@ -19,7 +19,7 @@ from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_norm, rms_rop from h3_blackwell_runtime.block import H3DiTBlock, gate_segments, modulate_segments from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.nvfp4 import Nvfp4Linear -from h3_blackwell_runtime.nvfp4_quant import nvfp4_activation_scale, vortex_quantize_nvfp4 +from h3_blackwell_runtime.nvfp4_quant import nvfp4_activation_scale, vortex_native_quantize_nvfp4, vortex_quantize_nvfp4 from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.rope import h3_rope_rotation from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas @@ -74,6 +74,8 @@ def quantize_activation(flat_x: torch.Tensor, quantizer: str, scale: torch.Tenso if scale is None: raise ValueError("vortex_precomputed_scale requires a precomputed scale") return vortex_quantize_nvfp4(flat_x, scale=scale, timings=timings) + if quantizer == "vortex_native": + return vortex_native_quantize_nvfp4(flat_x, timings=timings) raise ValueError(f"Unsupported quantizer: {quantizer}") @@ -245,7 +247,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--block-index", type=int, default=24) parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2") parser.add_argument("--linears", nargs="+", choices=("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2"), default=("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2")) - parser.add_argument("--quantizers", nargs="+", choices=("comfy", "vortex_recalculate", "vortex_precomputed_scale"), default=("comfy", "vortex_recalculate", "vortex_precomputed_scale")) + parser.add_argument("--quantizers", nargs="+", choices=("comfy", "vortex_recalculate", "vortex_precomputed_scale", "vortex_native"), default=("comfy", "vortex_recalculate", "vortex_precomputed_scale")) parser.add_argument("--warmup", type=int, default=2) parser.add_argument("--iterations", type=int, default=5) parser.add_argument("--profiler-iterations", type=int, default=2)