flashrt/diffusion-step-ops

Diffusion Step Ops

FlashRT CUDA kernels for small but frequent diffusion/runtime step operations.

These kernels target static-buffer and CUDA Graph friendly pipelines where PyTorch eager glue can become visible in the hot path.

Available Functions

• add_bf16(a, b): BF16 elementwise add.
• euler_step_bf16(latent, velocity, dt): BF16 Euler update.
• cfg_combine_into_residual_bf16(residual, v_cond, v_uncond, beta): in-place classifier-free guidance residual combine.
• cfg_combine_into_residual_fp16(residual, v_cond, v_uncond, beta): FP16 variant.
• teacher_force_first_frame_bf16(video_latent, cond_latent): copy conditioning frame into video_latent[:, :, 0].
• motus_decode_postprocess_bf16_to_fp32(decoded): drop first frame and map [-1, 1] to [0, 1].
• cast_bf16_to_fp32(src): BF16 to FP32 cast.
• pack_tail_bf16(tail, flat_dim): zero-pad a BF16 tail into a flat vector.
• add_bias_zero_tail_bf16(input, bias, valid_cols): add bias and zero padded columns.
• extract_tail_f32_to_bf16(flat, tail_numel): extract and cast an action tail.
• add_bias_pair_bf16(input, bias_a, bias_b): preserve two BF16 add-rounding stages.
• unipc_step_f32_bf16(...): fused UniPC corrector/predictor update.

Usage

from kernels import get_kernel

ops = get_kernel("flashrt/diffusion-step-ops")

latent = ops.euler_step_bf16(latent, velocity, dt=-0.125)
ops.cfg_combine_into_residual_bf16(residual, v_cond, v_uncond, beta=4.5)
ops.teacher_force_first_frame_bf16(video_latent, cond_latent)

next_sample, current_m, current_last = ops.unipc_step_f32_bf16(
    sample, velocity, prev_m1, prev_m2, prev_last,
    sigma, corrector_order, predictor_order,
    corrector_coefficients, predictor_coefficients,
)

All APIs require CUDA contiguous tensors. Unsupported shapes fail at the wrapper boundary.

The generic tail APIs cover the Cosmos3-Edge runtime contracts without model-specific aliases: pack_tail_bf16 is equivalent to the native fill-flat-velocity kernel, and extract_tail_f32_to_bf16 is equivalent to the native copy-action-tail kernel. Validation includes production flat_dim=1,201,920, tail_numel=3,840, and exact CUDA Graph replay.