Support vllm fp8 quant and gemm

torch_xla/csrc/BUILD

@@ -64,6 +64,7 @@ ptxla_cc_library(
         "token_handler.cpp",
         "torch_util.cpp",
         "view.cpp",
+        "vllm_fp8.cpp",
         "xla_backend_impl.cpp",
         "xla_graph_executor.cpp",
         "xla_lower_util.cpp",

torch_xla/csrc/dtype.cpp

@@ -106,6 +106,8 @@ at::ScalarType TorchTypeFromXlaType(xla::PrimitiveType xla_type) {
       return at::ScalarType::ComplexFloat;
     case xla::PrimitiveType::C128:
       return at::ScalarType::ComplexDouble;
+    case xla::PrimitiveType::F8E4M3FN:
+      return at::ScalarType::Float8_e4m3fn;
     default:
       XLA_ERROR() << "XLA type not supported: " << xla_type;
   }
@@ -137,6 +139,8 @@ xla::PrimitiveType XlaTypeFromTorchType(at::ScalarType scalar_type) {
       return xla::PrimitiveType::C64;
     case at::ScalarType::ComplexDouble:
       return xla::PrimitiveType::C128;
+    case at::ScalarType::Float8_e4m3fn:
+      return xla::PrimitiveType::F8E4M3FN;
     default:
       XLA_ERROR() << "Type not supported: " << scalar_type;
   }

torch_xla/csrc/ops/cuda_custom_call.cpp

@@ -0,0 +1,44 @@
+#include "torch_xla/csrc/ops/cuda_custom_call.h"
+
+#include "torch_xla/csrc/lowering_context.h"
+#include "torch_xla/csrc/ops/xla_ops.h"
+#include "torch_xla/csrc/xla_lower_util.h"
+
+namespace torch_xla {
+
+CudaCustomCall::CudaCustomCall(torch::lazy::OpList inputs, int num_outputs,
+                               xla::Shape output_shape,
+                               const std::string& call_target_name,
+                               const std::string& opaque)
+    : XlaNode(xla_cuda_custom_call, inputs, std::move(output_shape),
+              num_outputs, torch::lazy::MHash(call_target_name + opaque)),
+      call_target_name_(call_target_name),
+      opaque_(opaque) {}
+
+torch::lazy::NodePtr CudaCustomCall::Clone(torch::lazy::OpList operands) const {
+  return torch::lazy::MakeNode<CudaCustomCall>(
+      operands, num_outputs(), xla_shape(), call_target_name_, opaque_);
+}
+
+XlaOpVector CudaCustomCall::Lower(LoweringContext* loctx) const {
+  std::vector<xla::XlaOp> inputs;
+  inputs.reserve(operands().size());
+  for (auto& operand : operands()) {
+    inputs.push_back(loctx->GetOutputOp(operand));
+  }
+  std::vector<xla::XlaOp> output = BuildCudaCustomCall(
+      inputs, num_outputs(), xla_shape(), call_target_name_, opaque_);
+  if (num_outputs() == 1) {
+    return ReturnOp(output[0], loctx);
+  }
+  return ReturnOps(output, loctx);
+}
+
+std::string CudaCustomCall::ToString() const {
+  std::stringstream ss;
+  ss << XlaNode::ToString() << ", call_target_name=" << call_target_name_
+     << ", opaque=" << opaque_;
+  return ss.str();
+}
+
+}  // namespace torch_xla

torch_xla/csrc/ops/cuda_custom_call.h

@@ -0,0 +1,27 @@
+#ifndef XLA_TORCH_XLA_CSRC_OPS_CUDA_CUSTOM_CALL_H_
+#define XLA_TORCH_XLA_CSRC_OPS_CUDA_CUSTOM_CALL_H_
+
+#include "torch_xla/csrc/ir.h"
+
+namespace torch_xla {
+
+class CudaCustomCall : public XlaNode {
+ public:
+  CudaCustomCall(torch::lazy::OpList inputs, int num_outputs,
+                 xla::Shape output_shape, const std::string& call_target_name,
+                 const std::string& opaque);
+
+  torch::lazy::NodePtr Clone(torch::lazy::OpList operands) const override;
+
+  XlaOpVector Lower(LoweringContext* loctx) const override;
+
+  std::string ToString() const override;
+
+ private:
+  std::string call_target_name_;
+  std::string opaque_;
+};
+
+}  // namespace torch_xla
+
+#endif  // XLA_TORCH_XLA_CSRC_OPS_CUDA_CUSTOM_CALL_H_

torch_xla/csrc/ops/xla_ops.cpp

@@ -37,5 +37,6 @@ const OpKindWrapper xla_unselect("xla::unselect");
 const OpKindWrapper xla_update_slice("xla::update_slice");
 const OpKindWrapper xla_custom_sharding("xla::custom_sharding");
 const OpKindWrapper xla_tpu_custom_call("xla::tpu_custom_call");
+const OpKindWrapper xla_cuda_custom_call("xla::cuda_custom_call");
 
 }  // namespace torch_xla

torch_xla/csrc/ops/xla_ops.h

@@ -61,6 +61,7 @@ extern const OpKindWrapper xla_unselect;
 extern const OpKindWrapper xla_update_slice;
 extern const OpKindWrapper xla_custom_sharding;
 extern const OpKindWrapper xla_tpu_custom_call;
+extern const OpKindWrapper xla_cuda_custom_call;
 
 }  // namespace torch_xla
 

torch_xla/csrc/tensor_methods.cpp

@@ -38,6 +38,7 @@
 #include "torch_xla/csrc/ops/convolution_backward_overrideable.h"
 #include "torch_xla/csrc/ops/convolution_overrideable.h"
 #include "torch_xla/csrc/ops/count_nonzero.h"
+#include "torch_xla/csrc/ops/cuda_custom_call.h"
 #include "torch_xla/csrc/ops/cumprod.h"
 #include "torch_xla/csrc/ops/cumsum.h"
 #include "torch_xla/csrc/ops/dequant_tensor.h"
@@ -555,6 +556,19 @@ void tpu_custom_call_(XLATensorPtr& output,
       values, output->shape().get(), payload));
 }
 
+std::vector<XLATensorPtr> cuda_custom_call(
+    const std::vector<XLATensorPtr>& inputs, int num_outputs,
+    const xla::Shape& output_shape, const std::string& call_target_name,
+    const std::string& opaque) {
+  std::vector<torch::lazy::Value> values;
+  for (const auto& input : inputs) {
+    values.push_back(input->GetIrValue());
+  }
+  auto node = torch::lazy::MakeNode<CudaCustomCall>(
+      values, num_outputs, output_shape, call_target_name, opaque);
+  return inputs[0]->MakeOutputTensors(node, /*inherit_logical_type=*/false);
+}
+
 XLATensorPtr get_dimensions_size(const XLATensorPtr& input,
                                  std::vector<int64_t> dimensions) {
   return input->CreateFrom(torch::lazy::MakeNode<GetDimensionsSize>(

torch_xla/csrc/tensor_methods.h

@@ -90,6 +90,11 @@ void tpu_custom_call_(XLATensorPtr& output,
                       const std::vector<XLATensorPtr>& inputs,
                       const std::string& payload);
 
+std::vector<XLATensorPtr> cuda_custom_call(
+    const std::vector<XLATensorPtr>& inputs, int num_outputs,
+    const xla::Shape& output_shape, const std::string& call_target_name,
+    const std::string& opaque);
+
 XLATensorPtr get_dimensions_size(const XLATensorPtr& input,
                                  std::vector<int64_t> dimensions);
 

torch_xla/csrc/tensor_util.cpp

@@ -63,6 +63,20 @@ struct Caster<tsl::bfloat16> {
   }
 };
 template <>
+struct Caster<at::Float8_e4m3fn> {
+  template <typename D>
+  D cast(const at::Float8_e4m3fn& value) const {
+    return static_cast<D>(static_cast<float>(value));
+  }
+};
+template <>
+struct Caster<tsl::float8_e4m3fn> {
+  template <typename D>
+  D cast(const tsl::float8_e4m3fn& value) const {
+    return static_cast<D>(static_cast<float>(value));
+  }
+};
+template <>
 struct Caster<at::Half> {
   template <typename D>
   D cast(const at::Half& value) const {
@@ -181,6 +195,15 @@ struct NeedCast<tsl::bfloat16> {
   static constexpr bool value = true;
 };
 template <>
+struct NeedCast<tsl::float8_e4m3fn> {
+  static constexpr bool value = true;
+};
+
+template <>
+struct NeedCast<at::Float8_e4m3fn> {
+  static constexpr bool value = true;
+};
+template <>
 struct NeedCast<at::BFloat16> {
   static constexpr bool value = true;
 };
@@ -248,6 +271,18 @@ void CopyData<tsl::bfloat16, at::BFloat16>(tsl::bfloat16* dest,
                                            int64_t n, const CopyCasted&) {
   CheckedMemcpy<tsl::bfloat16, at::BFloat16>(dest, source, n);
 }
+template <>
+void CopyData<at::Float8_e4m3fn, tsl::float8_e4m3fn>(
+    at::Float8_e4m3fn* dest, const tsl::float8_e4m3fn* source, int64_t n,
+    const CopyCasted&) {
+  CheckedMemcpy<at::Float8_e4m3fn, tsl::float8_e4m3fn>(dest, source, n);
+}
+template <>
+void CopyData<tsl::float8_e4m3fn, at::Float8_e4m3fn>(
+    tsl::float8_e4m3fn* dest, const at::Float8_e4m3fn* source, int64_t n,
+    const CopyCasted&) {
+  CheckedMemcpy<tsl::float8_e4m3fn, at::Float8_e4m3fn>(dest, source, n);
+}
 
 std::vector<int64_t> GetIterationDimensions(const xla::Shape& shape) {
   // We want to favor the most minor dimension as core iteration dimension, as
@@ -458,6 +493,10 @@ void TensorToBufferSType(const at::Tensor& tensor, const xla::Shape& dest_shape,
       TensorToBuffer<SType, xla::complex128>(tensor, dest_shape, dest_buffer,
                                              dest_buffer_size, device);
       break;
+    case xla::PrimitiveType::F8E4M3FN:
+      TensorToBuffer<SType, tsl::float8_e4m3fn>(tensor, dest_shape, dest_buffer,
+                                                dest_buffer_size, device);
+      break;
     default:
       XLA_ERROR() << "Destination shape type not supported: " << dest_shape;
   }
@@ -537,6 +576,9 @@ at::Tensor XlaLiteralToTensorHelper(const xla::Literal& literal,
     case at::ScalarType::ComplexDouble:
       return XlaLiteralToTensor<SType, c10::complex<double>>(literal,
                                                              dest_element_type);
+    case at::ScalarType::Float8_e4m3fn:
+      return XlaLiteralToTensor<SType, at::Float8_e4m3fn>(literal,
+                                                          dest_element_type);
     default:
       XLA_ERROR() << "Unsupported scalar type: " << dest_element_type;
   }
@@ -597,6 +639,10 @@ void PopulateTensorBuffer(const at::Tensor& tensor,
       TensorToBufferSType<c10::complex<double>>(tensor, dest_shape, dest_buffer,
                                                 dest_buffer_size, device);
       break;
+    case at::ScalarType::Float8_e4m3fn:
+      TensorToBufferSType<at::Float8_e4m3fn>(tensor, dest_shape, dest_buffer,
+                                             dest_buffer_size, device);
+      break;
     default:
       XLA_ERROR() << "Tensor type not supported: " << tensor.type();
   }
@@ -648,6 +694,9 @@ at::Tensor MakeTensorFromXlaLiteral(const xla::Literal& literal,
     case xla::PrimitiveType::C128:
       return XlaLiteralToTensorHelper<xla::complex128>(literal,
                                                        dest_element_type);
+    case xla::PrimitiveType::F8E4M3FN:
+      return XlaLiteralToTensorHelper<tsl::float8_e4m3fn>(literal,
+                                                          dest_element_type);
     default:
       XLA_ERROR() << "Unsupported literal type: " << literal.shape();
   }