Python function

fbcache_conditional_execution

fbcache_conditional_execution()

max.pipelines.modeling.base.fbcache_conditional_execution(first_block_residual, prev_residual, prev_output, residual_threshold, run_remaining_blocks, run_remaining_kwargs, run_postamble, temb, output_types)

source

Handle FBCache F.cond branching pattern shared across DiT models.

The caller provides atomic DiT methods:

• run_remaining_blocks: runs blocks 1..N + single-stream blocks, returns pre-tail hidden states.
• run_postamble: applies final norm + projection.

The residual_threshold is a scalar Tensor (float32, shape=[]) passed as a graph input so it can be changed at runtime without recompilation.

Returns:

(first_block_residual, output) tensors.

Parameters:

• first_block_residual (Tensor)
• prev_residual (Tensor)
• prev_output (Tensor)
• residual_threshold (Tensor)
• run_remaining_blocks (Callable[[...], Tensor])
• run_remaining_kwargs (dict[str, Any])
• run_postamble (Callable[[Tensor, Tensor], Tensor])
• temb (Tensor)
• output_types (list[TensorType])

Return type:

tuple[Tensor, Tensor]

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit