flashinfer.cake_mhc_train.CakeMHCTrainPrepareFunction

class flashinfer.cake_mhc_train.CakeMHCTrainPrepareFunction(*args, **kwargs)

Autograd node of forward_prepare: forward(X, fn, scale, base, state) -> (U, *edges).

Its backward is backward_finish: it receives dU from the external branch and, as the gradients of the edges (outputs the finish node takes as inputs), dY and the gradient statistics from the finish node’s backward; it returns the single-cast dX (mapping, collapse and residual-mixing contributions summed in FP32) and the FP32 parameter gradients; the finish node returns no gradient for X.

__init__(*args, **kwargs)

Methods

__init__(*args, **kwargs)
apply
backward(ctx, dU, dY, grad_stats[, stats])	Define a formula for differentiating the operation with backward mode automatic differentiation.
forward(ctx, X, fn, scale, base, state)	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