Implementation:NVIDIA TransformerEngine Vectorized Pointwise
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Provides vectorized memory access utilities and pointwise CUDA kernel templates for efficient element-wise operations with optional FP8 quantization and activation function fusion.
Description
vectorized_pointwise.h implements a hierarchy of vectorized access classes that enable 128-bit (16-byte) coalesced memory access patterns:
- VectorizedStorage<DType, n>: Uses a union to store
nelements ofDTypeas a single wider load typeLType. - VectorizedAccessor<DType, nvec, aligned>: Handles alignment detection and safe boundary loads/stores, supporting both aligned and unaligned access patterns.
- VectorizedLoader: Specializes for read-only access with
load()method. - VectorizedStorer: Specializes for write access with
store()method.
The file also contains kernel templates for pointwise operations that fuse activation functions (GeLU, SiLU, etc.) with type casting and FP8 quantization, using the vectorized access pattern to maximize memory bandwidth utilization.
Usage
Use when implementing custom pointwise CUDA kernels that need optimal memory bandwidth. This is the foundation for all activation, cast, and quantization kernels in TransformerEngine.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/common/util/vectorized_pointwise.h- Lines
- 1--628
Signature
namespace transformer_engine {
template <typename DType, int n>
class VectorizedStorage {
public:
using LType = typename BytesToType<sizeof(DType) * n>::Type;
constexpr static int nvec = n;
// ...
};
template <typename DType, int nvec, bool aligned = false>
class VectorizedAccessor {
public:
using StorageType = VectorizedStorage<typename std::remove_const<DType>::type, nvec>;
// ...
};
template <typename DType, int nvec, bool aligned = false>
class VectorizedLoader : public VectorizedAccessor<const DType, nvec, aligned> {
inline __device__ void load(int id, int nvec);
};
template <typename DType, int nvec, bool aligned = false>
class VectorizedStorer : public VectorizedAccessor<DType, nvec, aligned> {
inline __device__ void store(int id, int nvec);
};
} // namespace transformer_engine
Import
#include "util/vectorized_pointwise.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
ptr |
DType* |
Yes | Pointer to the data buffer |
size |
size_t |
Yes | Number of elements in the buffer |
Outputs
| Name | Type | Description |
|---|---|---|
| vectorized data | StorageType |
Data loaded/stored as wide vectorized types |
Usage Examples
#include "util/vectorized_pointwise.h"
using namespace transformer_engine;
// In a CUDA kernel:
__global__ void my_kernel(const float* input, float* output, size_t n) {
VectorizedLoader<float, 4, true> loader(input, n);
VectorizedStorer<float, 4, true> storer(output, n);
int tid = blockIdx.x * blockDim.x + threadIdx.x;
loader.load(tid, 4);
// process loader.separate()...
storer.store(tid, 4);
}