CPU Overhead Optimizations

[vthumbe1503]

Description

CPU overhead optimizations

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

Python Optimizations

• TE Pybinded Enums like tex.FP8FwdTensors.GEMM1_INPUT are casted to int in each and every forward pass. Now we are caching the integer values in a constants file and using that instead.
• Getting tensor device in the helper function went through expensive tensor.device even for Quantized Tensor. Now we have the device declared as propery of QuantizedTensor, so it doesnt go through PyObject Lookup
• Defining device shape and is_cuda attributes as Quantized Tensor properties (since they are easy enough to compute in python) and it avoid the expensive PyObject Lookup
• Defining requires_grad and dtype as properties of base QuantizedTensor class. Here we cache the properties when values are being set to avoid the expensive PyObject Lookup. We still need to make sure setter goes through Pybind C++. For instance torch autograd engine in C++ needs to be aware of requires_grad changes.
• dtype of our Custom QuantizedTensor can change when we go through x.data = new_tensor. And so we make sure dtype is cached appropriately by defining appropriate _get_data and _set_data for the data property of QuantizedTensor

C++ Optimizations

• Caching symbol lookups in libcuda.so for driver calls like cuCtxGetCurrent, so we dont lookup the symbol in each and every forward/backward call.
• Caching nvte_non_tn_fp8_gemm_supported() function call
• Faster py object call without cxa_demangle to construct QuantizedTensor classes in C++
• Reduce Python work in QuantizedTensor object creation(calculating stride from shape and getting current cuda device can be done in C++ instead Python Constructor).

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[vthumbe1503]

/te-ci L1 pytorch

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

Greptile Summary

This PR implements comprehensive CPU overhead optimizations across Python and C++ layers of TransformerEngine, targeting hotspots identified through profiling.

Key Optimizations

Python Layer:

• Cached pybind11 enum values (FP8FwdTensorIdx, FP8BwdTensorIdx) as integers to avoid repeated enum-to-int conversions on every forward/backward pass
• Added cached _dtype and _requires_grad properties to QuantizedTensor base class with lazy initialization fallback to avoid expensive PyObject lookups
• Added device, shape, and is_cuda properties to all quantized tensor types to bypass parent class attribute resolution
• Modified QuantizedTensor.__new__ to accept pre-computed stride parameter, allowing C++ to compute strides without calling back into Python

C++ Layer:

• Replaced pybind11 named argument function calls with direct PyObject_Call using py::dict and py::tuple for lower overhead in tensor creation paths
• Added thread-safe symbol caching in cuda_driver.h using mutex-protected map to avoid repeated dlsym lookups for CUDA driver API functions
• Cached nvte_is_non_tn_fp8_gemm_supported() result per function invocation instead of calling multiple times for A and B matrix configurations
• Moved stride computation from Python to C++ to eliminate PyObjectVectorCall overhead
• Thread-safe extension initialization using std::call_once

Code Quality

The implementation demonstrates solid engineering practices:

• Proper RAII memory management with py::dict/py::tuple preventing reference leaks
• Thread safety through std::mutex and std::call_once
• Defensive error handling with RuntimeError for edge cases where tensor data is missing
• Fallback logic for cached properties to handle alternate tensor creation paths (unpickling, FSDP, etc.)

All critical issues from previous review rounds have been addressed, including memory leaks, thread safety concerns, and proper initialization of cached attributes.

Confidence Score: 5/5

• This PR is safe to merge with high confidence
• All previously identified critical issues (memory leaks, thread safety, initialization bugs) have been properly addressed. The optimizations are well-implemented with proper error handling, thread safety mechanisms, and defensive programming. The changes maintain backward compatibility while reducing CPU overhead through caching and reduced Python-C++ boundary crossings. Code quality is excellent with RAII memory management, mutex-protected shared state, and comprehensive error handling
• No files require special attention - all changes follow best practices and critical issues have been resolved

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Replaces pybind11 function calls with direct PyObject_Call for tensor creation, computes stride in C++ instead of Python, caches nvte_is_non_tn_fp8_gemm_supported() result. Memory management properly handled with py::dict/py::tuple RAII
transformer_engine/pytorch/quantized_tensor.py	Adds cached _dtype and _requires_grad properties with proper initialization and fallback logic. Adds stride parameter to __new__ to avoid Python call from C++. Implements data property setter to sync dtype cache
transformer_engine/common/util/cuda_driver.h	Adds symbol caching with thread-safe mutex-protected map to avoid repeated libcuda.so symbol lookups
transformer_engine/pytorch/constants.py	Adds FP8FwdTensorIdx and FP8BwdTensorIdx SimpleNamespace objects to cache enum-to-int conversions, avoiding expensive pybind11 enum casts on every forward/backward pass
transformer_engine/common/gemm/cublaslt_gemm.cu	Caches nvte_is_non_tn_fp8_gemm_supported() result at function start to avoid redundant calls for both A and B matrix configurations
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Replaces null-check guards with std::call_once for thread-safe initialization of Python extension classes

_{Last reviewed commit: e52a12d}

[greptile-apps]

_{24 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements CPU-side performance optimizations for FP8 operations by caching frequently accessed attributes and reducing redundant function calls. The optimizations target expensive PyObject attribute lookups on custom tensor types and repeated C++ function calls.

Key Changes:

• Caches requires_grad, dtype, shape, and is_cuda attribute accesses to avoid expensive PyObject lookups on custom tensors
• Reorders attribute checks in get_tensor_device() to prioritize internal quantized tensor attributes
• Makes num_devices static in nvte_is_non_tn_fp8_gemm_supported() to cache device count
• Stores GEMM support check results in local variables to avoid redundant function calls

Critical Issues Found:

• Variable redeclaration error in cublaslt_gemm.cu (line 224) will prevent compilation
• Logic bug in linear.py (line 484) changes FP8 state management from OR logic to AND logic, breaking functionality when bias is None or doesn't require grad

Confidence Score: 0/5

• This PR cannot be merged due to compilation error and critical logic bug
• Two critical issues prevent merging: (1) C++ compilation will fail due to variable redeclaration at line 224 of cublaslt_gemm.cu, and (2) logic bug at line 484 of linear.py breaks FP8 state management by requiring all three tensors to have requires_grad=True instead of any one of them
• Pay close attention to transformer_engine/common/gemm/cublaslt_gemm.cu (compilation error) and transformer_engine/pytorch/module/linear.py (logic bug)

Important Files Changed

File Analysis

Filename	Score	Overview
transformer_engine/common/gemm/cublaslt_gemm.cu	1/5	Caches function call result to reduce overhead, but contains variable redeclaration error that will cause compilation failure
transformer_engine/common/transformer_engine.cpp	5/5	Makes num_devices static to avoid redundant calls to cuda::num_devices() - valid optimization
transformer_engine/pytorch/module/linear.py	0/5	Caches requires_grad checks for performance, but contains critical logic bug at line 484 that changes FP8 state management behavior

Sequence Diagram

sequenceDiagram
    participant User as User Code
    participant Linear as Linear Module
    participant Quantizer as Quantizer/QuantizedTensor
    participant GEMM as GEMM Operations
    participant CPP as C++ Extensions

    Note over Linear,CPP: Performance Optimization Flow
    
    User->>Linear: forward(input, weight, bias)
    
    Note over Linear: Cache requires_grad checks
    Linear->>Linear: inp_requires_grad = inp.requires_grad<br/>weight_requires_grad = weight.requires_grad<br/>bias_requires_grad = bias.requires_grad
    
    Linear->>Quantizer: Check if quantized tensor
    alt QuantizedTensor
        Note over Quantizer: Use cached dtype property
        Quantizer->>Quantizer: return self._dtype
        Note over Quantizer: Use cached shape/is_cuda
        Quantizer->>Quantizer: return self._data.shape
    else Regular Tensor
        Quantizer->>Linear: Standard attribute access
    end
    
    Linear->>CPP: get_tensor_device(tensor)
    Note over CPP: Reordered attribute checks
    CPP->>CPP: Check _rowwise_data first<br/>Check _columnwise_data<br/>Check device last
    CPP-->>Linear: device_index
    
    Linear->>GEMM: Configure GEMM parameters
    Note over GEMM: Cache nvte_is_non_tn_fp8_gemm_supported
    GEMM->>CPP: nvte_is_non_tn_fp8_gemm_supported()
    Note over CPP: Static num_devices cached
    CPP-->>GEMM: support_flag
    GEMM->>GEMM: Store in local variable
    
    GEMM->>GEMM: Execute optimized GEMM
    GEMM-->>Linear: output
    
    Note over Linear: FP8 State Management
    alt FP8 enabled and requires_grad check
        Linear->>Linear: Update FP8 tensors<br/>based on cached flags
    end
    
    Linear-->>User: output

Loading

[greptile-apps]

Additional Comments (2)

transformer_engine/common/gemm/cublaslt_gemm.cu
variable redeclared in same scope - already declared at line 132

    // int is_nvte_non_tn_fp8_gemm_supported already declared at line 132

---

transformer_engine/pytorch/module/linear.py
logic change from original requires_grad(inp, weight, bias) which returns True if ANY tensor requires grad. New code requires ALL THREE to be True, breaking FP8 state management when bias is None or doesn't require grad

            if ctx.fp8 and (inp_requires_grad or weight_requires_grad or bias_requires_grad):

[vthumbe1503]

/te-ci L1 pytorch

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

_{8 files reviewed, 8 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{13 files reviewed, 13 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{19 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch

[greptile-apps]

_{23 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[vthumbe1503]

/te-ci pytorch

[greptile-apps]

_{23 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch

[greptile-apps]

Additional Comments (1)

transformer_engine/pytorch/tensor/float8_tensor.py, line 1019
cached _dtype not updated when copying Float8Tensor with different dtype

when copying from one Float8Tensor to another with different dtype (line 1004 condition), the code creates a dummy tensor with the new dtype and sets it using super(Float8Tensor, type(self)).data.__set__(self, dummy_tensor) (line 1019). this bypasses QuantizedTensor._set_data() which updates the cached _dtype attribute

result: cached _dtype becomes stale and won't match the actual tensor's dtype

add after line 1019:

self._dtype = tensor.dtype

[ksivaman]

/te-ci L0 L1