Skip to content

Commit 9dac78e

Browse files
vthumbe1503pre-commit-ci[bot]greptile-apps[bot]ksivaman
authored
CPU Overhead Optimizations (#2559)
* add all the optimizations Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * requires_grad optimization Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test if commenting out requires_grad works Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix minor bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix ci Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * missed a bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * Update transformer_engine/pytorch/csrc/quantizer.cpp Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * fix some bugs pointed to by copilot Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * linting error Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix the error Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * get rid of the change Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix the transpose shape bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * minor linter fix Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix linting error Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address copilot review comment regarding error check when both data and transpose are None Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix linting errors Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * missed a merge conflict Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * final optimizations Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ci error Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comment from greptile Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comment + stride optimization Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address linter issue Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix ci bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * another optimization to do at::native::empty_cuda directly instead of at::empty Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * cleanups Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * better solution for device Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * enum to int cache Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused function Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * Update transformer_engine/pytorch/tensor/float8_blockwise_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * index instead of device bug Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ci: Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * debug quantized tensor fix Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert cudnnt front end change Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
1 parent 3275e1a commit 9dac78e

23 files changed

Lines changed: 622 additions & 218 deletions

‎tests/pytorch/attention/test_attention.py‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
scaled_init_method_normal,
4545
)
4646
from transformer_engine.pytorch.utils import get_cudnn_version
47+
from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx
4748
import transformer_engine_torch as tex
4849
from transformer_engine.pytorch.quantized_tensor import (
4950
Quantizer,
@@ -2581,12 +2582,12 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
25812582
_2X_ACC_DGRAD = False
25822583
_2X_ACC_WGRAD = False
25832584

2584-
META_QKV = tex.FP8FwdTensors.GEMM1_OUTPUT
2585-
META_DQKV = tex.FP8BwdTensors.GRAD_OUTPUT1
2586-
META_O = tex.FP8FwdTensors.GEMM2_INPUT
2587-
META_DO = tex.FP8BwdTensors.GRAD_INPUT2
2588-
META_S = tex.FP8FwdTensors.GEMM3_OUTPUT
2589-
META_DP = tex.FP8BwdTensors.GRAD_INPUT3
2585+
META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT
2586+
META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1
2587+
META_O = FP8FwdTensorIdx.GEMM2_INPUT
2588+
META_DO = FP8BwdTensorIdx.GRAD_INPUT2
2589+
META_S = FP8FwdTensorIdx.GEMM3_OUTPUT
2590+
META_DP = FP8BwdTensorIdx.GRAD_INPUT3
25902591

25912592

25922593
class _custom_mha_fp8(torch.autograd.Function):
@@ -2614,14 +2615,14 @@ def forward(
26142615
d = in_features // h
26152616
b = cu_seqlens.numel() - 1
26162617

2617-
input_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
2618-
qkv_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM2_INPUT]
2619-
qkv_weight_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
2620-
o_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
2621-
dO_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
2622-
dQKV_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
2623-
s_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT2]
2624-
dP_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT3]
2618+
input_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT]
2619+
qkv_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_INPUT]
2620+
qkv_weight_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
2621+
o_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT]
2622+
dO_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1]
2623+
dQKV_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1]
2624+
s_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT2]
2625+
dP_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT3]
26252626

26262627
inp_fp8 = input_quantizer(inp)
26272628

‎tests/pytorch/test_custom_recipe.py‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import transformer_engine.pytorch as te
99
import transformer_engine_torch as tex
1010
from transformer_engine.common import recipe
11+
from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx
1112
from transformer_engine.pytorch import (
1213
autocast,
1314
Linear,
@@ -169,11 +170,11 @@ def test_custom_recipe_matches_current_scaling():
169170
with autocast(enabled=True, recipe=ref_recipe):
170171
out_ref = model_ref(inp_ref)
171172
# Assert dtypes for reference quantizers: HYBRID = E4M3 (fwd), E5M2 (bwd)
172-
ref_fwd_in = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
173-
ref_fwd_w = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
174-
ref_fwd_out = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
175-
ref_bwd_go = model_ref.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
176-
ref_bwd_gi = model_ref.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
173+
ref_fwd_in = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT]
174+
ref_fwd_w = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
175+
ref_fwd_out = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT]
176+
ref_bwd_go = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1]
177+
ref_bwd_gi = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1]
177178
assert ref_fwd_in.dtype == tex.DType.kFloat8E4M3
178179
assert ref_fwd_w.dtype == tex.DType.kFloat8E4M3
179180
assert ref_fwd_out.dtype == tex.DType.kFloat8E4M3
@@ -200,11 +201,11 @@ def quantizer_factory(role):
200201
with autocast(enabled=True, recipe=custom_recipe):
201202
out_custom = model_custom(inp_custom)
202203
# Assert dtypes for custom quantizers match reference mapping
203-
cus_fwd_in = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
204-
cus_fwd_w = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
205-
cus_fwd_out = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
206-
cus_bwd_go = model_custom.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
207-
cus_bwd_gi = model_custom.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
204+
cus_fwd_in = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT]
205+
cus_fwd_w = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
206+
cus_fwd_out = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT]
207+
cus_bwd_go = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1]
208+
cus_bwd_gi = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1]
208209
assert cus_fwd_in.dtype == tex.DType.kFloat8E4M3
209210
assert cus_fwd_w.dtype == tex.DType.kFloat8E4M3
210211
assert cus_fwd_out.dtype == tex.DType.kFloat8E4M3

‎transformer_engine/common/gemm/cublaslt_gemm.cu‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,10 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla
120120
// Set conditions for MXFP8 and NVFP4 gemm execution.
121121
const auto nvfp4 = is_nvfp_scaling(A.scaling_mode) && is_nvfp_scaling(B.scaling_mode);
122122
const auto mxfp8 = !nvfp4 && is_mxfp_scaling(A.scaling_mode) && is_mxfp_scaling(B.scaling_mode);
123+
int is_nvte_non_tn_fp8_gemm_supported = 0; // needed only for per tensor scaling
124+
if (is_tensor_scaling(A.scaling_mode) || is_tensor_scaling(B.scaling_mode)) {
125+
is_nvte_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported();
126+
}
123127

124128
// Configure A matrix
125129
if (is_tensor_scaling(A.scaling_mode)) {
@@ -129,7 +133,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla
129133
ret.Atype = A.data.dtype;
130134
ret.A_scale_inv = A.scale_inv.dptr;
131135
ret.lda = is_A_transposed ? k : m;
132-
if (!nvte_is_non_tn_fp8_gemm_supported() && !is_A_transposed) {
136+
if (!is_nvte_non_tn_fp8_gemm_supported && !is_A_transposed) {
133137
// Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data.
134138
if (A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype)) {
135139
ret.A = A.columnwise_data.dptr;
@@ -140,7 +144,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla
140144
} else {
141145
NVTE_CHECK(!is_fp8_dtype(ret.Atype), "Input A is missing column-wise usage");
142146
}
143-
} else if (nvte_is_non_tn_fp8_gemm_supported() && !A.has_data()) {
147+
} else if (is_nvte_non_tn_fp8_gemm_supported && !A.has_data()) {
144148
// Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed
145149
// data with the mirrored transpose-flag if we don't have row-wise data.
146150
NVTE_CHECK(A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype),
@@ -220,7 +224,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla
220224
ret.Btype = B.data.dtype;
221225
ret.B_scale_inv = B.scale_inv.dptr;
222226
ret.ldb = is_B_transposed ? n : k;
223-
if (!nvte_is_non_tn_fp8_gemm_supported() && is_B_transposed) {
227+
if (!is_nvte_non_tn_fp8_gemm_supported && is_B_transposed) {
224228
// Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data.
225229
if (B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype)) {
226230
ret.B = B.columnwise_data.dptr;
@@ -231,7 +235,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla
231235
} else {
232236
NVTE_CHECK(!is_fp8_dtype(ret.Btype), "Input B is missing column-wise usage");
233237
}
234-
} else if (nvte_is_non_tn_fp8_gemm_supported() && !B.has_data()) {
238+
} else if (is_nvte_non_tn_fp8_gemm_supported && !B.has_data()) {
235239
// Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed
236240
// data with the mirrored transpose-flag if we don't have row-wise data.
237241
NVTE_CHECK(B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype),

‎transformer_engine/common/util/cuda_driver.h‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@
99

1010
#include <cuda.h>
1111

12+
#include <mutex>
1213
#include <string>
14+
#include <unordered_map>
1315

1416
#include "../common.h"
1517
#include "../util/string.h"
@@ -29,13 +31,30 @@ void *get_symbol(const char *symbol, int cuda_version = 12010);
2931
* without GPUs. Indirect function calls into a lazily-initialized
3032
* library ensures we are accessing the correct version.
3133
*
34+
* Symbol pointers are cached to avoid repeated lookups.
35+
*
3236
* \param[in] symbol Function name
3337
* \param[in] args Function arguments
3438
*/
3539
template <typename... ArgTs>
3640
inline CUresult call(const char *symbol, ArgTs... args) {
3741
using FuncT = CUresult(ArgTs...);
38-
FuncT *func = reinterpret_cast<FuncT *>(get_symbol(symbol));
42+
43+
static std::unordered_map<std::string, void *> symbol_cache;
44+
static std::mutex cache_mutex;
45+
FuncT *func;
46+
47+
{
48+
std::lock_guard<std::mutex> lock(cache_mutex);
49+
auto it = symbol_cache.find(symbol);
50+
if (it == symbol_cache.end()) {
51+
void *ptr = get_symbol(symbol);
52+
symbol_cache[symbol] = ptr;
53+
func = reinterpret_cast<FuncT *>(ptr);
54+
} else {
55+
func = reinterpret_cast<FuncT *>(it->second);
56+
}
57+
}
3958
return (*func)(args...);
4059
}
4160

‎transformer_engine/debug/pytorch/debug_quantization.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -697,3 +697,12 @@ def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None
697697
raise RuntimeError(
698698
"Cannot recreate columnwise tensor from rowwise tensor is debug mode."
699699
)
700+
701+
@property
702+
def device(self):
703+
"""Return the device of the tensor. Define this to avoid expensive PyObject lookups."""
704+
if self.rowwise_gemm_tensor is not None:
705+
return self.rowwise_gemm_tensor.device
706+
if self.columnwise_gemm_tensor is not None:
707+
return self.columnwise_gemm_tensor.device
708+
raise RuntimeError("DebugQuantizedTensor has no data!")

‎transformer_engine/pytorch/constants.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# See LICENSE for license information.
44

55
"""Enums for e2e transformer"""
6+
from types import SimpleNamespace
67
import torch
78
import torch.distributed
89
import transformer_engine_torch as tex
@@ -40,6 +41,25 @@
4041
tex.DType.kBFloat16: torch.bfloat16,
4142
}
4243

44+
# Cache enum -> int conversions to avoid repeated PyObject lookups.
45+
FP8FwdTensorIdx = SimpleNamespace(
46+
GEMM1_INPUT=int(tex.FP8FwdTensors.GEMM1_INPUT),
47+
GEMM1_WEIGHT=int(tex.FP8FwdTensors.GEMM1_WEIGHT),
48+
GEMM1_OUTPUT=int(tex.FP8FwdTensors.GEMM1_OUTPUT),
49+
GEMM2_INPUT=int(tex.FP8FwdTensors.GEMM2_INPUT),
50+
GEMM2_WEIGHT=int(tex.FP8FwdTensors.GEMM2_WEIGHT),
51+
GEMM2_OUTPUT=int(tex.FP8FwdTensors.GEMM2_OUTPUT),
52+
GEMM3_OUTPUT=int(tex.FP8FwdTensors.GEMM3_OUTPUT),
53+
)
54+
FP8BwdTensorIdx = SimpleNamespace(
55+
GRAD_INPUT1=int(tex.FP8BwdTensors.GRAD_INPUT1),
56+
GRAD_INPUT2=int(tex.FP8BwdTensors.GRAD_INPUT2),
57+
GRAD_INPUT3=int(tex.FP8BwdTensors.GRAD_INPUT3),
58+
GRAD_OUTPUT1=int(tex.FP8BwdTensors.GRAD_OUTPUT1),
59+
GRAD_OUTPUT2=int(tex.FP8BwdTensors.GRAD_OUTPUT2),
60+
GRAD_OUTPUT3=int(tex.FP8BwdTensors.GRAD_OUTPUT3),
61+
)
62+
4363
AttnMaskTypes = (
4464
"no_mask",
4565
"padding",

‎transformer_engine/pytorch/cpp_extensions/fused_attn.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
NVTE_Fused_Attn_Backend,
1717
)
1818
from ..quantized_tensor import Quantizer
19+
from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx
1920

2021

2122
__all__ = [
@@ -103,12 +104,12 @@
103104
BACKEND_F16m512_FP8_THREADS_PER_CTA = 128
104105
BACKEND_F16arb_ELTS_PER_THREADS = 16
105106

106-
META_QKV = tex.FP8FwdTensors.GEMM1_OUTPUT
107-
META_DQKV = tex.FP8BwdTensors.GRAD_OUTPUT1
108-
META_O = tex.FP8FwdTensors.GEMM2_INPUT
109-
META_DO = tex.FP8BwdTensors.GRAD_INPUT2
110-
META_S = tex.FP8FwdTensors.GEMM3_OUTPUT
111-
META_DP = tex.FP8BwdTensors.GRAD_INPUT3
107+
META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT
108+
META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1
109+
META_O = FP8FwdTensorIdx.GEMM2_INPUT
110+
META_DO = FP8BwdTensorIdx.GRAD_INPUT2
111+
META_S = FP8FwdTensorIdx.GEMM3_OUTPUT
112+
META_DP = FP8BwdTensorIdx.GRAD_INPUT3
112113

113114

114115
def fused_attn_fwd(

‎transformer_engine/pytorch/cpp_extensions/gemm.py‎

Lines changed: 2 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -67,28 +67,6 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float:
6767
return 0.0
6868

6969

70-
def get_tensor_device(tensor: torch.Tensor) -> int:
71-
"""
72-
Returns tensor device as an integer.
73-
74-
This method is used because checking instances of
75-
QuantizedTensor or Storage incurs more CPU overhead.
76-
The order of attributes checked is important to also
77-
minimize overhead.
78-
"""
79-
if hasattr(tensor, "device"):
80-
return tensor.device.index
81-
if hasattr(tensor, "_rowwise_data") and tensor._rowwise_data is not None:
82-
return tensor._rowwise_data.device.index
83-
if hasattr(tensor, "_columnwise_data") and tensor._columnwise_data is not None:
84-
return tensor._columnwise_data.device.index
85-
if hasattr(tensor, "_data") and tensor._data is not None:
86-
return tensor._data.device.index
87-
if hasattr(tensor, "_transpose") and tensor._transpose is not None:
88-
return tensor._transpose.device.index
89-
return torch.cuda.current_device()
90-
91-
9270
def general_gemm(
9371
A: torch.Tensor,
9472
B: torch.Tensor,
@@ -117,7 +95,7 @@ def general_gemm(
11795

11896
alpha = validate_gemm_scale(alpha, True)
11997
beta = validate_gemm_scale(beta, accumulate)
120-
workspace = get_cublas_workspace(get_tensor_device(A), ub is not None, False)
98+
workspace = get_cublas_workspace(A.device.index, ub is not None, False)
12199

122100
if ub_type is not None:
123101
assert ub is not None, (
@@ -235,7 +213,7 @@ def general_grouped_gemm(
235213
out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype
236214

237215
sm_count = get_sm_count()
238-
workspaces = get_cublas_workspace(get_tensor_device(A[0]), False, True)
216+
workspaces = get_cublas_workspace(A[0].device.index, False, True)
239217

240218
if grad and use_bias:
241219
grad_bias = [

‎transformer_engine/pytorch/csrc/extensions/pybind.cpp‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,10 @@ PyTypeObject *Float8BlockwiseQuantizerClass = nullptr;
3535
PyTypeObject *NVFP4TensorPythonClass = nullptr;
3636
PyTypeObject *NVFP4TensorStoragePythonClass = nullptr;
3737
PyTypeObject *NVFP4QuantizerClass = nullptr;
38+
std::once_flag extension_init_flag;
3839
PyTypeObject *GroupedTensorStoragePythonClass = nullptr;
3940

4041
void init_float8_extension() {
41-
if (Float8TensorPythonClass) return;
4242
auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.float8_tensor");
4343
Float8QuantizerClass =
4444
reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "Float8Quantizer"));
@@ -55,7 +55,6 @@ void init_float8_extension() {
5555
}
5656

5757
void init_mxfp8_extension() {
58-
if (MXFP8TensorPythonClass) return;
5958
auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.mxfp8_tensor");
6059
MXFP8QuantizerClass =
6160
reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "MXFP8Quantizer"));
@@ -70,7 +69,6 @@ void init_mxfp8_extension() {
7069
}
7170

7271
void init_float8blockwise_extension() {
73-
if (Float8BlockwiseQTensorStoragePythonClass) return;
7472
auto fp8_module =
7573
py::module_::import("transformer_engine.pytorch.tensor.float8_blockwise_tensor");
7674
auto fp8_base_module = py::module_::import(
@@ -91,7 +89,6 @@ void init_float8blockwise_extension() {
9189
}
9290

9391
void init_nvfp4_extensions() {
94-
if (NVFP4TensorPythonClass) return;
9592
auto nvfp4_module = py::module_::import("transformer_engine.pytorch.tensor.nvfp4_tensor");
9693
NVFP4QuantizerClass = reinterpret_cast<PyTypeObject *>(
9794
PyObject_GetAttrString(nvfp4_module.ptr(), "NVFP4Quantizer"));
@@ -116,11 +113,13 @@ void init_grouped_tensor_extension() {
116113
}
117114

118115
void init_extension() {
119-
init_float8_extension();
120-
init_mxfp8_extension();
121-
init_float8blockwise_extension();
122-
init_nvfp4_extensions();
123-
init_grouped_tensor_extension();
116+
std::call_once(extension_init_flag, []() {
117+
init_float8_extension();
118+
init_mxfp8_extension();
119+
init_float8blockwise_extension();
120+
init_nvfp4_extensions();
121+
init_grouped_tensor_extension();
122+
});
124123
}
125124

126125
} // namespace transformer_engine::pytorch

0 commit comments

Comments
 (0)