flashinfer.cake_causal_conv1d_train.CakeCausalConv1dFunction

class flashinfer.cake_causal_conv1d_train.CakeCausalConv1dFunction(*args, **kwargs)

Autograd wrapper saving only (x, weight, cu_seqlens, bias) for the backward.

forward(x, weight, cu_seqlens, bias=None, silu_grad_preactivation="bf16") returns y; the backward runs the fused dx + deterministic dweight (+ dbias) pair of launches of cake_causal_conv1d_train_backward() on the incoming gradient (materialized as a contiguous BF16 tensor when autograd hands over another layout or dtype) and returns the gradients of weight and bias in the parameters’ dtypes.

__init__(*args, **kwargs)

Methods

__init__(*args, **kwargs)
apply
backward(ctx, grad_y)	Define a formula for differentiating the operation with backward mode automatic differentiation.
forward(ctx, x, weight, cu_seqlens[, bias, ...])	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