CUDA implementation of the fused MatMul-Bias-GELU operation. More...

Inheritance diagram for Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >:

[legend]

Collaboration diagram for Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >:

[legend]

Public Types
using	MR = typename CudaDevice::MR
using	UnaryOperationBase = UnaryOperation<DeviceType::Cuda, TInput, TOutput>
Public Types inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cuda, float, float >
using	MR
using	TensorInputType
using	TensorOutputType
Public Types inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
using	DataTypeTraits

Public Member Functions
	CudaMatMulBiasGeluOp ()
	Constructs a new CUDA MatMul-Bias-GELU fused operation with the default device context.
void	backward (const Tensor< TInput, MR > &input, const Tensor< TOutput, MR > &output, const Tensor< TOutput, MR > &output_gradient, const std::vector< std::shared_ptr< ITensor > > &parameters, std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &parameter_gradients, Tensor< TInput, MR > &input_gradient, const std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &output_state) const
	Performs the backward pass of the fused MatMul-Bias-GELU operation.
void	forward (const Tensor< TInput, MR > &input, const std::vector< std::shared_ptr< ITensor > > &parameters, Tensor< TOutput, MR > &output, std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &output_state) const override
	Performs the forward pass of the fused MatMul-Bias-GELU operation on CUDA.
std::string	getName () const override
	Gets the name of this operation.
Public Member Functions inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cuda, float, float >
virtual	~UnaryOperation ()=default
virtual void	backward (const ITensor &input, const ITensor &output_grad, ITensor &input_grad) const=0
	Backward pass: compute gradient wrt input given output gradient.
virtual void	forward (const ITensor &input, ITensor &output) const=0
	Forward pass: compute output = f(input).
Public Member Functions inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
virtual	~Operation ()=default
virtual void	build (const BuildContext &build_context)
	Prepare the operation for a concrete input shape.
virtual void	clearGradients () noexcept
	Clear any cached gradient pointers held by the operation.
virtual TensorDataType	getDataType () const
	Tensor data type for this operation.
virtual DeviceType	getDeviceType () const
	Device type for this operation.
virtual OperationType	getOperationType () const=0
	Operation type identifier.
virtual std::size_t	getStateMemorySize () const
	Returns the number of bytes of state memory allocated by this operation.
virtual bool	isBuilt () const
	Whether build() completed successfully for a concrete input shape.
virtual bool	isEvalMode () const
	Query whether operation is configured for training.
virtual void	setGradients (ITensor weight_grad, ITensor bias_grad)
	Bind module-owned gradient tensors to the operation.
virtual void	setParameters (ITensor weight, ITensor bias)
	Bind module-owned parameter tensors to the operation.
virtual void	setTrainingMode (TrainingMode training_mode)
	Configure operation training-mode behavior.

Static Public Member Functions
static const std::string &	className ()
	Gets the class name of this operation.

Additional Inherited Members
Static Public Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
static constexpr TensorDataType	data_type
static constexpr DeviceType	device_type
Static Protected Member Functions inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cuda, float, float >
static const TensorInputType &	asInputTensor (const ITensor &t)
static TensorOutputType &	asOutputTensor (ITensor &t)
Protected Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
bool	is_built_
TrainingMode	training_mode_

Detailed Description

template<typename TInput = float, typename TOutput = TInput>
requires ValidFloatTensorTypes<TInput, TOutput>
class Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >

CUDA implementation of the fused MatMul-Bias-GELU operation.

This class provides a CUDA-based implementation of a fused operation that combines matrix multiplication, bias addition, and GELU activation in a single operation. Fusing these operations improves performance by reducing memory traffic and kernel launch overhead.

The implementation is optimized for NVIDIA GPUs using cuBLASLt for high-performance computation of the fused operation.

Template Parameters

TPrecision	The data type of the input tensor elements.
TDataType	The data type for computation and output (defaults to the input type).

Member Typedef Documentation

◆ MR

template<typename TInput = float, typename TOutput = TInput>

using Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::MR = typename CudaDevice::MR

◆ UnaryOperationBase

template<typename TInput = float, typename TOutput = TInput>

using Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::UnaryOperationBase = UnaryOperation<DeviceType::Cuda, TInput, TOutput>

Constructor & Destructor Documentation

◆ CudaMatMulBiasGeluOp()

template<typename TInput = float, typename TOutput = TInput>

Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::CudaMatMulBiasGeluOp ( )

inline

Constructs a new CUDA MatMul-Bias-GELU fused operation with the default device context.

Initializes the operation with a CUDA device context (defaults to CUDA:0).

Member Function Documentation

◆ backward()

template<typename TInput = float, typename TOutput = TInput>

void Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::backward	(	const Tensor< TInput, MR > &	input,
		const Tensor< TOutput, MR > &	output,
		const Tensor< TOutput, MR > &	output_gradient,
		const std::vector< std::shared_ptr< ITensor > > &	parameters,
		std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &	parameter_gradients,
		Tensor< TInput, MR > &	input_gradient,
		const std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &	output_state ) const

inline

Performs the backward pass of the fused MatMul-Bias-GELU operation.

Computes gradients with respect to inputs, weights, and biases.

Parameters

input	Input tensor from the forward pass.
output	Output tensor from the forward pass.
output_gradient	Gradient of the loss with respect to the output.
parameters	Parameters tensor from forward pass [weights, bias].
parameter_gradients	Gradients for parameters [d_weights, d_bias].
input_gradient	Gradient of the loss with respect to the input.
properties	Additional attributes for the operation.
output_state	Cache tensors from forward pass.

Here is the call graph for this function:

◆ className()

template<typename TInput = float, typename TOutput = TInput>

const std::string & Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::className ( )

inlinestatic

Gets the class name of this operation.

Returns: const std::string& The class name of the operation.

◆ forward()

template<typename TInput = float, typename TOutput = TInput>

void Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::forward	(	const Tensor< TInput, MR > &	input,
		const std::vector< std::shared_ptr< ITensor > > &	parameters,
		Tensor< TOutput, MR > &	output,
		std::vector< std::shared_ptr< Tensor< TOutput, MR > > > &	output_state ) const

inlineoverride

Performs the forward pass of the fused MatMul-Bias-GELU operation on CUDA.

This method efficiently computes a matrix multiplication followed by a bias addition and GELU activation in a single fused operation. The implementation uses cuBLASLt for optimal performance on NVIDIA GPUs.

Parameters

input	Input tensor of shape [B, S, K], where B is batch size, S is sequence length, and K is the input dimension.
parameters	Vector of parameter tensors where: parameters[0]: Weights tensor of shape [K, N] parameters[1]: Bias tensor of shape [N]
properties	Additional attributes for the operation.
output	Output tensor of shape [B, S, N] to store the result.
output_state	Cache for intermediate results (not used in this operation).

Here is the call graph for this function:

◆ getName()

template<typename TInput = float, typename TOutput = TInput>

std::string Mila::Dnn::Compute::Cuda::MatMulBiasGelu::CudaMatMulBiasGeluOp< TInput, TOutput >::getName ( ) const

inlineoverridevirtual

Gets the name of this operation.

Returns: std::string The name of the operation ("Cuda::MatMulBiasGeluOp").

Implements Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

The documentation for this class was generated from the following file:

/__w/Mila/Mila/Mila/Src/Dnn/Compute/Devices/Cuda/Operations/Linear/CudaLinearGeluOp.ixx

Public Types

Public Member Functions

Static Public Member Functions

Additional Inherited Members

Detailed Description

Member Typedef Documentation

◆ MR

◆ UnaryOperationBase

Constructor & Destructor Documentation

◆ CudaMatMulBiasGeluOp()

Member Function Documentation

◆ backward()

◆ className()

◆ forward()

◆ getName()