mlx 0.30.7.2 → 0.30.7.6

data/lib/mlx-onnx/webgpu_harness.rb

@@ -0,0 +1,289 @@
+# frozen_string_literal: true
+
+require "json"
+require "open3"
+require "fileutils"
+
+module MLX
+  module ONNX
+    module WebGPUHarness
+      module_function
+
+      def export_onnx_webgpu_harness(
+        target_dir,
+        payload_or_source,
+        opset: 18,
+        model_name: "mlx_graph",
+        execution_providers: %w[webgpu wasm],
+        benchmark_warmup_runs: 2,
+        benchmark_measure_runs: 10,
+        external_data: false,
+        external_data_size_threshold: 1024,
+        external_data_file: nil
+      )
+        output_dir = file_path(target_dir)
+        raise ArgumentError, "target_dir must not be empty" if output_dir.empty?
+
+        providers = normalize_web_execution_providers(execution_providers)
+        warmup_runs = normalize_non_negative_integer(
+          benchmark_warmup_runs,
+          "benchmark_warmup_runs"
+        )
+        measure_runs = normalize_positive_integer(
+          benchmark_measure_runs,
+          "benchmark_measure_runs"
+        )
+
+        FileUtils.mkdir_p(output_dir)
+        model_filename = "model.onnx"
+        model_path = File.join(output_dir, model_filename)
+        onnx_json = MLX::ONNX.graph_ir_to_onnx_json(
+          payload_or_source,
+          opset: opset,
+          model_name: model_name
+        )
+        MLX::ONNX.graph_ir_to_onnx(
+          model_path,
+          payload_or_source,
+          opset: opset,
+          model_name: model_name,
+          external_data: external_data,
+          external_data_size_threshold: external_data_size_threshold,
+          external_data_file: external_data_file
+        )
+
+        stub = JSON.parse(onnx_json)
+        input_specs = stub.fetch("graph").fetch("inputs")
+        input_examples = build_input_examples(input_specs)
+
+        manifest = {
+          "format" => "onnx_webgpu_harness_v1",
+          "model" => model_filename,
+          "execution_providers" => providers,
+          "benchmark" => {
+            "warmup_runs" => warmup_runs,
+            "measure_runs" => measure_runs
+          },
+          "inputs" => input_specs.map do |spec|
+            {
+              "name" => spec.fetch("name"),
+              "shape" => spec.fetch("shape"),
+              "dtype" => spec.fetch("dtype")
+            }
+          end
+        }
+        if external_data
+          manifest["external_data"] = [
+            external_data_file.nil? ? "model.data" : external_data_file.to_s
+          ]
+        end
+
+        File.binwrite(
+          File.join(output_dir, "harness.manifest.json"),
+          JSON.pretty_generate(manifest)
+        )
+        File.binwrite(
+          File.join(output_dir, "inputs.example.json"),
+          JSON.pretty_generate(input_examples)
+        )
+        copy_assets!(output_dir)
+
+        manifest
+      end
+
+      def smoke_test_onnx_webgpu_harness(
+        harness_dir,
+        timeout_seconds: 30,
+        mock_ort: false,
+        local_ort: true,
+        node_bin: ENV.fetch("NODE", "node")
+      )
+        directory = file_path(harness_dir)
+        raise ArgumentError, "harness_dir must not be empty" if directory.empty?
+
+        directory = File.expand_path(directory)
+        unless Dir.exist?(directory)
+          raise ArgumentError, "harness_dir does not exist: #{directory}"
+        end
+
+        timeout = normalize_positive_integer(timeout_seconds, "timeout_seconds")
+        mock = normalize_boolean(mock_ort, "mock_ort")
+        local = normalize_boolean(local_ort, "local_ort")
+        node = node_bin.to_s
+        raise ArgumentError, "node_bin must not be empty" if node.empty?
+
+        smoke_script = web_harness_smoke_script_path
+        unless File.file?(smoke_script)
+          raise RuntimeError, "missing web harness smoke script: #{smoke_script}"
+        end
+
+        argv = [
+          node,
+          smoke_script,
+          "--harness-dir",
+          directory,
+          "--timeout-seconds",
+          timeout.to_s
+        ]
+        argv << "--mock-ort" if mock
+        argv << (local ? "--local-ort" : "--no-local-ort")
+
+        stdout, stderr, status = Open3.capture3(*argv, chdir: web_root_dir)
+        unless status.success?
+          raise RuntimeError, <<~MSG
+            web harness smoke test failed: #{argv.join(" ")}
+            stdout:
+            #{stdout}
+            stderr:
+            #{stderr}
+          MSG
+        end
+
+        telemetry = begin
+          JSON.parse(stdout)
+        rescue JSON::ParserError => e
+          raise RuntimeError, <<~MSG
+            web harness smoke test produced invalid JSON: #{e.message}
+            stdout:
+            #{stdout}
+            stderr:
+            #{stderr}
+          MSG
+        end
+
+        unless telemetry.is_a?(Hash)
+          raise RuntimeError, "web harness smoke test produced non-object telemetry"
+        end
+        unless telemetry.fetch("format", nil) == "onnx_webgpu_telemetry_v1"
+          raise RuntimeError, "unexpected web harness telemetry format: #{telemetry.fetch('format', nil).inspect}"
+        end
+
+        telemetry
+      end
+
+      def file_path(file)
+        if file.respond_to?(:to_path)
+          file.to_path.to_s
+        else
+          file.to_s
+        end
+      end
+      private_class_method :file_path
+
+      def normalize_web_execution_providers(value)
+        providers = if value.is_a?(::Array)
+          value
+        else
+          [value]
+        end
+        providers = providers.map(&:to_s)
+        raise ArgumentError, "execution_providers must contain at least one provider" if providers.empty?
+
+        allowed = %w[webgpu wasm]
+        providers.each do |provider|
+          unless allowed.include?(provider)
+            raise ArgumentError, "execution_providers contains unsupported provider #{provider.inspect}"
+          end
+        end
+        providers.uniq
+      end
+      private_class_method :normalize_web_execution_providers
+
+      def normalize_non_negative_integer(value, label)
+        integer = begin
+          Integer(value)
+        rescue ArgumentError, TypeError
+          raise ArgumentError, "#{label} must be a non-negative Integer"
+        end
+        raise ArgumentError, "#{label} must be a non-negative Integer" if integer.negative?
+
+        integer
+      end
+      private_class_method :normalize_non_negative_integer
+
+      def normalize_positive_integer(value, label)
+        integer = begin
+          Integer(value)
+        rescue ArgumentError, TypeError
+          raise ArgumentError, "#{label} must be a positive Integer"
+        end
+        raise ArgumentError, "#{label} must be a positive Integer" unless integer.positive?
+
+        integer
+      end
+      private_class_method :normalize_positive_integer
+
+      def normalize_boolean(value, label)
+        unless value == true || value == false
+          raise ArgumentError, "#{label} must be true or false"
+        end
+
+        value
+      end
+      private_class_method :normalize_boolean
+
+      def build_input_examples(input_specs)
+        input_specs.each_with_object({}) do |spec, out|
+          out[spec.fetch("name")] = build_zero_tensor_values(
+            spec.fetch("shape"),
+            spec.fetch("dtype")
+          )
+        end
+      end
+      private_class_method :build_input_examples
+
+      def build_zero_tensor_values(shape, dtype)
+        if shape.empty?
+          zero_leaf_value_for_dtype(dtype)
+        else
+          ::Array.new(shape.first) { build_zero_tensor_values(shape[1..], dtype) }
+        end
+      end
+      private_class_method :build_zero_tensor_values
+
+      def zero_leaf_value_for_dtype(dtype)
+        if dtype == "bool" || dtype == "bool_"
+          false
+        elsif dtype == "complex64"
+          { "__mlx_complex__" => [0.0, 0.0] }
+        elsif dtype.start_with?("float") || dtype == "bfloat16"
+          0.0
+        else
+          0
+        end
+      end
+      private_class_method :zero_leaf_value_for_dtype
+
+      def copy_assets!(output_dir)
+        template_dir = web_harness_template_dir
+        unless Dir.exist?(template_dir)
+          raise RuntimeError, "missing web harness template directory: #{template_dir}"
+        end
+
+        %w[index.html harness.js].each do |file_name|
+          source = File.join(template_dir, file_name)
+          unless File.file?(source)
+            raise RuntimeError, "missing web harness template file: #{source}"
+          end
+          FileUtils.cp(source, File.join(output_dir, file_name))
+        end
+      end
+      private_class_method :copy_assets!
+
+      def web_harness_template_dir
+        File.expand_path("../../web/onnx_webgpu_harness", __dir__)
+      end
+      private_class_method :web_harness_template_dir
+
+      def web_harness_smoke_script_path
+        File.join(web_harness_template_dir, "browser_smoke.mjs")
+      end
+      private_class_method :web_harness_smoke_script_path
+
+      def web_root_dir
+        File.expand_path("../../web", __dir__)
+      end
+      private_class_method :web_root_dir
+    end
+  end
+end

data/{mlx → submodules/mlx}/mlx/backend/cuda/cublas_utils.cpp

@@ -105,13 +105,6 @@ void CublasMatmulBase::init_base(
   CHECK_CUBLAS_ERROR(
       cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));
 
-  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_POINTER_MODE,
-      &pointer_mode,
-      sizeof(int32_t)));
-
   // In cublasLt matrices use column-major layout, while it is possible to use
   // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
   // epilogue does not work with the option. So instead we swap A and B to make

data/{mlx → submodules/mlx}/mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -73,6 +73,14 @@ CublasGemm::CublasGemm(
       batch_count,
       a_batch_stride,
       b_batch_stride);
+
+  // alpha and beta are both host pointers
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
 }
 
 CublasGemm::CublasGemm(
@@ -215,8 +223,8 @@ void CublasGemm::execute(
     const void* a,
     const void* b,
     const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
   const void* alpha_ptr = α
   const void* beta_ptr = β
   complex64_t alpha_c, beta_c;

data/{mlx → submodules/mlx}/mlx/backend/cuda/quantized/cublas_qqmm.cpp

@@ -13,39 +13,26 @@ namespace mlx::core {
 
 namespace {
 
-// Currently cublas supports only mxfp8 and nvfp4
-// quantization modes for block scaled quantization
-cudaDataType_t qmode_to_cublas_scale_dtype(std::string mode) {
-  if (mode == "mxfp8") {
-    return CUDA_R_8F_UE8M0;
-  } else if (mode == "nvfp4") {
-    return CUDA_R_8F_UE4M3;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
-  }
-}
-
-cudaDataType_t qmode_to_cublas_dtype(std::string mode) {
-  if (mode == "mxfp8") {
-    return CUDA_R_8F_E4M3;
-  } else if (mode == "nvfp4") {
-    return CUDA_R_4F_E2M1;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
-  }
-}
+struct QuantModeConfig {
+  cudaDataType_t data_type;
+  cudaDataType_t scale_dtype;
+  cublasLtMatmulMatrixScale_t scale_mode;
+};
 
-cublasLtMatmulMatrixScale_t qmode_to_cublas_scale_mode(std::string mode) {
+QuantModeConfig get_quant_mode_config(const std::string& mode) {
   if (mode == "mxfp8") {
-    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;
+    return {
+        CUDA_R_8F_E4M3,
+        CUDA_R_8F_UE8M0,
+        CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0};
   } else if (mode == "nvfp4") {
-    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+    return {
+        CUDA_R_4F_E2M1,
+        CUDA_R_8F_UE4M3,
+        CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3};
   }
+  throw std::runtime_error(
+      fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
 }
 
 } // namespace
@@ -64,21 +51,21 @@ CublasQQMM::CublasQQMM(
     int64_t a_batch_stride,
     int64_t b_batch_stride,
     Dtype out_dtype,
-    std::string qmode) {
+    const std::string& qmode) {
+  auto config = get_quant_mode_config(qmode);
+
   // The compute type must be CUBLAS_COMPUTE_32F.
   // The scale type must be CUDA_R_32F.
   cudaDataType_t scale_type = CUDA_R_32F;
   cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F;
   cudaDataType_t output_type =
       cublas_utils::dtype_to_cublas_type(out_dtype, "CublasQQMM");
-  cudaDataType_t data_type = qmode_to_cublas_dtype(qmode);
-  quantization_mode_ = std::string(qmode);
 
   init_base(
       device,
       scale_type,
       gemm_compute_type,
-      data_type,
+      config.data_type,
       output_type,
       a_transposed,
       a_rows,
@@ -92,8 +79,8 @@ CublasQQMM::CublasQQMM(
       a_batch_stride,
       b_batch_stride);
 
-  a_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
-  b_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+  a_scale_mode_ = config.scale_mode;
+  b_scale_mode_ = config.scale_mode;
 
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
       matmul_desc_,
@@ -123,7 +110,7 @@ CublasQQMM::CublasQQMM(
     int64_t b_batch_stride,
     int64_t c_batch_stride,
     Dtype out_dtype,
-    std::string qmode)
+    const std::string& qmode)
     : CublasQQMM(
           device,
           a_transposed,
@@ -158,11 +145,14 @@ void CublasQQMM::run(
     const array& b,
     const array& a_scale,
     const array& b_scale,
-    float alpha) {
+    const array& alpha,
+    const array& beta) {
   encoder.set_input_array(a);
   encoder.set_input_array(b);
   encoder.set_input_array(a_scale);
   encoder.set_input_array(b_scale);
+  encoder.set_input_array(alpha);
+  encoder.set_input_array(beta);
   encoder.set_output_array(out);
 
   execute(
@@ -173,19 +163,37 @@ void CublasQQMM::run(
       gpu_ptr<void>(a_scale),
       gpu_ptr<void>(b_scale),
       nullptr,
-      alpha);
+      gpu_ptr<void>(alpha),
+      gpu_ptr<void>(beta));
 }
 
-void CublasQQMM::execute(
+void CublasQQMM::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(a_scale);
+  encoder.set_input_array(b_scale);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(a_scale),
+      gpu_ptr<void>(b_scale),
+      nullptr);
+}
+
+void CublasQQMM::set_scales_ptrs(
     cu::CommandEncoder& encoder,
-    void* out,
-    const void* a,
-    const void* b,
     const void* a_scale,
-    const void* b_scale,
-    const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const void* b_scale) {
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
       matmul_desc_,
       CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
@@ -196,6 +204,49 @@ void CublasQQMM::execute(
       CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
       &a_scale,
       sizeof(a_scale)));
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    const void* alpha,
+    const void* beta) {
+  set_scales_ptrs(encoder, a_scale, b_scale);
+  // alpha and beta are both should be device pointers for nvfp4
+  // by default cublas uses host pointers
+  // https://docs.nvidia.com/cuda/cublas/#cublasltpointermode-t
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
+  execute_matmul(encoder, out, a, b, c, alpha, beta);
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
+  set_scales_ptrs(encoder, a_scale, b_scale);
+  // alpha and beta are both should be host pointers
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
 
   const void* alpha_ptr = &alpha;
   const void* beta_ptr = &beta;

data/{mlx → submodules/mlx}/mlx/backend/cuda/quantized/cublas_qqmm.h

@@ -25,7 +25,7 @@ class CublasQQMM : public CublasMatmulBase {
       int64_t a_batch_stride,
       int64_t b_batch_stride,
       Dtype out_dtype,
-      std::string quantization_mode);
+      const std::string& quantization_mode);
 
   CublasQQMM(
       cu::Device& device,
@@ -43,7 +43,7 @@ class CublasQQMM : public CublasMatmulBase {
       int64_t b_batch_stride,
       int64_t c_batch_stride,
       Dtype out_dtype,
-      std::string quantization_mode);
+      const std::string& quantization_mode);
 
   void run(
       cu::CommandEncoder& encoder,
@@ -52,20 +52,33 @@ class CublasQQMM : public CublasMatmulBase {
       const array& b,
       const array& a_scale,
       const array& b_scale,
-      float alpha = 1.0f);
+      const array& alpha,
+      const array& beta);
 
- private:
-  void run_batched(
+  void run(
       cu::CommandEncoder& encoder,
       array& out,
       const array& a,
       const array& b,
       const array& a_scale,
-      const array& b_scale,
-      const Shape& batch_shape,
-      const Strides& a_batch_strides,
-      const Strides& b_batch_strides,
-      float alpha);
+      const array& b_scale);
+
+ private:
+  void set_scales_ptrs(
+      cu::CommandEncoder& encoder,
+      const void* a_scale,
+      const void* b_scale);
+
+  void execute(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* a_scale,
+      const void* b_scale,
+      const void* c,
+      const void* alpha,
+      const void* beta);
 
   void execute(
       cu::CommandEncoder& encoder,
@@ -75,10 +88,9 @@ class CublasQQMM : public CublasMatmulBase {
       const void* a_scale,
       const void* b_scale,
       const void* c,
-      float alpha = 1,
-      float beta = 0);
+      const float alpha = 1.0f,
+      const float beta = 0.0f);
 
-  std::string quantization_mode_;
   cublasLtMatmulMatrixScale_t a_scale_mode_;
   cublasLtMatmulMatrixScale_t b_scale_mode_;
   cublasLtMatmulMatrixScale_t c_scale_mode_;