[Do not review] Python Bindings for Block Quantization

csrc/tensor_metadata.cpp

@@ -344,9 +344,10 @@ inferAndValidateAllocationSizesAndStrides(
   auto [allocation_sizes, allocation_strides] =
       inferAllocationSizesAndStrides(tensor, tv, ee);
   // Only validate final sizes and strides when we have a non-empty tensor.
-  if (tensor.numel() != 0) {
-    validateAllocationSizesAndStrides(tv, allocation_sizes, allocation_strides);
-  }
+  // if (tensor.numel() != 0) {
+  //   validateAllocationSizesAndStrides(tv, allocation_sizes,
+  //   allocation_strides);
+  // }
   return {std::move(allocation_sizes), std::move(allocation_strides)};
 }
 

python/python_direct/ops.cpp

@@ -3486,6 +3486,78 @@ tuple[TensorView, TensorView, TensorView]
       py::return_value_policy::reference);
 }
 
+namespace {
+
+// Helper function to apply swizzle transformation to block scaling factors
+// for FP4 quantization. This transforms the memory layout to optimize access
+// patterns.
+void swizzleBlockScale(TensorView* tv_block_scale_fp8) {
+  // auto original_loop = tv_block_scale_fp8->getLoopDomain();
+  tv_block_scale_fp8->split(0, 128);
+  // m/128, 128, k
+  tv_block_scale_fp8->split(1, 32);
+  // m/128, 4(m_o), 32(m_i), k
+  tv_block_scale_fp8->split(3, 4);
+  // m/128, 4(m_o), 32(m_i), k/4, 4(k)
+  std::vector<IterDomain*> tv_block_scale_fp8_alloc{
+      tv_block_scale_fp8->axis(0),
+      tv_block_scale_fp8->axis(3),
+      tv_block_scale_fp8->axis(2),
+      tv_block_scale_fp8->axis(1),
+      tv_block_scale_fp8->axis(4)};
+  // m/128, k/4, 32(m_i), 4(m_o), 4(k)
+  tv_block_scale_fp8->setAllocationDomain(tv_block_scale_fp8_alloc, true);
+
+  // back to a 2D logical domain.
+  // m/128, 4(m_o), 32(m_i), k/4, 4(k) ->
+  // m/32, 32, k/4, 4(k)
+  tv_block_scale_fp8->merge(0);
+  // m/32, 32, k/4, 4(k) -> m, k/4, 4(k)
+  tv_block_scale_fp8->merge(0);
+  // m, k/4, 4(k) -> m, k
+  tv_block_scale_fp8->merge(-2);
+}
+
+} // namespace
+
+void bindQuantizationOps(py::module_& ops) {
+  ops.def(
+      "nv_block_quantize",
+      [](TensorView* input,
+         TensorView* global_scale,
+         bool swizzle_block_scales,
+         int64_t block_size,
+         PrimDataType dtype) -> py::tuple {
+        auto output = blockQuantize(input, global_scale, block_size, dtype);
+        if (swizzle_block_scales) {
+          swizzleBlockScale(output.block_scales);
+        }
+        return py::make_tuple(output.quantized_tensor, output.block_scales);
+      },
+      py::arg("input"),
+      py::arg("global_scale").none(true) = py::none(),
+      py::arg("swizzle_block_scales") = false,
+      py::arg("block_size") = 16,
+      py::arg("dtype") = DataType::Float4_e2m1fn,
+      R"(
+Block quantize tensor to NVFP4 format.
+Parameters
+----------
+input : TensorView
+    Input tensor to quantize. Must be a floating point tensor.
+global_scale : Val or TensorView, optional
+block_size : int, optional
+    Block size for quantization. Default is 16.
+Returns
+-------
+tuple[TensorView, TensorView]
+    A tuple containing (block_scales, quantized_tensor) where:
+    - block_scales: Per-block scaling factors
+    - quantized_tensor: Quantized tensor in NVFP4 format
+      )",
+      py::return_value_policy::reference);
+}
+
 template <
     class ShapeType,
     TensorView* (*RandomFuncWithSeed)(
@@ -3638,6 +3710,7 @@ void bindOperations(py::module& nvfuser) {
   bindSearchOps(nvf_ops);
   bindSdpaOps(nvf_ops);
   bindRandomOps(nvf_ops);
+  bindQuantizationOps(nvf_ops);
 }
 
 } // namespace nvfuser::python

python/python_direct/python_translate.cpp

@@ -1589,6 +1589,33 @@ class PythonTranslator : public OptInConstDispatch {
         {out_tv});
   }
 
+  // Map BlockQuantizationOp to python frontend
+  void handle(const BlockQuantizationOp* bqop) final {
+    NVF_ERROR(bqop != nullptr);
+    visited_vals_.insert(bqop->output(0));
+    visited_vals_.insert(bqop->output(1));
+
+    static const auto default_args = std::make_tuple(
+        KeywordArgument<decltype(bqop->globalScale())>{"global_scale", nullptr},
+        KeywordArgument<int64_t>{"block_size", 16},
+        KeywordArgument<bool>{"swizzle_block_scales", false},
+        KeywordArgument<DataType>{"dtype", DataType::Float4_e2m1fn});
+
+    auto dtype = bqop->quantizedOutput()->as<TensorView>()->dtype();
+    auto swizzled_block_scale =
+        bqop->blockScales()->as<TensorView>()->hasAllocation();
+    printer_.generateKwargsOperation(
+        "fd.ops.nv_block_quantize",
+        std::make_tuple(bqop->in()),
+        default_args,
+        std::make_tuple(
+            bqop->globalScale(),
+            bqop->blockSize(),
+            swizzled_block_scale,
+            dtype),
+        std::vector<const nvfuser::Val*>{bqop->output(0), bqop->output(1)});
+  }
+
   // Map EmbeddingFwdOp to python frontend
   void handle(const EmbeddingFwdOp* eop) final {
     NVF_ERROR(eop != nullptr);

runtime/block_quantization_kernels.cu

@@ -86,7 +86,7 @@ __device__ void block_quantize_to_nvfp4(
     if constexpr (std::is_same<T, float>::value) {
       vec_in[i] = input[i];
     } else if constexpr (std::is_same<T, __bfloat>::value) {
-      vec_in[i] = __bfloat2float(input[i]);
+      vec_in[i] = __bfloat2float(__float2bfloat(__bfloat2float(input[i])));
     } else if constexpr (std::is_same<T, __half>::value) {
       vec_in[i] = __half2float(input[i]);
     }
@@ -107,22 +107,23 @@ __device__ void block_quantize_to_nvfp4(
 
   // This division should be replaced with a multiplication
   // by a reciprocal for better performance.
-  float scaled_max = block_max / 6.000000000e+00f;
+  // float scaled_max = block_max / 6.000000000e+00f;
 
+  constexpr float rcp_6f = 1.0f / 6.0f;
+
+  float scaled_max = 0.0f;
   if constexpr (USE_GLOBAL_SCALE) {
-    scaled_max = scaled_max * global_scale[0];
+    scaled_max = block_max * global_scale[0] * rcp_6f;
+  } else {
+    scaled_max = block_max / 6.000000000e+00f;
   }
 
-  float clamped_max = clamp(
-      scaled_max, 1.562500000e-02f, 4.480000000e+02f); // Clamp between 0 and 1
-
-  __e4m3 clamped_max_fp8 = __float2e4m3(clamped_max);
+  __e4m3 clamped_max_fp8 = __float2e4m3(scaled_max);
 
-  // Convert back from FP8 to float using __e4m32float
-  float clamped_max_converted = __e4m32float(clamped_max_fp8);
+  float clamped_max = __e4m32float(clamped_max_fp8);
 
   if constexpr (USE_GLOBAL_SCALE) {
-    clamped_max_converted = clamped_max_converted / global_scale[0];
+    clamped_max = global_scale[0] / clamped_max;
   }
 
   // Write out the block scaling factor to global memory.
@@ -165,8 +166,8 @@ __device__ void block_quantize_to_nvfp4(
   Array<float, ITEMS_PER_THREAD, ITEMS_PER_THREAD> clamped_vals;
 #pragma unroll
   for (int i = 0; i < ITEMS_PER_THREAD; ++i) {
-    float scaled_val = vec_in[i] / clamped_max_converted;
-    clamped_vals[i] = clamp(scaled_val, -6.000000000e+00f, 6.000000000e+00f);
+    // float scaled_val = vec_in[i] / clamped_max;
+    clamped_vals[i] = vec_in[i] * clamped_max;
   }
 
   Array<__e2m1, ITEMS_PER_THREAD, 1> fp4_vals;