flashinfer.cake_rmsnorm_train.CakeRMSNormFunction¶

class flashinfer.cake_rmsnorm_train.CakeRMSNormFunction(*args, **kwargs)¶

Autograd wrapper saving only (x_or_h_new, w, rstd) for the backward.

forward(x, w, eps, residual=None, deterministic=True) returns y or (y, h_new). In residual mode the gradient of x and of residual is the same tensor (dx + g_h_new). The weight gradient is accumulated in FP32 by the kernel and returned in w.dtype.

__init__(*args, **kwargs)¶

Methods

`__init__`(args, *kwargs)
`apply`
`backward`(ctx, grad_y, *grad_rest)	Define a formula for differentiating the operation with backward mode automatic differentiation.
`forward`(ctx, x, w, eps[, residual, ...])	Define the forward of the custom autograd Function.
`jvp`(ctx, *grad_inputs)	Define a formula for differentiating the operation with forward mode automatic differentiation.
`mark_dirty`(*args)	Mark given tensors as modified in an in-place operation.
`mark_non_differentiable`(*args)	Mark outputs as non-differentiable.
`mark_shared_storage`(*pairs)
`maybe_clear_saved_tensors`
`name`
`register_hook`
`register_prehook`
`save_for_backward`(*tensors)	Save given tensors for a future call to `backward()`.
`save_for_forward`(*tensors)	Save given tensors for a future call to `jvp()`.
`set_materialize_grads`(value)	Set whether to materialize grad tensors.
`set_output_grad_dtype`(*dtypes)	Declare the gradient dtype for each of this Function's outputs.
`setup_context`(ctx, inputs, output)	There are two ways to define the forward pass of an autograd.Function.
`vjp`(ctx, *grad_outputs)	Define a formula for differentiating the operation with backward mode automatic differentiation.
`vmap`(info, in_dims, *args)	Define the behavior for this autograd.Function underneath `torch.vmap()`.

Attributes

`boxed_grads_call`
`clear_saved_tensors_on_access`	Bool that specifies if backward should receive grads as a single mutable list argument instead of individual args in an immutable tuple.
`dirty_tensors`
`generate_vmap_rule`
`materialize_grads`
`metadata`
`needs_input_grad`
`next_functions`
`non_differentiable`
`output_grad_dtypes`
`requires_grad`
`saved_for_forward`
`saved_tensors`
`saved_variables`
`to_save`