Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy > Member List

This is the complete list of members for Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >, including all inherited members.

backward(const TensorType &input, const TensorType &output_grad)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
build(const BuildContext &context) finalMila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
build_context_Mila::Dnn::Component< TDeviceType, TComputePrecision >protected
Component(const std::string &name)Mila::Dnn::Component< TDeviceType, TComputePrecision >inlineexplicit
ComponentBase typedef (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >
decode(const TensorType &q, const TensorType &k, const TensorType &v, dim_t position_offset)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
forward(const TensorType &input)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
getConfig() const noexcept (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
getDeviceId() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
getDeviceType()Mila::Dnn::Component< TDeviceType, TComputePrecision >inlinestatic
getExecutionContext() constMila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
getGradients() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
getMemoryStats() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
getModelDim() const noexcept (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
getName() constMila::Dnn::Component< TDeviceType, TComputePrecision >inline
getNumHeads() const noexcept (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
getNumKvHeads() const noexcept (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
getParameterNames() constMila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
getParameters() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
getPrecision() noexceptMila::Dnn::Component< TDeviceType, TComputePrecision >inlinestatic
getRequiredMemory(const BuildContext &context) const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
getTrainingMode() const noexceptMila::Dnn::Component< TDeviceType, TComputePrecision >inline
getType() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
GroupedQueryAttention(const std::string &name, const GqaConfig &config, std::optional< DeviceId > device_id=std::nullopt)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlineexplicit
hasExecutionContext() const noexceptMila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
HostStagingMemoryResource typedefMila::Dnn::Component< TDeviceType, TComputePrecision >protected
installSharedOutput(std::shared_ptr< TensorType > output)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
isBuilt() const finalMila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
kCacheDtype (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >static
kKvCompressed (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >static
KvCacheTensorType typedef (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >
load_(ModelArchive &archive, SerializationMode mode)Mila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
loadParameter(const std::string &, const Serialization::ITensorBlob &)Mila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
loadParameterFromBlob(const std::string &param_name, const Serialization::ITensorBlob &blob, Tensor< TParameterPrecision, TMemoryResource > &target, const shape_t &expected_shape)Mila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
MR typedef (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >
onBuilding(const BuildContext &context) overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlineprotectedvirtual
onExecutionContextSet() overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlineprotectedvirtual
onTrainingModeChanging(TrainingMode training_mode) overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlineprotectedvirtual
OpType typedef (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >
parameterCount() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
prefill(const TensorType &q, const TensorType &k, const TensorType &v, dim_t position_offset)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
requireSerializableParameters() constMila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
resetKVCache()Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
rewindKvCache(dim_t position)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
save_(ModelArchive &, SerializationMode) const override (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
saveFlatTensors(Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) constMila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
saveParameterToArchive(ModelArchive &archive, const std::string &parameter_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter) constMila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
saveParameterToWriter(Serialization::SafeTensorsWriter &writer, const std::string &flat_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter, Serialization::TensorSavePass pass) constMila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
setExecutionContext(IExecutionContext *context)Mila::Dnn::Component< TDeviceType, TComputePrecision >inlineprotected
setState(const GqaState &state)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
setTrainingMode(TrainingMode mode)Mila::Dnn::Component< TDeviceType, TComputePrecision >inline
setUseFlashDecode(bool enabled)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
setUseFlashPrefill(bool enabled)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
supportsKVCache() const noexceptMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inline
synchronize() overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
TensorType typedef (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >
toString() const overrideMila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >inlinevirtual
zeroGradients()Mila::Dnn::Component< TDeviceType, TComputePrecision >inlinevirtual
~GroupedQueryAttention() override=default (defined in Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >)Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >