flashinfer.cake_causal_conv1d_train.KcpRecomputeHandle

class flashinfer.cake_causal_conv1d_train.KcpRecomputeHandle

The boundary tensors of one checkpointed call, carried from its original forward to torch’s recompute of that forward, so the recompute runs without cross-rank communication.

Create one handle per convolution call and micro-step, outside the outermost checkpoint region, and pass it as an argument of the checkpointed function, so the recompute sees the same object:

def layer(x, weight, handle):
    return cake_causal_conv1d_kcp(x, weight, ctx, recompute_handle=handle)

y = torch.utils.checkpoint.checkpoint(layer, x, w, KcpRecomputeHandle(), use_reentrant=False)

entries maps an operation key to that operation’s entry: "conv" is this package’s (the forward’s history pool hist, the very tensor the saved set holds); other packages that accept a handle with an entries dict use their own keys, so the handle of a convolution may also serve the recurrence of the same layer. A handle created inside a checkpoint region is a new, empty one in that region’s recompute, which then exchanges again (correct, not communication-free). The entries live as long as the handle; nbytes() is the retained size.

The forward halo exchange is a collective. Every rank of ctx.group must pass a handle for the same call, or no rank passes one, and every rank creates it at the same checkpoint nesting level. A mismatch pairs one rank’s recompute exchange with another rank’s backward exchange: the ranks hang or exchange wrong rows. No check is made. Never share one handle between two convolution calls (also not in nested regions). Never call clear() on a handle, or reuse it, while a backward that needs its entry is pending.

__init__() → None

Methods

__init__()
clear()	Drop every entry (the retained tensors go with the last reference).
nbytes()	Bytes of the distinct tensors the entries hold.

Attributes

entries