Add Cosmos3-Edge BF16 production gates
Browse files
README.md
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Diffusion Step Ops
|
| 2 |
+
|
| 3 |
+
FlashRT CUDA kernels for small but frequent diffusion/runtime step operations.
|
| 4 |
+
|
| 5 |
+
These kernels target static-buffer and CUDA Graph friendly pipelines where
|
| 6 |
+
PyTorch eager glue can become visible in the hot path.
|
| 7 |
+
|
| 8 |
+
## Available Functions
|
| 9 |
+
|
| 10 |
+
- `add_bf16(a, b)`: BF16 elementwise add.
|
| 11 |
+
- `euler_step_bf16(latent, velocity, dt)`: BF16 Euler update.
|
| 12 |
+
- `cfg_combine_into_residual_bf16(residual, v_cond, v_uncond, beta)`: in-place classifier-free guidance residual combine.
|
| 13 |
+
- `cfg_combine_into_residual_fp16(residual, v_cond, v_uncond, beta)`: FP16 variant.
|
| 14 |
+
- `teacher_force_first_frame_bf16(video_latent, cond_latent)`: copy conditioning frame into `video_latent[:, :, 0]`.
|
| 15 |
+
- `motus_decode_postprocess_bf16_to_fp32(decoded)`: drop first frame and map `[-1, 1]` to `[0, 1]`.
|
| 16 |
+
- `cast_bf16_to_fp32(src)`: BF16 to FP32 cast.
|
| 17 |
+
- `pack_tail_bf16(tail, flat_dim)`: zero-pad a BF16 tail into a flat vector.
|
| 18 |
+
- `add_bias_zero_tail_bf16(input, bias, valid_cols)`: add bias and zero padded columns.
|
| 19 |
+
- `extract_tail_f32_to_bf16(flat, tail_numel)`: extract and cast an action tail.
|
| 20 |
+
- `add_bias_pair_bf16(input, bias_a, bias_b)`: preserve two BF16 add-rounding stages.
|
| 21 |
+
- `unipc_step_f32_bf16(...)`: fused UniPC corrector/predictor update.
|
| 22 |
+
|
| 23 |
+
## Usage
|
| 24 |
+
|
| 25 |
+
```python
|
| 26 |
+
from kernels import get_kernel
|
| 27 |
+
|
| 28 |
+
ops = get_kernel("flashrt/diffusion-step-ops")
|
| 29 |
+
|
| 30 |
+
latent = ops.euler_step_bf16(latent, velocity, dt=-0.125)
|
| 31 |
+
ops.cfg_combine_into_residual_bf16(residual, v_cond, v_uncond, beta=4.5)
|
| 32 |
+
ops.teacher_force_first_frame_bf16(video_latent, cond_latent)
|
| 33 |
+
|
| 34 |
+
next_sample, current_m, current_last = ops.unipc_step_f32_bf16(
|
| 35 |
+
sample, velocity, prev_m1, prev_m2, prev_last,
|
| 36 |
+
sigma, corrector_order, predictor_order,
|
| 37 |
+
corrector_coefficients, predictor_coefficients,
|
| 38 |
+
)
|
| 39 |
+
```
|
| 40 |
+
|
| 41 |
+
All APIs require CUDA contiguous tensors. Unsupported shapes fail at the
|
| 42 |
+
wrapper boundary.
|
| 43 |
+
|
| 44 |
+
The generic tail APIs cover the Cosmos3-Edge runtime contracts without
|
| 45 |
+
model-specific aliases: `pack_tail_bf16` is equivalent to the native
|
| 46 |
+
fill-flat-velocity kernel, and `extract_tail_f32_to_bf16` is equivalent to the
|
| 47 |
+
native copy-action-tail kernel. Validation includes production
|
| 48 |
+
`flat_dim=1,201,920`, `tail_numel=3,840`, and exact CUDA Graph replay.
|