For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
ReduceRMSNormRoPE
struct ReduceRMSNormRoPE
Fuses RMS normalization and Rotary Position Embedding (RoPE) into one operation.
Computes per-row RMS normalization scaled by weight, then applies RoPE to
the normalized values using the provided cosine and sine tables. The last
dimension of the input must be an even number.
Implemented traits
AnyType,
ImplicitlyDestructible
Methods
execute
static def execute[dtype: DType, cos_sin_dtype: DType, rank: Int, target: StringSlice[StaticConstantOrigin], multiply_before_cast: Bool = True](output: ManagedTensorSlice[FusedOutput, static_spec=output.static_spec], input: ManagedTensorSlice[FusedInput, static_spec=input.static_spec], weight: ManagedTensorSlice[Input, static_spec=weight.static_spec], epsilon: Scalar[dtype], weight_offset: Scalar[dtype], cos_vals: ManagedTensorSlice[FusedInput, static_spec=cos_vals.static_spec], sin_vals: ManagedTensorSlice[FusedInput, static_spec=sin_vals.static_spec], ctx: DeviceContext)
shape
static def shape[dtype: DType, cos_sin_dtype: DType, rank: Int](input: ManagedTensorSlice[Input, static_spec=input.static_spec], weight: ManagedTensorSlice[Input, static_spec=weight.static_spec], epsilon: Scalar[dtype], weight_offset: Scalar[dtype], cos_vals: ManagedTensorSlice[Input, static_spec=cos_vals.static_spec], sin_vals: ManagedTensorSlice[Input, static_spec=sin_vals.static_spec]) -> IndexList[rank]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!