Prototype Vortex native NVFP4 quantizer

src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp

@@ -1,7 +1,10 @@
 #include <torch/extension.h>
 
+#include <vector>
+
 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<torch::Tensor> 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<torch::Tensor> 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");
 }

src/h3_blackwell_runtime/csrc/nvfp4_scale.cu

@@ -2,14 +2,21 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <torch/extension.h>
 
+#include <cuda_fp8.h>
+
 #include <cmath>
 #include <cstdint>
 #include <limits>
+#include <vector>
 
 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<uint32_t>(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<uint8_t>(code | 0x8u) : code;
+}
+
+__inline__ __device__ uint8_t encode_fp8_e4m3(float value) {
+  __nv_fp8_e4m3 encoded(value);
+  return *reinterpret_cast<uint8_t*>(&encoded);
+}
+
+__inline__ __device__ float decode_fp8_e4m3(uint8_t value) {
+  __nv_fp8_e4m3 encoded;
+  *reinterpret_cast<uint8_t*>(&encoded) = value;
+  return static_cast<float>(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<uint8_t>((even << 4) | odd);
+      }
+    }
+  }
+}
+
 template <typename scalar_t>
 __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<torch::Tensor> 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<<<scale_rows, threads, 0, stream>>>(
+      reinterpret_cast<const uint16_t*>(input.data_ptr()),
+      scale.data_ptr<float>(),
+      qdata.data_ptr<uint8_t>(),
+      reinterpret_cast<uint8_t*>(block_scale.data_ptr()),
+      rows,
+      cols,
+      q_rows,
+      q_cols,
+      scale_rows,
+      scale_cols);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {qdata, block_scale};
+}

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,
+            ),
+        ),
+    )

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)