flashinfer.cake_mhc_train.CakeMHCTrainFinishFunction

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

Autograd node of forward_finish: forward(X, B, state, *edges) -> Y.

Its backward is backward_prepare: it returns dB for the external branch and hands dY and the gradient statistics (under save_stats=False also the statistics it replayed) to the prepare node as the gradients of edges, the prepare node’s edge outputs, so the autograd engine releases them once the prepare node’s backward has run. When a checkpoint region that starts after the prepare reruns this forward, the statistics the state already released are replayed bitwise from X; a reentrant region reruns it without the edges, and its backward leaves them on the state for the prepare node instead.

__init__(*args, **kwargs)

Methods

__init__(*args, **kwargs)
apply
backward(ctx, dY)	Define a formula for differentiating the operation with backward mode automatic differentiation.
forward(ctx, X, B, state, *edges)	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