pybind.h 3.59 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*************************************************************************
 * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#define PYBIND11_DETAILED_ERROR_MESSAGES  // TODO remove

#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_H_
#define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_H_
#include <pybind11/detail/common.h>
#include <pybind11/functional.h>
#include <pybind11/pybind11.h>
#include <torch/torch.h>

#include "common.h"
#include "transformer_engine/transformer_engine.h"

namespace transformer_engine::pytorch {

extern PyTypeObject *Float8TensorPythonClass;
extern PyTypeObject *Float8TensorBasePythonClass;
extern PyTypeObject *Float8QuantizerClass;
24
extern PyTypeObject *Float8CurrentScalingQuantizerClass;
25
26
27
extern PyTypeObject *MXFP8TensorPythonClass;
extern PyTypeObject *MXFP8TensorBasePythonClass;
extern PyTypeObject *MXFP8QuantizerClass;
28
29
30
extern PyTypeObject *Float8BlockwiseQTensorPythonClass;
extern PyTypeObject *Float8BlockwiseQTensorBasePythonClass;
extern PyTypeObject *Float8BlockwiseQuantizerClass;
31
32
33
34
35
36
37
38
39

void init_extension();

void init_float8_extension();

void init_mxfp8_extension();

namespace detail {

40
41
42
43
44
inline bool IsFloat8Quantizers(PyObject *obj) { return Py_TYPE(obj) == Float8QuantizerClass; }

inline bool IsFloat8CurrentScalingQuantizers(PyObject *obj) {
  return Py_TYPE(obj) == Float8CurrentScalingQuantizerClass;
}
45
46
47
48
49

inline bool IsFloat8Tensor(PyObject *obj) {
  return Py_TYPE(obj) == Float8TensorPythonClass || Py_TYPE(obj) == Float8TensorBasePythonClass;
}

50
inline bool IsMXFP8Quantizers(PyObject *obj) { return Py_TYPE(obj) == MXFP8QuantizerClass; }
51
52
53
54
55

inline bool IsMXFP8Tensor(PyObject *obj) {
  return Py_TYPE(obj) == MXFP8TensorPythonClass || Py_TYPE(obj) == MXFP8TensorBasePythonClass;
}

56
57
58
59
60
61
62
63
64
inline bool IsFloat8BlockwiseQuantizers(PyObject *obj) {
  return Py_TYPE(obj) == Float8BlockwiseQuantizerClass;
}

inline bool IsFloat8BlockwiseQTensor(PyObject *obj) {
  return Py_TYPE(obj) == Float8BlockwiseQTensorPythonClass ||
         Py_TYPE(obj) == Float8BlockwiseQTensorBasePythonClass;
}

65
66
67
68
69
70
71
72
73
74
75
TensorWrapper NVTETensorFromFloat8Tensor(py::handle tensor, Quantizer *quantizer);

template <typename T>
std::unique_ptr<Quantizer> CreateQuantizer(const py::handle quantizer) {
  return std::make_unique<T>(quantizer);
}

TensorWrapper NVTETensorFromMXFP8Tensor(py::handle tensor, Quantizer *quantization_params);

std::unique_ptr<Quantizer> CreateMXFP8Params(const py::handle params);

76
77
78
TensorWrapper NVTETensorFromFloat8BlockwiseQTensor(py::handle tensor,
                                                   Quantizer *quantization_params);

79
80
81
82
83
inline bool IsFloatingPointType(at::ScalarType type) {
  return type == at::kFloat || type == at::kHalf || type == at::kBFloat16;
}

constexpr std::array custom_types_converters = {
84
    std::make_tuple(IsFloat8Tensor, IsFloat8Quantizers, NVTETensorFromFloat8Tensor,
85
                    CreateQuantizer<Float8Quantizer>),
86
87
88
    std::make_tuple(IsFloat8Tensor, IsFloat8CurrentScalingQuantizers, NVTETensorFromFloat8Tensor,
                    CreateQuantizer<Float8CurrentScalingQuantizer>),
    std::make_tuple(IsMXFP8Tensor, IsMXFP8Quantizers, NVTETensorFromMXFP8Tensor,
89
90
91
                    CreateQuantizer<MXFP8Quantizer>),
    std::make_tuple(IsFloat8BlockwiseQTensor, IsFloat8BlockwiseQuantizers,
                    NVTETensorFromFloat8BlockwiseQTensor, CreateQuantizer<Float8BlockQuantizer>)};
92
93
94
95
96
97

}  // namespace detail

}  // namespace transformer_engine::pytorch

#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_PYBIND_H_