diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index e2cd6646a..6e0588c43 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -259,6 +259,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_H3_VSA_FP4_K_SMOOTH` | bool | `0` | attention | Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD training numerics. Exact for the softmax; it only changes quantization error. | | `FASTVIDEO_H3_GATE_CHUNK_ROWS` | int | `0` | performance | Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back to token rows. 0 keeps the single full-size index_select; a positive value caps the transient at chunk_rows\*heads\*head_dim and is bit-identical. | | `FASTVIDEO_H3_VSA_FINE_ATTENTION` | one of , emulator-mix8, emulator-mixmx, emulator-qk, emulator-bf16 | `""` | attention | Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over FASTVIDEO_H3_VSA_FP4. Empty keeps the native route. | +| `FASTVIDEO_WAN22_VAE_DECODE_DIR` | path | `""` | performance | Directory of the Wan2.2 VAE decoder kernels (fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds on first use, when its config sets decode_kernel. Empty takes the copy in this checkout. | | `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE` | bool | `0` | performance | VSA-H3 64/128-token tiles: scatter rows straight into the heads-first layout the sparse kernels read, removing their per-call transpose copy of query, key and value. Bit-identical; no-grad, uncompiled CUDA only. | | `FASTVIDEO_H3_VSA_TILE_FIRST` | bool | `0` | attention | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters. | | `FASTVIDEO_H3_VSA_SM89_KERNEL` | one of original, bf16, int8 | `original` | attention | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV). | @@ -328,6 +329,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. | | `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. | | `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. | +| `FASTVIDEO_TEST_PAN2_MODEL_DIR` | path | unset | test | PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU tests; unset skips them. | | `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. | | `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR` | str | unset | test | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`. | diff --git a/docs/inference/pan2.md b/docs/inference/pan2.md index 570594a6c..8e65bb5e2 100644 --- a/docs/inference/pan2.md +++ b/docs/inference/pan2.md @@ -19,7 +19,7 @@ from the `pan2_t2v` and `pan2_i2v` presets: | Resolution | 704x1248 (height and width must be multiples of 32) | | Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) | | Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) | -| Guidance | 3.0, classifier-free guidance. Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) | +| Guidance | 2.5, classifier-free guidance, the reference launcher's default (3 before 2026-10-23). Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) | | Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps | PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw @@ -67,6 +67,20 @@ The trimmed sequence can use either FlashAttention or Torch SDPA. PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py ``` +## Performance switches + +PAN2 sampling turns these on by default (`PAN2PipelineConfig`); pass any of them as `False` to +`VideoGenerator.from_pretrained` to turn it off: + +| Switch | What it does | +|---|---| +| `vae_decode_kernel` | Decodes with fastvideo-kernel's Wan2.2 decoder on B200 / GB200 / B300 (704x1248, batch 1): bf16 with fp32 accumulation, 7x faster than the fp32 decode on a B300 at PSNR 61-63 dB and SSIM 0.998 against it, at a peak of about 20 GiB for a 5 s video. It builds on first use (minutes) into `WAN22_VAE_DECODE_BUILD_DIR`, against fastvideo-kernel's CUTLASS submodule (`git submodule update --init fastvideo-kernel/include/cutlass`, or set `CUTLASS_DIR`). Other GPUs and sizes keep the fp32 decode. | +| `vae_decode_kernel_parallel` | With 2, 4 or 8 GPUs, splits that decode over them instead of decoding the whole video on each: a 20 s video takes 3.0 s on one B300, 1.55 s on two, 0.79 s on four and 0.47 s on eight. | +| `dit_fused_kernels` | Runs the DiT's fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation. Their rounding differs slightly from the unfused path. | + +On GPUs with room for every component (B200, B300), keep the text encoder resident as well: +`text_encoder_cpu_offload=False`. + ## Memory and offloading The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index b5f08734f..a60b1ff76 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -21,6 +21,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image): | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) | | FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below | | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen | +| Wan2.2 VAE decoder (TI2V-5B, 48 latent channels) at 704 x 1248 on `sm_100a` / `sm_103a`, on 1-8 GPUs of a node | `wan22_vae_decode/`, built from `wan22_vae_decode/solution/` on first use against `include/cutlass` | `wan22_vae_decode.WanVAEDecoder`; FastVideo's `WanVAEConfig.decode_kernel`, on for PAN2 sampling (see `wan22_vae_decode/README.md`) | ## What gets built where, and when diff --git a/fastvideo-kernel/tests/test_wan22_vae_decode.py b/fastvideo-kernel/tests/test_wan22_vae_decode.py new file mode 100644 index 000000000..e1d115bcd --- /dev/null +++ b/fastvideo-kernel/tests/test_wan22_vae_decode.py @@ -0,0 +1,97 @@ +"""The Wan2.2 VAE decoder (fastvideo-kernel/wan22_vae_decode) through its wrapper, against the reference. + +Needs an sm_100a or sm_103a GPU (B200 / GB200, B300) and the VAE's weights: WAN22_VAE_WEIGHTS names either a diffusers AutoencoderKLWan +directory (its diffusion_pytorch_model.safetensors) or a .pt file holding the 110 tensors in weight_names() order. The +first test builds the extension (about a minute). The kernel's own checks are wan22_vae_decode/tools/bench.py (nine +timed workloads and six short clips at 1, 2 and 4 GPUs, on real latents); these tests cover the wrapper on one GPU: +the weights it reads from a state dict, and decodes against the fp32 reference, each within three times what the bf16 +reference errs by on the same latents (the bench's rule, calibrated per input instead of per clip). +""" + +import functools +import os +import sys + +import pytest +import torch + +PACKAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "wan22_vae_decode") +sys.path.insert(0, PACKAGE) +sys.path.insert(0, os.path.join(PACKAGE, "tools")) + +WEIGHTS = os.environ.get("WAN22_VAE_WEIGHTS", "") + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="the Wan2.2 VAE decoder is built for sm_100a / sm_103a (B200 / GB200, B300)"), + pytest.mark.skipif(not WEIGHTS or not os.path.exists(WEIGHTS), + reason="set WAN22_VAE_WEIGHTS to the VAE's diffusers directory or a weights .pt file"), +] + + +@functools.lru_cache(maxsize=None) +def _state_dict(): + """name -> tensor of the VAE's parameters, under diffusers' names.""" + from wan22_vae_decode import weight_names + + if os.path.isdir(WEIGHTS): + from safetensors.torch import load_file + + return load_file(os.path.join(WEIGHTS, "diffusion_pytorch_model.safetensors")) + tensors = torch.load(WEIGHTS, map_location="cpu", weights_only=True) + return tensors if isinstance(tensors, dict) else dict(zip(weight_names(), tensors)) + + +@functools.lru_cache(maxsize=None) +def _decoder(): + from wan22_vae_decode import WanVAEDecoder + + return WanVAEDecoder(_state_dict(), device="cuda") + + +def test_weights_follow_diffusers_names(): + from wan22_vae_decode import decoder_weights, weight_names + + names = weight_names() + assert len(names) == 110 and names[0] == "post_quant_conv.weight" and names[-1] == "decoder.conv_out.bias" + weights = decoder_weights(_state_dict(), "cpu") + assert all(w.dtype == torch.bfloat16 and w.is_contiguous() for w in weights) + assert all(torch.equal(w, _state_dict()[n].to(torch.bfloat16)) for n, w in zip(names, weights)) + + +@pytest.mark.parametrize("latent_frames", [1, 2, 5]) +def test_decode_against_reference(latent_frames): + import reference + from wan22_vae_decode import weight_names + + generator = torch.Generator().manual_seed(latent_frames) + z = torch.randn((1, 48, latent_frames, 44, 78), generator=generator).to("cuda", torch.bfloat16) + video = _decoder().decode(z) + assert video.shape == (1, 3, 4 * latent_frames - 3, 704, 1248) and video.dtype == torch.bfloat16 + + weights = {n: _state_dict()[n].to("cuda", torch.bfloat16).float() for n in weight_names()} + ref = reference.decode(z.float(), weights, torch.float32) + release = reference.decode(z, weights, torch.bfloat16).float() + d, d_release = video.float() - ref, release - ref + assert bool(torch.isfinite(video).all()) and float(video.abs().max()) <= 1.0 + assert float(d.abs().max()) <= 3 * float(d_release.abs().max()) + assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_release.pow(2).mean().sqrt()) + + +def test_decodes_are_independent(): + """No state carries from one call to the next: the same latents decode to the same video after another clip.""" + z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(7)).to("cuda", torch.bfloat16) + first = _decoder().decode(z) + _decoder().decode(torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(8)).to("cuda", + torch.bfloat16)) + assert torch.equal(_decoder().decode(z), first) + + +def test_pass_size_does_not_change_the_video(): + """The causal history crosses passes, so any number of latent frames per pass decodes the same video.""" + from wan22_vae_decode import WanVAEDecoder + + z = torch.randn((1, 48, 6, 44, 78), generator=torch.Generator().manual_seed(9)).to("cuda", torch.bfloat16) + reference_pass = _decoder().decode(z) + for frames in (1, 16): + assert torch.equal(WanVAEDecoder(_state_dict(), device="cuda", frames_per_pass=frames).decode(z), reference_pass) diff --git a/fastvideo-kernel/wan22_vae_decode/README.md b/fastvideo-kernel/wan22_vae_decode/README.md new file mode 100644 index 000000000..5f66abd5c --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/README.md @@ -0,0 +1,98 @@ +# Wan2.2 VAE decoder (sm_100a, sm_103a) + +The decoder of the Wan2.2 VAE (the VAE of Wan2.2-TI2V-5B: 48 latent channels, 16 x 16 spatial and 4x temporal +compression) at 704 x 1248, on one, two or four B200 / GB200 GPUs of a node. It decodes latents [1, 48, T, 44, 78] +into the video [1, 3, 4T - 3, 704, 1248] for any T >= 1, as diffusers' `AutoencoderKLWan.decode(z).sample` does; +`reference.py` defines the result exactly, a plain-PyTorch decoder bit-exact with diffusers in fp32 and in bf16. + +## Use + +```python +import sys +sys.path.insert(0, "fastvideo-kernel/wan22_vae_decode") +from wan22_vae_decode import WanVAEDecoder + +decoder = WanVAEDecoder(vae.state_dict(), device="cuda") # diffusers' AutoencoderKLWan, or its state dict +video = decoder.decode(latents) # [1, 48, T, 44, 78], as the VAE's decode takes them -> [1, 3, 4T-3, 704, 1248] bf16 +``` + +- **Over several GPUs**, every rank builds a decoder with the same process group (`WanVAEDecoder(sd, device, pg)`, + default: torch.distributed's default group) and calls `decode` with the full latents; rank 0 gets the video and the + other ranks `None`. On one GPU without torch.distributed, the decoder initializes a single-rank NCCL group. +- The first call builds the extension from `solution/` (about a minute) into `$WAN22_VAE_DECODE_BUILD_DIR` (default + `~/.cache/fastvideo/wan22_vae_decode`, one directory per `LOCAL_RANK`), and later processes reuse the build. +- `WanVAEDecoder(..., frames_per_pass=4)` sets the latent frames each GPU carries through one pass (see below). +- `WanVAEDecoder` prepares the weights once (layout, phase and Winograd transforms of the weights). The raw entry + points are `prepare(weights, rank, world, pg)`, `run(z, weights, rank, world, pg, out)` (4 frames per pass) and + `run_with(..., out, frames_per_pass)` in `solution/binding.cpp`, destination-passing, on the current stream; + `weights` are the 110 tensors of `weight_names()` in bf16. + +## Performance and memory + +At PAN2's default 704 x 1248, 24 fps, on GB200 GPUs, against PAN2's decode today (diffusers' `AutoencoderKLWan` in bf16 +under autocast, one GPU, one latent frame at a time). Memory is the peak allocated and, after it, the peak reserved by +PyTorch's caching allocator (what nvidia-smi shows), per GPU; the default is 4 latent frames per GPU per pass: + +| Decoder | 8 s (193 frames) | 12 s (289 frames) | 20 s (481 frames) | +|---|---|---|---| +| Release, 1 GPU | 5.88 s, 12.9 / 17.1 GiB | 8.84 s, 13.9 / 19.4 GiB | 14.76 s, 15.8 / 91.5 GiB | +| **This decoder, 1 GPU** | **1.17 s** (x5.0), 20.7 / 36.3 GiB | **1.74 s** (x5.1), 21.2 / 36.8 GiB | **2.89 s** (x5.1), 22.2 / 37.8 GiB | +| **This decoder, 2 GPUs** | **619 ms**, 18.7 / 33.0 GiB | **929 ms**, 19.2 / 33.5 GiB | **1.54 s**, 20.2 / 34.5 GiB | +| **This decoder, 4 GPUs** | **328 ms**, 17.8 / 41.4 GiB | **475 ms**, 18.3 / 34.1 GiB | **787 ms** (x18.8), 19.2 / 42.9 GiB | + +`frames_per_pass` trades memory for speed and leaves the video unchanged. On one GPU: + +| Latent frames per pass | 8 s | 12 s | 20 s | +|---|---|---|---| +| 2 | 1.25 s, 14.4 / 24.4 GiB | 1.88 s, 14.9 / 24.9 GiB | 3.13 s, 15.9 / 25.8 GiB | +| 4 (default) | 1.17 s, 20.7 / 36.3 GiB | 1.74 s, 21.2 / 36.8 GiB | 2.89 s, 22.2 / 37.8 GiB | +| 8 | 1.10 s, 33.7 / 62.1 GiB | 1.64 s, 34.2 / 62.6 GiB | 2.72 s, 35.2 / 63.6 GiB | +| 16 | 1.06 s, 59.8 / 121.0 GiB | 1.58 s, 60.2 / 121.7 GiB | 2.67 s, 61.2 / 122.5 GiB | + +- With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, the default reserves 30.6-32.0 GiB on one GPU and 2 frames + per pass 19.6-21.1 GiB, at the same speed. +- On the optimization run's nine workloads (`tools/bench.py`: 121, 241 and 481 frames at 1, 2 and 4 GPUs, SM clocks + 1.6-1.95 GHz), the default's geomean is 758 ms, against 1043 ms for the run's first solution (x1.38) and 698 ms at + 16 frames per pass. +- The CUTLASS Winograd component products take 60 % of the GPU time, the transforms and epilogues 30 % + (`notes/optimization.md`). + +## Correctness + +Against the fp32 decode of real latents, with `d = video - reference` per clip: every element finite and in [-1, 1], +max |d| at most three times what the release's bf16 decode errs by, and the RMS of d likewise. On this tree, every +clip at worlds 1, 2 and 4 and the short clips `f1`, `f5`, `f17` at worlds 1 and 4 pass, using 0.30-0.66 of the max +allowance and 0.35-0.42 of the RMS one (PSNR 56-60 dB against fp32; the release's bf16 decode has 58-61 dB). + +```bash +cd fastvideo-kernel/wan22_vae_decode +DATA=/path/to/data python tools/bench.py --solution solution # every workload this job has the GPUs for +DATA=/path/to/data python tools/bench.py --workload f5x1 f17x1 --no-time # correctness of the short clips +cd .. && WAN22_VAE_WEIGHTS=/path/to/vae python -m pytest tests/test_wan22_vae_decode.py # the wrapper, one GPU +``` + +`tools/bench.py` reads `$DATA`: the latents (`latents_.pt`), the weights in `weight_names()` order +(`decoder_weights_bf16.pt`, `weight_names.json`), the fp32 decodes (`ref_.pt`) and the release decode's errors +(`calibration.json`); `tools/spec.py` documents each. The tests need only the VAE's weights. + +## Limits + +- sm_100a and sm_103a (B200 / GB200, B300): `tcgen05` with tensor memory, TMA and 2-CTA clusters; the loader builds + for the GPU it runs on. CUDA 13, torch 2.12 tested; the tests pass on B300 too, built for sm_103a. +- 704 x 1248 video (latents 44 x 78) and batch 1; any T >= 1. +- 1, 2 or 4 GPUs of one node were checked; the code accepts up to 8 ranks. +- bf16 weights and activations, fp32 accumulation and fp32 RMS-norm statistics. +- In FastVideo, `WanVAEConfig.decode_kernel` routes `AutoencoderKLWan.decode` here where it applies (on one GPU per + rank, `WanVAEDecoder(..., local=True)`), and PAN2 sampling turns it on (`PAN2PipelineConfig.vae_decode_kernel`). On + a B300 with PAN2 latents it decodes 7x faster than the fp32 decode, at PSNR 61-63 dB and SSIM 0.998 against it. + +## Layout + +- `solution/`: the kernels and their binding; `solution.json` holds the build flags. The CUTLASS and CuTe headers + come from fastvideo-kernel's CUTLASS submodule (`include/cutlass`, or `$CUTLASS_DIR`); the decoder's structure follows + diffusers' (`WAN_LICENSE.txt`). +- `wan22_vae_decode.py`: `WanVAEDecoder`, `decoder_weights` and `weight_names`. +- `reference.py`: the definition of the result. +- `tools/`: `bench.py` (check, audit and time), `loader.py` (the build), `spec.py` (workloads, weights, data), + `audit.py`, `bounds.py`. +- `notes/optimization.md`: the design, what each merge bought, what did not pay, and where to go next. diff --git a/fastvideo-kernel/wan22_vae_decode/notes/optimization.md b/fastvideo-kernel/wan22_vae_decode/notes/optimization.md new file mode 100644 index 000000000..fcdae9f2b --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/notes/optimization.md @@ -0,0 +1,99 @@ +# How the decoder got here + +The solution is the main of a multi-agent optimization run (bifrost, 2026-10-25/26): 32 DeepSeek V4.1 Flash peers +and two GPT-6 Astra architects on GB200, starting from a first solution written before the run. Every merge passed the +run's CI (correct on every workload and short clip at worlds 1, 2 and 4, and at least 1 % faster than main by the +geomean in the same job, unless waived) and an independent review. The comments in `solution/` cite the run's pull +requests (PRn), issues (In) and jobs (Jn). + +## The design + +The decoder runs layer by layer, as kernels of its own and CUTLASS kernels instantiated in `solution/` (the +headers are fastvideo-kernel's CUTLASS submodule). There is no megakernel: see "What did not pay" below. + +- **Frames in passes.** A pass carries several latent frames per GPU instead of the reference's one (4 by default, + `frames_per_pass`; the run measured its speed at 16, see "Memory"), so its GEMMs are large, and the causal history + (the last two frames of each 3 x 3 x 3 convolution's input) is carried from pass to pass as an offset, never by + concatenation. +- **The 18 convolutions of up blocks 1-3** (most of the FLOPs) are F(4x4, 3x3) Winograd in space, keeping the causal + 3-tap convolution in time, in three steps: + 1. `wino_in.cu`: the input transform, two blocks per tile; + 2. 36 component products, each a CUTLASS SM100 implicit-GEMM 3 x 1 x 1 convolution (`conv.cu`: 2-CTA 256 x 256 x 64 + tiles, TMA im2col, `tcgen05` MMA, fp32 accumulation); + 3. `wino_epilogue.cu`: the inverse transform with the bias and the residual add, and the *next* convolution's RMS + norm, SiLU and causal padding, so that convolution reads its input ready-made. +- **The other convolutions** at 44 x 78 are direct CUTLASS implicit GEMMs; a few small ones still go through cuDNN + and cuBLAS, under 1 % of the GPU time. +- **The upsamplers** never materialize the upsampled tensor: nearest x2 followed by a 3 x 3 convolution is four 2 x 2 + phase convolutions at the low resolution, and `resample.cu` interleaves them with the bias and the DupUp3D shortcut + in one pass. The temporal upsampler (`timeconv.cu`) writes its two frames straight into place. +- **Pointwise and the rest:** `norm.cu` (RMS norm + SiLU + causal pad in one pass, where the epilogue does not already + apply it), `attn.cu` (the mid block's single-head attention over each frame's 3,432 positions), and `head.cu` (the + output head: `conv_out` as a 3 x 1 x 1 convolution to 112 channels, then a 9-tap stencil with the bias, the clamp + and the unpatchify write). +- **Over GPUs** (`comm.cu`, `prefix.cu`): the decoder's prefix (post_quant_conv, conv_in, the mid block's first + residual block and its attention) is split by time, each rank computing its frames with exact causal context from + the full latents, then transposed to bands of rows (`prefix.cu`). From there each rank owns 44 / world latent rows; + every 3 x 3 convolution exchanges one halo row with the neighbours only (NCCL send/recv in one group), and rank 0 + receives the other ranks' pixels per chunk. + +## The merges + +| Merge | Change | CI speedup | Geomean after | +|---|---|---|---| +| kernel#0 | the first solution: Winograd up blocks, phase upsamplers, fused pointwise, row bands | | 1043 ms | +| PR3 | 1 x 1 channel shortcuts as own kernels, their bias folded into the residual | x1.017 | 1025 ms | +| PR24 | mid-block attention as own kernels instead of ATen SDPA (8.5 -> 2.6 ms of f121x1) | x1.013 | 1012 ms | +| PR15 | RMS norm + SiLU + pad in one pass | x1.110 | 913 ms | +| PR204 | halo exchange with the neighbours only, one coalesced group | x1.014 | 901 ms | +| PR392 | the fused Winograd input transform and the fused inverse epilogue (with the next norm) | x1.189 | 757 ms | +| PR90 | output head as a temporal convolution and a 9-tap stencil (waived at parity) | x1.003 | 755 ms | +| PR438 | the prefix split by time before the row bands | x1.019 | 740 ms | +| PR452 | world 8 uses all-gather, not the ring (worlds 1/2/4 unchanged; waived) | x1.001 | 740 ms | +| PR488 | the upsampler's phases interleaved in one pass, and the temporal upsampler in place | x1.017 | 728 ms | +| PR513 | chunks of min(16 x world, 61) latent frames instead of 8 x world | x1.023 | 711 ms | +| PR490 | a check that no fused input leaks across chunks (waived) | x0.999 | 712 ms | +| PR592 | no dead store of y at each residual block's first convolution | x1.010 | 704 ms | +| PR517 | the input transform at two blocks per tile (waived at x1.007) | x1.007 | 699 ms | + +## Memory + +The run optimized time alone, and its main decoded 16 latent frames per GPU per pass at a peak of 100-102 GiB on one +GPU, against 13-16 GiB for the release's decode, which carries one latent frame. After the run: + +- **Tensors are freed at their last use.** At the peak, up block 3's first convolution (512 -> 256 channels at + 352 x 624) held its padded input after the Winograd transform had read it, the transform (2.25x the input) after + every product was enqueued, and the block's input after its shortcut and first norm had read it. Releasing each of + them then takes the peak from 101 to 60 GiB at 16 frames per pass, at the same speed, with a bit-identical video. +- **The pass is 4 latent frames per GPU by default** (`frames_per_pass`), since memory grows with it at 3.5-4.2 GiB per + latent frame at 704 x 1248: 21-22 GiB on one GPU and 18-20 GiB per GPU at 2 and 4 GPUs, for 8-9 % more time than 16 + frames per pass. The time a smaller pass costs is in the low resolutions (44 x 78 and 88 x 156), whose GEMMs do not + fill the GPU with few frames, while the memory is in the two high ones (176 x 312 and 352 x 624). Passes of 16 frames + through up block 1 and of 2-4 through up blocks 2 and 3 should keep most of the speed at this memory; it is not built. + +## What did not pay + +- **Megakernels.** Before the run, a persistent kernel per residual block measured 4.1-7.5x slower than two fused + per-convolution kernels. During the run, a cooperative kernel for one whole Winograd convolution (forward transform, + products and inverse with the next norm, in one launch of 152 CTAs in 2-CTA clusters) was correct but took + 6.2-7.2 ms at T 8, 352 x 624, 256 channels, against 3.3 ms for the same work in three launches (0.46-0.53x). The gaps + between the kernels of a decode are about 5 us each, so fewer launches buy little; what a single kernel loses is + the occupancy and the tiling each stage picks for itself. +- **FP8.** e4m3 operands in the Winograd component products cut the products from 262 to 162 ms at f121x1 (10-11 % + of the decode), but fail the error budget on all three clips (RMS at 2.3-3.3x the allowance, even with a power-of-two + scale per component): the transformed operands need about one more mantissa bit. Before the run, FP8 in the direct + convolutions passed in up blocks 1-3 but bought only 1 % of the decode, since quantization was a pass of its own. +- **Splitting by time throughout.** A chunk needs 38 latent frames of causal context, so recomputing it is a poor fit; + only the prefix is split by time (PR438). + +## Where to go next + +- **The products** are about 60 % of the GPU time (59.6 % at f121x1: 1,383 CUTLASS launches, 384 ms of 644), at + 65-68 % of the tensor cores' peak, bound by operand fetch. The transforms and the epilogues are another 30 %. A grouped products kernel (all 36 components of a convolution, persistent, in + 2-CTA clusters) passed correctness during the run; the step after it is to compute the input transform inside the + products' operand producer, which removes a write and a read of the transformed tensor (2.25x the activation) per + convolution. +- **Eight GPUs across two trays:** the code accepts up to 8 ranks (PR452 fixed world 8's all-gather), but world 8 + was never in CI. A frame pipeline that passes each + layer's two-frame context beat row bands at eight GPUs before the run (1.40-1.55 s against 2.12 s on f241, with + library convolutions). diff --git a/fastvideo-kernel/wan22_vae_decode/reference.py b/fastvideo-kernel/wan22_vae_decode/reference.py new file mode 100644 index 000000000..cc036fa29 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/reference.py @@ -0,0 +1,251 @@ +"""Wan2.2 VAE decoder: a standalone PyTorch reference of diffusers' AutoencoderKLWan._decode (Wan2.2-TI2V-5B VAE). + +decode(z, weights, dtype) maps latents z [B, 48, T, H/16, W/16] to a video [B, 3, 4T-3, H, W] with values in [-1, 1], +as AutoencoderKLWan.decode(z).sample does with tiling and slicing off (its defaults). It needs torch and nothing else: +`weights` is a plain dict name -> tensor of the decoder's and post_quant_conv's parameters under diffusers' names +(the 110 of wan22_vae_decode.weight_names()), such as the VAE's state dict: + + weights = {k: v.cuda().float() for k, v in vae.state_dict().items()} + video = decode(latents.cuda(), weights) # fp32 + video = decode(latents.cuda(), weights, torch.bfloat16) # a pipeline's bf16 decode + +What is computed +---------------- +The decoder of the TI2V-5B VAE (z_dim 48, decoder_base_dim 256, dim_mult [1, 2, 4, 4], num_res_blocks 2, +is_residual, patch_size 2, temporal upsampling in the first two up blocks, as the VAE's config.json says): + +1. post_quant_conv: a 1x1x1 conv 48 -> 48, over all T latent frames at once. +2. The decoder, run on one latent frame at a time, in order (chunk i is latent frame i). A chunk sees the earlier + ones only through the causal caches (below). With h x w = H/16 x W/16: + conv_in causal 3x3x3 conv 48 -> 1024 at h x w + mid block residual block 1024, attention 1024, residual block 1024 at h x w + up block 0 3 residual blocks 1024; upsample3d 1024 (t x2, h w x2); + DupUp3D 1024 -> 1024 to 2h x 2w + up block 1 3 residual blocks 1024; upsample3d 1024 (t x2, h w x2); + DupUp3D 1024 -> 1024 to 4h x 4w + up block 2 residual blocks 1024 -> 512, 512, 512; upsample2d 512 (h w x2); + DupUp3D 1024 -> 512 to 8h x 8w + up block 3 residual blocks 512 -> 256, 256, 256; no resampler, no shortcut at 8h x 8w + norm_out RMS norm 256, then SiLU + conv_out causal 3x3x3 conv 256 -> 12 at 8h x 8w + Up blocks 0 to 2 add their shortcut, DupUp3D of the block's input, to the output of their resampler. + The chunks' outputs are concatenated along time. +3. unpatchify, [B, 12, F, 8h, 8w] -> [B, 3, F, H, W]: the 12 channels are (colour k, column offset j, row offset i) + of 2 x 2 pixel patches, video[b, k, f, 2y + i, 2x + j] = out[b, 4k + 2j + i, f, y, x]. +4. clamp to [-1, 1]. + +The parts, with diffusers' class names: +- causal conv (WanCausalConv3d), kernel kt x kh x kw, stride 1: zero padding of kh//2 and kw//2 on both sides in + space; in time, 2 * (kt//2) frames before the chunk and none after: the frames of the conv's cache, and zeros for + those the cache lacks. +- RMS norm (WanRMS_norm) over channels: F.normalize, x / max(||x||_2, 1e-12), times sqrt(C), times gamma. +- residual block (WanResidualBlock), in -> out channels: x1 = conv1(SiLU(norm1(x))), x2 = conv2(SiLU(norm2(x1))), + causal 3x3x3 convs; it returns x2 + shortcut, the shortcut being x, or a 1x1x1 conv in -> out of x when in != out. + (Dropout is 0, and diffusers decodes in eval mode.) +- attention (WanAttentionBlock), per frame: n = RMS norm of x; q, k and v are the channels [0, C), [C, 2C) and + [2C, 3C) of a 1x1 conv C -> 3C of n; single-head attention softmax(q k^T / sqrt(C)) v over the h*w positions of the + frame (head dim 1024, no mask: diffusers' docstring calls it causal, it is not); a 1x1 conv C -> C; plus x. +- upsample3d (WanResample "upsample3d"): the time conv, a causal 3x1x1 conv C -> 2C whose output channel c becomes + frame 2t and channel C + c frame 2t + 1, then upsample2d. upsample2d (WanResample "upsample2d"): per frame, nearest + x2 in h and w (mode nearest-exact, WanUpsample) and a 3x3 conv2d C -> C with zero padding 1. +- DupUp3D, in -> out channels, factor ft in time (2 in up blocks 0 and 1, 1 in up block 2) and 2 in h and w: each + input channel repeated r = out * ft * 4 / in times, the copies laid out over the (ft, 2, 2) sub-positions: + dup[c, ft*t + a, 2y + i, 2x + j] = x[(4*ft*c + 4a + 2i + j) // r, t, y, x]. In up blocks 0 and 1 (r = 8) that is + nearest x2 in t, h and w; in up block 2 (r = 2) it is dup[c, t, 2y + i, 2x + j] = x[2c + i, t, y, x]: channel pairs + become row pairs, each value repeated in two columns. + +The first latent frame yields one frame +--------------------------------------- +On the first chunk (diffusers' first_chunk): +- upsample3d skips the time conv: the chunk's frame passes through as one frame (diffusers marks the conv's cache + "Rep"). The time conv's input sequence therefore begins at the second chunk, after zeros: it never sees latent + frame 0. (Diffusers runs the second chunk's time conv without history and caches a zero frame in front of that + chunk's frame; giving it no history, as here, and zero padding yields the same frames.) +- DupUp3D keeps only the last of its ft copies in time (x[:, :, ft - 1:]): one frame, as on the main path. +So chunk 0 gives one video frame and every later chunk four (each upsample3d doubles them): 1 + 4(T - 1) = 4T - 3. + +Causal caches +------------- +Each causal conv with a 3-frame kernel (conv_in, conv1 and conv2 of every residual block, the two time convs, +conv_out) keeps the last two frames of its input, the activations it convolves; after a one-frame chunk, the previous +chunk's last frame and that one. The next chunk's conv sees them before its own frames, zero padding standing in for +what is missing (the first chunk has no history, a one-frame history lacks a frame). The loop thus computes, chunk by +chunk, what one causal conv over the whole sequence would, with zeros before the first frame. The 1x1x1 convs +(post_quant_conv, the residual shortcuts) need no cache. + +Over the whole sequence at once, then, the decoder runs every causal conv over all the frames of its level, two zero +frames before the first; upsample3d's time conv over the frames after the first only, the first passing through; and +DupUp3D doubling every frame but the first. That form equals the frame loop up to rounding, not bit for bit, its convs +having other shapes (fp32: max |diff| 1.7e-5, mean 6e-7 on scripts/check_reference.py's latents). + +Precision +--------- +dtype=torch.float32 computes everything in fp32, as diffusers does in fp32 (both under the same torch.backends +settings; scripts/check_reference.py turns TF32 off). +dtype=torch.bfloat16 computes what diffusers does with the VAE in bf16 under torch.autocast("cuda", dtype=bfloat16), +which is how a video pipeline decodes with this VAE in bf16: the weights and the latents are rounded to bf16 (autocast +casts the latents at post_quant_conv); every conv and the attention take bf16 operands and return bf16 (autocast's +lower-precision ops; the kernels accumulate in fp32; for head dim 1024 PyTorch runs SDPA's memory-efficient kernel); +the RMS norm normalizes in fp32 and rounds to bf16 before the products with sqrt(C) and gamma, each of which rounds +again (diffusers' upcast for 16-bit inputs); the nearest upsampling runs in fp32 (WanUpsample's upcast; it only copies +values); SiLU, the sums, DupUp3D, unpatchify and clamp run in bf16. The output is bf16. decode turns autocast off +itself: the casts are explicit. + +The operations and their order are diffusers', so that the same kernels see the same operands: on a GB200 with torch +2.12 and cuDNN 9.20, scripts/check_reference.py finds the reference bit-exact against diffusers in both modes. (The +bf16 decode itself is max 0.07, mean 0.003 from the fp32 one there.) + +Credits +------- +Derived from the decoder of diffusers' AutoencoderKLWan (src/diffusers/models/autoencoders/autoencoder_kl_wan.py, +diffusers 0.39.0), Copyright 2025 The Wan Team and The HuggingFace Team, licensed under the Apache License, Version 2.0 +(http://www.apache.org/licenses/LICENSE-2.0); diffusers' port of the Wan2.2 VAE of Wan-AI. Its structure, parameter +names and caching rules are theirs. Changes: the modules are rewritten as functions of a weight dict, the cache is +keyed by parameter name instead of a running index, the "Rep" marker is expressed as the time conv starting at the +second chunk, and the encoder, tiling and slicing are left out. +""" + +import torch +import torch.nn.functional as F + +CACHE_T = 2 # frames of history kept by a causal conv with a 3-frame kernel +PATCH_SIZE = 2 +RESNETS_PER_UP_BLOCK = 3 # num_res_blocks + 1 +# decoder.up_blocks.{i}: (in channels, out channels, temporal upsampling, spatial upsampling) +UP_BLOCKS = ((1024, 1024, True, True), (1024, 1024, True, True), (1024, 512, False, True), (512, 256, False, False)) +PREFIXES = ("decoder.", "post_quant_conv.") + + +@torch.no_grad() +def decode(z, weights, dtype=torch.float32): + """Decodes latents z [B, 48, T, H/16, W/16] into a video [B, 3, 4T-3, H, W] in [-1, 1], of `dtype`. + + `weights` maps diffusers' parameter names to tensors; those of the decoder and post_quant_conv are cast to `dtype` + on z's device (no copy when they already are). dtype is torch.float32 or torch.bfloat16 (see the module docstring). + """ + if dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"dtype must be torch.float32 or torch.bfloat16, not {dtype}") + weights = {name: t.to(z.device, dtype) for name, t in weights.items() if name.startswith(PREFIXES)} + + with torch.autocast(z.device.type, enabled=False): + x = causal_conv(z.to(dtype), weights, "post_quant_conv") + cache = {} + chunks = [decode_chunk(x[:, :, i : i + 1], weights, cache, first=i == 0) for i in range(x.shape[2])] + return unpatchify(torch.cat(chunks, dim=2), PATCH_SIZE).clamp(-1.0, 1.0) + + +def decode_chunk(x, weights, cache, first): + """The decoder on one latent frame x [B, 48, 1, h, w]: [B, 12, 1, 8h, 8w] on the first chunk, [B, 12, 4, 8h, 8w] + on the others. `cache` holds the causal convs' history, by conv name, and is updated.""" + x = cached_conv(x, weights, "decoder.conv_in", cache) + x = residual_block(x, weights, "decoder.mid_block.resnets.0", cache) + x = attention(x, weights, "decoder.mid_block.attentions.0") + x = residual_block(x, weights, "decoder.mid_block.resnets.1", cache) + for i in range(len(UP_BLOCKS)): + x = up_block(x, weights, i, cache, first) + x = F.silu(rms_norm(x, weights["decoder.norm_out.gamma"])) + return cached_conv(x, weights, "decoder.conv_out", cache) + + +def up_block(x, weights, i, cache, first): + """WanResidualUpBlock i: the residual blocks, then (up blocks 0 to 2) the resampler plus DupUp3D of x.""" + _, out_channels, temporal, spatial = UP_BLOCKS[i] + name = f"decoder.up_blocks.{i}" + block_input = x + for j in range(RESNETS_PER_UP_BLOCK): + x = residual_block(x, weights, f"{name}.resnets.{j}", cache) + if not spatial: + return x + + # The first chunk skips the time conv, whose history therefore starts empty at the second chunk. + if temporal and not first: + x = time_upsample(x, weights, f"{name}.upsampler.time_conv", cache) + x = spatial_upsample(x, weights, f"{name}.upsampler.resample.1") + return x + dup_up(block_input, out_channels, 2 if temporal else 1, 2, first) + + +def residual_block(x, weights, name, cache): + """WanResidualBlock: conv2(SiLU(norm2(conv1(SiLU(norm1(x)))))) plus x, or plus a 1x1x1 conv of x.""" + shortcut = causal_conv(x, weights, f"{name}.conv_shortcut") if f"{name}.conv_shortcut.weight" in weights else x + x = cached_conv(F.silu(rms_norm(x, weights[f"{name}.norm1.gamma"])), weights, f"{name}.conv1", cache) + x = cached_conv(F.silu(rms_norm(x, weights[f"{name}.norm2.gamma"])), weights, f"{name}.conv2", cache) + return x + shortcut + + +def attention(x, weights, name): + """WanAttentionBlock: single-head attention over the h*w positions of each frame, plus x.""" + identity = x + b, c, t, h, w = x.shape + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = rms_norm(x, weights[f"{name}.norm.gamma"]) + qkv = F.conv2d(x, weights[f"{name}.to_qkv.weight"], weights[f"{name}.to_qkv.bias"]) + qkv = qkv.reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous() # [b*t, 1 head, h*w, 3C] + q, k, v = qkv.chunk(3, dim=-1) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = F.conv2d(x, weights[f"{name}.proj.weight"], weights[f"{name}.proj.bias"]) + return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4) + identity + + +def time_upsample(x, weights, name, cache): + """upsample3d's time conv: a causal 3x1x1 conv C -> 2C; channel c becomes frame 2t, channel C + c frame 2t + 1.""" + b, c, t, h, w = x.shape + x = cached_conv(x, weights, name, cache) + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0], x[:, 1]), 3) + return x.reshape(b, c, t * 2, h, w) + + +def spatial_upsample(x, weights, name): + """Per frame: nearest x2 in h and w (in fp32, as WanUpsample does; it only copies values), then a 3x3 conv2d.""" + b, c, t, h, w = x.shape + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = F.interpolate(x.float(), scale_factor=(2.0, 2.0), mode="nearest-exact").type_as(x) + x = F.conv2d(x, weights[f"{name}.weight"], weights[f"{name}.bias"], padding=1) + return x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + +def dup_up(x, out_channels, factor_t, factor_s, first): + """DupUp3D: each channel repeated r = out * ft * fs * fs / in times, the copies spread over the (ft, fs, fs) + sub-positions of time, rows and columns; on the first chunk only the last of the ft copies in time is kept.""" + b, c, t, h, w = x.shape + x = x.repeat_interleave(out_channels * factor_t * factor_s * factor_s // c, dim=1) + x = x.view(b, out_channels, factor_t, factor_s, factor_s, t, h, w) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(b, out_channels, t * factor_t, h * factor_s, w * factor_s) + return x[:, :, factor_t - 1 :] if first else x + + +def rms_norm(x, gamma): + """WanRMS_norm over channels (dim 1): F.normalize, in fp32 for a 16-bit x and rounded back, times sqrt(C), times + gamma.""" + x_hat = F.normalize(x.float() if x.dtype in (torch.float16, torch.bfloat16) else x, dim=1).to(x.dtype) + return x_hat * x.shape[1] ** 0.5 * gamma + + +def cached_conv(x, weights, name, cache): + """causal_conv of one chunk with the history in cache[name], which it replaces by the last CACHE_T frames of its + input: after a one-frame chunk, the previous chunk's last frame and this one.""" + history = cache.get(name) + keep = x[:, :, -CACHE_T:].clone() + if keep.shape[2] < CACHE_T and history is not None: + keep = torch.cat([history[:, :, -1:], keep], dim=2) + cache[name] = keep + return causal_conv(x, weights, name, history) + + +def causal_conv(x, weights, name, history=None): + """WanCausalConv3d `name` on x [B, C, t, h, w]: zero padding of kh//2 and kw//2 on both sides in space and, in + time, 2 * (kt//2) frames before x: `history` (the frames before x, at most that many), then zeros.""" + weight, bias = weights[f"{name}.weight"], weights[f"{name}.bias"] + kt, kh, kw = weight.shape[2:] + pad = [kw // 2, kw // 2, kh // 2, kh // 2, 2 * (kt // 2), 0] + if history is not None and pad[4] > 0: + x = torch.cat([history, x], dim=2) + pad[4] -= history.shape[2] + return F.conv3d(F.pad(x, pad), weight, bias) + + +def unpatchify(x, p): + """[B, 3*p*p, F, h, w] -> [B, 3, F, p*h, p*w]: channel (k*p + j)*p + i is colour k at row offset i, column j.""" + b, cp, f, h, w = x.shape + x = x.view(b, cp // (p * p), p, p, f, h, w) + x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() + return x.view(b, cp // (p * p), f, h * p, w * p) diff --git a/fastvideo-kernel/wan22_vae_decode/solution/attn.cu b/fastvideo-kernel/wan22_vae_decode/solution/attn.cu new file mode 100644 index 000000000..4559bc4e3 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/attn.cu @@ -0,0 +1,217 @@ +// Mid-block self-attention with our own kernels, replacing ATen's +// scaled_dot_product_attention (fmha_cutlassF, 8.5 ms of f121x1's 954 ms). +// +// The reference computes, per latent frame, one head of softmax(q k^T/32) v over +// the H*W=3432 positions with C=1024 channels and no mask. ATen's path cannot use +// flash attention here (head dim 1024 > 256) and falls back to the xformers +// cutlass kernel, which reads and writes far more per frame than the products +// need: the decode's whole attention is 1.5 TFLOP (0.63 ms at 16-bit peak for +// f121x1) against 8.5 ms measured, ~13x off. +// +// Three stages, each at its own bound: +// 1. pack: q,k into contiguous [HW,C] per frame and v transposed into the +// [C,PAD] weight layout the convolution GEMM wants (~21 MB/frame). +// 2. S = q k^T per frame (fp32 out) and O = p v per frame (bf16 out), both as +// 1x1x1 implicit GEMMs through the vendored CUTLASS SM100 convolution +// collective (2-CTA tcgen05 MMA, TMA, fp32 accumulation) -- the same +// machinery conv.cu uses, so no new GEMM path is introduced. +// 3. row softmax: S row -> exp(S/sqrt(C) - max)/sum in bf16, one CTA per row, +// the row held in shared memory (one read of S, one write of P). +// P's rows are padded to a multiple of 64 columns and v's transpose is zero +// there, so the GEMM's tail contributes nothing. +#include +#include +#include +#include +#include "attn.h" + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/convnd_problem_shape.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/conv/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/conv/device/conv_universal_adapter.hpp" +#include "cutlass/conv/kernel/conv_universal.hpp" + +using namespace cute; +using ElementBF16 = cutlass::bfloat16_t; +using LayoutNDHWC = cutlass::layout::TensorNDHWC; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassTensorOp; +constexpr auto Fprop = cutlass::conv::Operator::kFprop; + +// One 1x1x1 implicit GEMM, y = x (per position) times w, fp32 accumulation and +// ElementD out. x: channels-last3d [1,C,D,1,1] (may be a strided view); w: +// [K,C,1,1,1] with explicit strides; y: packed channels-last3d [1,K,D,1,1]. +template +static void gemm1x1(const at::Tensor& x, const at::Tensor& w, const at::Tensor& y) { + using Tile = Shape, Int, Shape<_64>>; + using Cluster = Shape, _1, _1>; + using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto, + float, float, ElementD, LayoutNDHWC, 8, ElementD, LayoutNDHWC, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + using Mainloop = typename cutlass::conv::collective::CollectiveBuilder< + ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float, + Tile, Cluster, + cutlass::conv::collective::StageCountAutoCarveout(sizeof(typename Epilogue::SharedStorage))>, + cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp; + using Problem = cutlass::conv::ConvProblemShape; + using Kernel = cutlass::conv::kernel::ConvUniversal; + using Conv = cutlass::conv::device::ConvUniversalAdapter; + int C = x.size(1), D = x.size(2); + int K = w.size(0); + TORCH_CHECK(y.size(1) == K && y.size(2) == D, "attn gemm: output shape mismatch"); + Problem problem(cutlass::conv::Mode::kCrossCorrelation, + {1, D, 1, 1, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)}, + {K, 1, 1, 1, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)}, + {0, 0, 0}, {0, 0, 0}, {1, 1, 1}, {1, 1, 1}, 1); + typename Kernel::StrideC sc{}; + typename Kernel::StrideD sd{}; + cute::for_each(cute::make_seq(sc)>{}, [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; }); + cute::for_each(cute::make_seq(sd)>{}, [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; }); + typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()}, + {{1.f, 0.f}, (const ElementD*)y.data_ptr(), sc, (ElementD*)y.data_ptr(), sd}}; + Conv conv; + TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess, + "attn gemm: CUTLASS cannot implement D=", D, " N=", K, " C=", C); + size_t size = Conv::get_workspace_size(args); + auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte)); + auto stream = at::cuda::getCurrentCUDAStream(); + TORCH_CHECK(conv.initialize(args, workspace.data_ptr(), stream) == cutlass::Status::kSuccess, + "attn gemm: initialization failed"); + TORCH_CHECK(conv.run(stream) == cutlass::Status::kSuccess, "attn gemm: run failed"); +} + +// Channels-last-3d views with explicit strides: a size-1 dim built by unsqueeze +// carries stride 1, which the im2col TMA descriptor rejects. +static at::Tensor nhwc_in(const at::Tensor& m, int64_t C, int64_t D) { // [D,C] -> [1,C,D,1,1] + return m.as_strided({1, C, D, 1, 1}, {C * D, 1, C, C, C}); +} +static at::Tensor nhwc_w(const at::Tensor& m, int64_t K, int64_t C) { // [K,C] -> [K,C,1,1,1] + return m.as_strided({K, C, 1, 1, 1}, {C, 1, C, C, C}); +} +static at::Tensor nhwc_out(const at::Tensor& m, int64_t K, int64_t D) { // [D,K] -> [1,K,D,1,1] + return m.as_strided({1, K, D, 1, 1}, {K * D, 1, K, K, K}); +} + +// v [T,1,HW,3C] at offset 2C -> vt [T,C,PAD] (transposed, [T,C,PAD] zeroed first). +static __global__ void transpose_v_kernel(const __nv_bfloat16* __restrict__ qkv, int T, int HW, int C, int PAD, + __nv_bfloat16* __restrict__ vt) { + __shared__ __nv_bfloat16 tile[32][33]; + int i0 = blockIdx.x * 32, c0 = blockIdx.y * 32, t = blockIdx.z; + { + int i = i0 + threadIdx.y, c = c0 + threadIdx.x; + if (i < HW && c < C) tile[threadIdx.y][threadIdx.x] = qkv[((size_t)t * HW + i) * (3 * C) + 2 * C + c]; + } + __syncthreads(); + { + int i = i0 + threadIdx.x, c = c0 + threadIdx.y; + if (i < HW && c < C) vt[((size_t)t * C + c) * PAD + i] = tile[threadIdx.x][threadIdx.y]; + } +} + +// One CTA per row: p = exp(s*scale - max)/sum in bf16, then zero the padded tail. +static __global__ void row_softmax_kernel(const float* __restrict__ s, int cols, int PAD, float scale, + __nv_bfloat16* __restrict__ p) { + extern __shared__ float row[]; + const int n4 = cols / 4; + const float4* src = (const float4*)(s + (size_t)blockIdx.x * cols); + __nv_bfloat16* dst = p + (size_t)blockIdx.x * PAD; + float4* r4 = (float4*)row; + for (int j = threadIdx.x; j < n4; j += blockDim.x) { + float4 v = src[j]; + v.x *= scale; + v.y *= scale; + v.z *= scale; + v.w *= scale; + r4[j] = v; + } + __syncthreads(); + float m = -INFINITY; + for (int j = threadIdx.x; j < n4; j += blockDim.x) { + float4 v = r4[j]; + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + __shared__ float red[32]; + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = m; + __syncthreads(); + if (threadIdx.x < 32) { + m = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : -INFINITY; + for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + if (threadIdx.x == 0) red[0] = m; + } + __syncthreads(); + m = red[0]; + __syncthreads(); + float sum = 0.f; + for (int j = threadIdx.x; j < n4; j += blockDim.x) { + float4 v = r4[j]; + v.x = __expf(v.x - m); + v.y = __expf(v.y - m); + v.z = __expf(v.z - m); + v.w = __expf(v.w - m); + r4[j] = v; + sum += (v.x + v.y) + (v.z + v.w); + } + for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off); + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = sum; + __syncthreads(); + if (threadIdx.x < 32) { + sum = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : 0.f; + for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off); + if (threadIdx.x == 0) red[0] = sum; + } + __syncthreads(); + const float inv = 1.f / red[0]; + const int npair = cols / 2; + for (int j = threadIdx.x; j < npair; j += blockDim.x) + *(__nv_bfloat162*)(dst + 2 * j) = __floats2bfloat162_rn(row[2 * j] * inv, row[2 * j + 1] * inv); + for (int j = cols + threadIdx.x; j < PAD; j += blockDim.x) dst[j] = __float2bfloat16(0.f); +} + +at::Tensor attn_forward(const at::Tensor& qkv) { + TORCH_CHECK(qkv.dim() == 4 && qkv.size(1) == 1 && qkv.size(3) % 3 == 0, "attn: expected [T,1,HW,3C]"); + int T = qkv.size(0), HW = qkv.size(2), C = qkv.size(3) / 3; + TORCH_CHECK(C % 64 == 0, "attn: channels must be a multiple of 64"); + int PAD = (HW + 63) / 64 * 64; + auto stream = at::cuda::getCurrentCUDAStream(); + auto bf = qkv.options(); + auto vt = at::zeros({T, C, PAD}, bf); + auto s = at::empty({T, HW, HW}, bf.dtype(at::kFloat)); + auto p = at::empty({T, HW, PAD}, bf); + auto o = at::empty({T, 1, HW, C}, bf); + const __nv_bfloat16* qkv_p = (const __nv_bfloat16*)qkv.data_ptr(); + + transpose_v_kernel<<>>( + qkv_p, T, HW, C, PAD, (__nv_bfloat16*)vt.data_ptr()); + + // S = q k^T : M = positions, N = keys, K = channels. q and k are read in place: + // their channel stride is 1 and their position stride 3C, which TMA accepts. + for (int t = 0; t < T; ++t) { + auto qv = qkv.select(0, t).select(0, 0); // [HW, 3C] + auto x = qv.narrow(1, 0, C).as_strided({1, C, HW, 1, 1}, {(int64_t)3 * C * HW, 1, 3 * C, 3 * C, 3 * C}); + auto w = qv.narrow(1, C, C).as_strided({HW, C, 1, 1, 1}, {3 * C, 1, 3 * C, 3 * C, 3 * C}); + auto y = nhwc_out(s.select(0, t), HW, HW); + gemm1x1(x, w, y); + } + + float* s_p = (float*)s.data_ptr(); + float scale = 1.f / sqrtf((float)C); + row_softmax_kernel<<>>( + s_p, HW, PAD, scale, (__nv_bfloat16*)p.data_ptr()); + + // O = p v : M = positions, N = channels, K = keys (padded). + for (int t = 0; t < T; ++t) { + auto x = nhwc_in(p.select(0, t), PAD, HW); + auto w = nhwc_w(vt.select(0, t), C, PAD); + auto y = nhwc_out(o.select(0, t).select(0, 0), C, HW); + gemm1x1(x, w, y); + } + return o; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/attn.h b/fastvideo-kernel/wan22_vae_decode/solution/attn.h new file mode 100644 index 000000000..de1df1567 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/attn.h @@ -0,0 +1,6 @@ +#pragma once +#include +// One head of self-attention over the H*W positions of every latent frame of +// one decoder chunk: qkv is [T,1,HW,3C] bf16 (q, k, v concatenated on the last +// dim, as binding.cpp's attention() builds it); the result is [T,1,HW,C] bf16. +at::Tensor attn_forward(const at::Tensor& qkv); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/binding.cpp b/fastvideo-kernel/wan22_vae_decode/solution/binding.cpp new file mode 100644 index 000000000..a8c59c67b --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/binding.cpp @@ -0,0 +1,459 @@ +// Wan2.2 VAE decoder seed. Derived from the reference architecture and +// indexing in diffusers' AutoencoderKLWan (Wan Team / Hugging Face Team, +// Apache-2.0; see WAN_LICENSE.txt). Native CUDA execution is implemented here. +// All invocation state, including temporal history, belongs to Decoder below. +#include +#include +#include +#include +#include +#include +#include "kernels.h" +#include "prefix.h" +#include "comm.h" +#include "resample.h" +#include "attn.h" +#include "head.h" +#include +#include +#include +#include + +at::Tensor custom_conv(const at::Tensor&,const at::Tensor&,int64_t,int64_t,int64_t); + +// The temporal upsampler, writing the frame-interleaved video directly (solution/timeconv.cu). +at::Tensor time_upsample(const at::Tensor&,const at::Tensor&,const at::Tensor&,const at::Tensor&,int64_t,int64_t,bool); + +using PG = c10::intrusive_ptr; +using at::Tensor; +using Weights = std::vector; + +Weights prepare(const Weights& weights,int64_t rank,int64_t world,const PG& pg) { + TORCH_CHECK(weights.size()==110,"Expected the 110 decoder tensors"); + at::NoGradGuard guard; + Weights result; + result.reserve(weights.size()); + for(const auto& w:weights) { + // A 1x1x1 tensor is "trivially channels-last", so contiguous() leaves it with unit strides on + // its size-1 spatial dims; CUTLASS's im2col TMA needs those strides to be 16-byte multiples. + // The same bytes with dense strides: as_strided keeps the NDHWC order but gives every spatial + // dim the channel-count stride. + if(w.dim()==5 && w.size(2)==1 && w.size(3)==1 && w.size(4)==1) { + auto wc=w.contiguous();int64_t C=wc.size(1); + result.push_back(wc.as_strided({wc.size(0),C,1,1,1},{C,1,C,C,C})); + } + else if(w.dim()==5) result.push_back(w.contiguous(at::MemoryFormat::ChannelsLast3d)); + else if(w.dim()==4 && w.size(2)>0 && w.size(1)>1) + result.push_back(w.unsqueeze(2).contiguous(at::MemoryFormat::ChannelsLast3d)); + else result.push_back(w.contiguous()); + } + // The output head: conv_out's nine spatial taps become the output channels of a 3x1x1 + // (temporal-only) convolution, so its 112 outputs are (tap i,j) x (channel c) and head.cu's + // epilogue sums the nine shifted copies. W'[(i*3+j)*12+c, ch, t] = W[c, ch, t, i, j]; the four + // channels past 108 are zero and are never read. N=112 rather than 108 keeps the epilogue's + // channel count a multiple of 16. + { + auto orig=result[108]; + auto t=orig.permute({3,4,0,1,2}).contiguous(); // [i,j,c,ch,t] = [3,3,12,256,3] + auto filled=at::zeros({112,256,3},orig.options()); // [n,ch,t], n=(i*3+j)*12+c + filled.narrow(0,0,108).copy_(t.reshape({108,256,3})); + auto kt=filled.transpose(1,2).contiguous(); // [n,t,ch] contiguous + // The layout a [K,C,3,3,3] weight has under ChannelsLast3d: C innermost with stride 1, every + // other stride a multiple of 16 bytes. torch normalises a size-1 H/W away, hence as_strided. + result[108]=kt.transpose(1,2).as_strided({112,256,3,1,1},{768,1,256,256,256}); + } + // Nearest-x2 + 3x3 convolution equals four 2x2 low-resolution phases. + // Form the sums in fp32 and store the transformed weights in bf16. + for(int spatial : {39,61,85}) { + auto orig=result[spatial].to(at::kFloat); + int C=orig.size(0); + for(int a=0;a<2;++a)for(int b=0;b<2;++b) { + auto transformed=at::zeros({C,C,1,2,2},orig.options()); + for(int u=0;u<2;++u)for(int v=0;v<2;++v) { + auto sum=at::zeros({C,C},orig.options()); + for(int r=0;r<3;++r)for(int c=0;c<3;++c) { + int dy=a+r-1,dx=b+c-1; + int iy=dy<0?-1:dy/2,ix=dx<0?-1:dx/2; + if(iy==u+a-1&&ix==v+b-1)sum.add_(orig.select(2,0).select(2,r).select(2,c)); + } + transformed.select(2,0).select(2,u).select(2,v).copy_(sum); + } + result.push_back(transformed.to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d)); + } + } + // F(4x4,3x3) Winograd weights for up blocks 1, 2, and 3. + // The 36 components each retain the causal 3-tap temporal convolution. + const float G[6][3]={{.25f,0,0},{-1.f/6,-1.f/6,-1.f/6},{-1.f/6,1.f/6,-1.f/6},{1.f/24,1.f/12,1.f/6},{1.f/24,-1.f/12,1.f/6},{0,0,1}}; + for(int idx : {44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105}) { + auto orig=result[idx].to(at::kFloat); + for(int a=0;a<6;a++)for(int b=0;b<6;b++) { + auto sum=at::zeros({orig.size(0),orig.size(1),3},orig.options()); + for(int r=0;r<3;r++)for(int c=0;c<3;c++)if(G[a][r]*G[b][c]!=0) + sum.add_(orig.select(3,r).select(3,c),G[a][r]*G[b][c]); + result.push_back(sum.unsqueeze(3).unsqueeze(4).to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d)); + } + } + // The 1x1x1 shortcut of a channel-changing resnet is added to the second convolution's + // bias + output, out = round(acc2 + b2) + round(acc1 + b1). Fold b1 into b2 instead (one more + // bf16 rounding of the bias sum and none of the operand): the shortcut tensor then needs no + // separate bias pass, which was a full read+write of it. + for(auto pair : {std::pair{72,70},std::pair{94,92}}) { + auto sum=result[pair.second].to(at::kFloat)+result[pair.first].to(at::kFloat); + result[pair.second]=sum.to(at::kBFloat16).contiguous(); + } + return result; +} + +struct Decoder { + // This must not implicitly bind to a fused-epilogue weight index when + // independently extended conv() interfaces are merged. + struct TemporalTrim { + int frames; + explicit constexpr TemporalTrim(int n) : frames(n) {} + }; + + const Weights& w; + int rank,world; + PG pg; + std::array cache; + bool sharded=false; + // A convolution whose input was already normalized, activated and padded by the previous + // Winograd convolution's fused epilogue: (tensor, the conv index it is prepared for, and the + // sharding it was prepared under — a mismatch there would be silently wrong, not a crash). + Tensor prepared; int prepared_for=-1; bool prepared_sharded=false; + // The conv index whose y store was suppressed because this conv reserved it for `prepared`. + // It is the RUNTIME half of the y_dead contract: the producer states an intent and this records + // it, so a later edit that stops the consumer taking the `pre` edge aborts instead of handing a + // never-written tensor to norm_pad. Cleared when the edge is taken, and asserted closed at the + // end of every chunk beside prepared_for, because a conv index restarts at the same values. + int y_suppressed_for=-1; + Decoder(const Weights& weights,int r,int n,const PG& p):w(weights),rank(r),world(n),pg(p){} + + Tensor conv(Tensor x,int idx,int gamma=-1,bool activate=true,bool add_bias=true,const Tensor& residual={},TemporalTrim trim=TemporalTrim{0},int fg=-1,int fnext=-1,bool y_dead=false) { + RECORD_FUNCTION("wan.conv."+std::to_string(idx), std::vector()); + const int temporal_trim=trim.frames; + int kt=w[idx].size(2),kh=w[idx].size(3),kw=w[idx].size(4); + TORCH_CHECK(temporal_trim>=0 && temporal_trim1&&kh==3)?1:0; + Tensor p=x; + bool pre=(prepared_for==idx); + // The consumer asserts the producer's y_dead contract AT RUNTIME, not in prose. If a conv had + // its y store suppressed for this idx and this conv does not take the fused `prepared` edge, + // it is about to read a tensor that was never written -- so abort instead of producing a wrong + // video. The existing metadata checks at `pre` cannot catch this, because they live inside the + // branch that is not being taken. + TORCH_CHECK(pre || y_suppressed_for!=idx, + "conv ",idx,": its y store was suppressed FOR THIS CONV, which then did not take " + "the prepared edge; the tensor it is about to read is unwritten"); + // Prefix pruning never uses the fused, full-length prepared-input edge. + // A future caller must implement a trimmed prepared layout before mixing them. + TORCH_CHECK(temporal_trim==0 || (!pre && fg<0 && fnext<0), + "Causal prefix trim cannot use or produce a fused prepared input"); + if(pre) { + TORCH_CHECK(prepared.defined(),"conv ",idx," is marked prepared but no tensor was stored"); + TORCH_CHECK(prepared_sharded==sharded,"prepared tensor crosses a change of sharding"); + // The graph edge checks itself. pt/ph here are the consumer's own, read from this conv's + // weight, and the producer computed the same two from w[fnext] with fnext==idx -- so these + // equalities are true by construction today and are a proof obligation on any later edit: + // a wrong fnext would otherwise hand a wrong-shaped tensor to the convolution, and a + // wrong-but-same-shaped one would pass every check the kernel makes. + TORCH_CHECK(prepared.size(1)==x.size(1) && prepared.size(2)==x.size(2)+pt && + prepared.size(3)==x.size(3)+2*ph && prepared.size(4)==x.size(4), + "conv ",idx,": prepared ",prepared.sizes()," does not fit input ",x.sizes(), + " (pt ",pt,", ph ",ph,", sharded ",(int)sharded,")"); + p=prepared; prepared=Tensor(); prepared_for=-1; prepared_sharded=false; + y_suppressed_for=-1; // the suppressed y was reserved for this edge and this edge took it + } + else if(pt||ph||gamma>=0) p=norm_pad(x,gamma>=0?w[gamma]:Tensor(),cache[idx],pt,ph,activate); + // Memory: the unpadded input is dead once the padded copy exists (p holds it when no copy was made). + x=Tensor(); + // The causal cache carries this convolution's input history to the next chunk. `p` may be the + // tensor the previous fused epilogue produced and this conv has just consumed, so the copy is + // ordered after that kernel by the stream they share -- a side stream here would race it. + if(pt) cache[idx]=last_frames(p,pt,ph); + if(ph) { + RECORD_FUNCTION("wan.halo",std::vector()); + auto recv=halo_exchange(halo_pack(p),pg,rank,world); + halo_unpack(p,recv,rank,world); + } + const int Ho=p.size(3)-2*ph,Wo=p.size(4); + std::vector padding={0,ph?0:kh/2,kw/2}; + bool use_custom=(p.size(4)==78 && kt==3 && kh==3 && w[idx].size(0)==1024 && w[idx].size(1)==1024) || idx==61 || idx==85 + || idx==71 || idx==93; + // Only the expensive spatial 3x3 convolutions use Winograd. Low-resolution + // convolutions remain direct implicit GEMMs; 1x1 and temporal-only layers + // use cuDNN. Each transformed GEMM accumulates bf16 operands in fp32. + int wi=-1; + const int wids[18]={44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105}; + for(int i=0;i<18;i++)if(idx==wids[i])wi=i; + Tensor y; + if(wi>=0) { + auto transformed=winograd_input_fast(p,ph); + // Memory: the padded input is dead once transformed (its history went to the cache above), and the transform + // (2.25x the input) once every component product is enqueued; the caching allocator reuses both in stream order. + p=Tensor(); + std::vector products; + for(int a=0;a<36;a++) { + auto input=transformed.select(0,a).unsqueeze(0).permute({0,4,1,2,3}); + products.push_back(custom_conv(input,w[122+36*wi+a],0,0,3)); + } + transformed=Tensor(); + // The SiLU this epilogue applies is the CONSUMER's norm activation, not this conv's + // `activate`: fusing is valid only when the consumer would have activated. Every fused + // consumer in the decoder (the two convs of a residual block) has activate=true, so the + // test below holds by construction; a non-activating consumer would have to pass its own. + bool fuse=(fg>=0&&fnext>=0&&activate); + // write_y is false exactly where the caller asserts nobody reads this launch's y AND the + // fused `prepared` edge exists to carry the activation instead. It is recorded so the + // consumer can ASSERT the edge rather than trust it; see y_suppressed_for above. + const bool write_y=!(y_dead && fuse); + int pt2=pt,ph2=ph; + if(fuse) { pt2=w[fnext].size(2)-1; ph2=(sharded&&world>1&&w[fnext].size(3)==3)?1:0; } + auto out=wino_epilogue(products,Ho,Wo,add_bias?w[idx+1]:Tensor(),residual, + fuse?w[fg]:Tensor(),fuse?cache[fnext]:Tensor(),pt2,ph2,write_y); + y=out.y; + if(fuse) { prepared=out.n; prepared_for=fnext; prepared_sharded=sharded; + if(!write_y) y_suppressed_for=fnext; } + } else y=use_custom?custom_conv(p,w[idx],padding[1],padding[2],3):at::cudnn_convolution(p,w[idx],padding,{1,1,1},{1,1,1},1,false,false,false); + if(add_bias && wi<0) finish(y,w[idx+1],residual); + return y; + } + + Tensor resnet(Tensor x,int base,bool change=false,int nx=-1,int nxg=-1) { + RECORD_FUNCTION("wan.resnet."+std::to_string(base),std::vector()); + Tensor skip=change?conv(x,base+6,-1,true,false):x; + // Memory: moved into the first convolution, the block's input is released there once its norm has read it + // (skip still holds it when it is the residual). + auto y=conv(std::move(x),base+1,base,true,true,Tensor(),TemporalTrim{0},base+3,base+4, /*y_dead=*/true); + return conv(y,base+4,base+3,true,true,skip,TemporalTrim{0},nxg,nx); + } + + Tensor attention(const Tensor& x) { + RECORD_FUNCTION("wan.attention",std::vector()); + int T=x.size(2),H=x.size(3),W=x.size(4),C=x.size(1); + Tensor qkv=conv(x,5,4,false); + auto p=qkv.permute({0,2,3,4,1}).view({T,1,H*W,3*C}); + Tensor y=attn_forward(p); + y=y.view({1,T,H,W,C}).permute({0,4,1,2,3}); + return conv(y,7,-1,true,true,x); + } + + // The output head: RMS norm + SiLU + the causal padding, the 3x1x1 convolution to 112 channels + // (108 real) and the 3x3 stencil with the bias and clamp (head.cu). Replaces conv(x,108,107), + // whose 12 output channels waste 244 of a 256-wide tile's columns. + Tensor output_head(const Tensor& x) { + RECORD_FUNCTION("wan.head",std::vector()); + const int pt=2; + int ph=world>1?1:0; + Tensor p=norm_pad(x,w[107],cache[108],pt,ph,true); + if(ph) { + RECORD_FUNCTION("wan.halo",std::vector()); + // The third halo site. PR204 deleted pack_edges/unpack_edges, and the resolution that + // compiles-and-passes while silently dropping the vectorised pack is to re-add them: + // this is the same two pad rows through comm.cu's halo_pack/halo_unpack, which is what + // the other two sites (conv/resnet) already use. Census: post_edges == wait_edges == 3. + auto recv=halo_exchange(halo_pack(p),pg,rank,world); + halo_unpack(p,recv,rank,world); + } + cache[108]=last_frames(p,pt,ph); + Tensor g=conv_n128(p,w[108]); + int rows=x.size(3); + auto out=at::empty({1,12,g.size(2),rows,g.size(4)},x.options().memory_format(at::MemoryFormat::ChannelsLast3d)); + // 352 is the pixel-plane height (44 latent rows x 8) and lo this rank's first pixel-plane row; + // both are the fixed 704x1248 geometry, and lo must stay 8*(44*rank/world) because run() + // passes lo*16 - the same row on the 704-tall video grid - to write_video. + head_stencil(g,w[109],out,352,ph,8*(44*rank/world)); + return out; + } + + Tensor up(Tensor input,int block,bool first) { + RECORD_FUNCTION("wan.up_block."+std::to_string(block),std::vector()); + const int bases[4]={21,43,65,87}; + int base=bases[block]; + bool change=block>=2; + int b2=base+(change?8:6),b3=base+(change?14:12); + // Memory: only the resampler's DupUp3D shortcut reads the block's input again, and block 3 has none. + auto x=resnet(block==3?std::move(input):input,base,change,b2+1,b2); + x=resnet(x,b2,false,b3+1,b3); + x=resnet(x,b3,false); + if(block==3) return x; + int spatial=base+(change?20:18); + if(block<2) { + // The 3x1x1 causal time convolution writes the frame-interleaved tensor directly (a stride-2 D epilogue + // in solution/timeconv.cu), so `up_time` and `finish` are gone from this seam. `activate=false` is + // INERT: with no gain the dispatcher takes norm.cu's copy-only mode and never reads the flag. + int start=first?1:0; + Tensor head=first?x.narrow(2,0,1):Tensor(); + int nf=x.size(2)-start; + Tensor p; + if(nf>0) { + p=norm_pad(x.narrow(2,start,nf),Tensor(),cache[spatial+2],2,0,false); + cache[spatial+2]=last_frames(p,2,0); + } + x=time_upsample(p,head,w[spatial+2],w[spatial+3],x.size(2),nf,first); + } + if(world>1) { + x=norm_pad(x,Tensor(),Tensor(),0,1,false); + auto recv=halo_exchange(halo_pack(x),pg,rank,world); + halo_unpack(x,recv,rank,world); + } + // Four-phase spatial upsampler: the four 2x2 phase convolutions, then one bandwidth-bound pass that + // interleaves them with the bias and the DupUp3D shortcut (resample.cu). + std::vector pw; + for(int a=0;a<2;++a)for(int b=0;b<2;++b) pw.push_back(w[110+4*block+2*a+b]); + return resample_upsample(x,pw,w[spatial+1],input,world>1,block<2,first); + } + + // Stateless exact prefix interval. Three causal convolutions require six + // preceding latent frames, but the support is pruned by two after each conv. + Tensor prefix_interval(const Tensor& z,int begin,int count) { + sharded=false; + const int end=begin+count; + const int lo6=std::max(0,begin-6),lo4=std::max(0,begin-4),lo2=std::max(0,begin-2); + auto raw=z.narrow(2,lo6,end-lo6).contiguous(at::MemoryFormat::ChannelsLast3d); + auto x=conv(raw,0); + // Reducing left temporal padding prunes the invalid outputs without ever + // computing them. Near global frame zero, the remaining padding is real. + x=conv(x,2,-1,true,true,Tensor(),TemporalTrim{lo4-lo6}); + auto y=conv(x,10,9,true,true,Tensor(),TemporalTrim{lo2-lo4}); + // Batch is one and T is the outer physical NDHWC dimension, so this + // frame slice remains dense for finish()'s linear residual indexing. + auto skip=x.narrow(2,begin-lo4,count); + TORCH_CHECK(skip.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix residual must be dense NDHWC"); + y=conv(y,13,12,true,true,skip,TemporalTrim{begin-lo2}); + return attention(y); + } + + Tensor distributed_prefix(const Tensor& z,int start,int count) { + // No persistent prefix history: temporal ownership changes from batch to + // batch. The full latents on each rank provide exact predecessor context. + Decoder prefix(w,rank,world,pg); + Tensor x; + if(count>=2*world) { + const int lo=count*rank/world,hi=count*(rank+1)/world; + x=prefix.prefix_interval(z,start+lo,hi-lo); + x=prefix_temporal_to_rows(x,count,rank,world,pg); + } else { + // Tiny clips/tails do not amortize the transpose, and may have empty + // temporal owners. A stateless replicated prefix remains exact. + x=prefix.prefix_interval(z,start,count); + const int lo=44*rank/world,hi=44*(rank+1)/world; + x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d); + } + return x; + } + + Tensor tail(Tensor x,bool first) { + // Every prepared edge is consumed within its residual/up-block chain. + // Tail history survives chunks, but no prepared current-chunk tensor may. + TORCH_CHECK(prepared_for<0 && !prepared.defined() && y_suppressed_for<0, + "Prepared input crossed a tail-chunk boundary"); + sharded=true; + x=resnet(x,15); + for(int i=0;i<4;++i) x=up(std::move(x),i,first); + return output_head(x); + } + + Tensor chunk(const Tensor& z,bool first) { + sharded=false; + // The invariant tail() checks must be CHECKED here too, not cleared. Clearing it first makes + // tail()'s check vacuous on the only path world 1 takes, and a leaked fused tensor is not + // merely retained: prepared_for is a conv INDEX and every chunk restarts at the same indices, + // so the next chunk's conv of that index CONSUMES it (`pre`) and the video is silently wrong + // at world 1 while world 4 errors out. + // + // The reset was dead on this tree, so the check costs one compare per chunk and nothing else, + // but the closure argument is NOT "every fused pair lives inside one resnet()" -- that is + // false. resnet(b)'s second conv is called with (nxg=b2, nx=b2+1) and produces for + // fnext=nx=b2+1, which is the FIRST conv of the next resnet; the same edge appears at every + // up block (47->50, 53->56 at block 1; 69->74, 77->80 at block 2). The chain closes at the + // END OF THE UP BLOCK instead: up() runs resnet(base, nxg=b2, nx=b2+1), then + // resnet(b2, nxg=b3, nx=b3+1), then resnet(b3) with the DEFAULT nx=-1, so the third resnet's + // last conv has fg=fnext=-1, produces nothing, and no fused tensor leaves up(). up() is the + // only Winograd caller and nothing after it is Winograd (output_head goes through conv_n128), + // so the triple is clear at every chunk entry. + // The predicate is tail()'s, EXACTLY, and that identity is the point rather than a + // coincidence: chunk() must accept everything tail() accepts, or the subset property this + // change exists to establish is not established. A third conjunct would also be unreachable -- + // all three of prepared / prepared_for / prepared_sharded are written together at the single + // producer site (conv()'s `if(fuse)`) and cleared together at the single consumption site, so + // `prepared_sharded` can never differ from `prepared_for` and cannot fail on its own. + TORCH_CHECK(prepared_for<0 && !prepared.defined(), + "Prepared input crossed a chunk boundary"); + auto x=conv(z,0); + x=conv(x,2); + x=resnet(x,9); + x=attention(x); + if(world>1) { + int lo=44*rank/world,hi=44*(rank+1)/world; + x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d); + } + return tail(x,first); + } +}; + +void decode_video(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out, + int64_t frames_per_pass) { + at::NoGradGuard no_grad; + c10::cuda::CUDAGuard device_guard(z.device()); + TORCH_CHECK(z.scalar_type()==at::kBFloat16&&z.dim()==5&&z.size(0)==1&&z.size(1)==48&&z.size(2)>=1&&z.size(3)==44&&z.size(4)==78,"Expected bf16 [1,48,T,44,78]"); + TORCH_CHECK(world>=1&&world<=8&&rank>=0&&rank=1,"frames_per_pass must be at least 1, not ",frames_per_pass); + halo_stats()=HaloStats{}; // per call: the debug tail below prints THIS run's exchange, not the process's + Decoder decoder(weights,rank,world,pg); + // A pass carries frames_per_pass latent frames per GPU, frames_per_pass * world in all. The working memory of a GPU + // grows with it (3.5-4.2 GiB per latent frame at 704 x 1248), and so does the speed: a pass's GEMMs at 44 x 78 and + // 88 x 156 fill the GPU only with enough frames, and above world 1 every pass also pays a fixed cost (its prefix + // and its halo exchanges). The causal history crosses passes in `cache`, so the video does not depend on the size. + // At most 61 latent frames: the largest pass this tree was measured with (wino_epilogue.cu's indices are 64-bit). + static const int64_t kMaxChunk=61; + const int chunk_size=(int)std::min(frames_per_pass*world,kMaxChunk); + int T=z.size(2); + for(int start=0;start1) pixels=decoder.tail(decoder.distributed_prefix(z,start,count),start==0); + else { + auto input=z.narrow(2,start,count).contiguous(at::MemoryFormat::ChannelsLast3d); + pixels=decoder.chunk(input,start==0); + } + int frame_offset=start==0?0:4*start-3; + if(world==1) write_video(pixels,out,frame_offset,0); + else if(rank==0) { + write_video(pixels,out,frame_offset,0); + for(int peer=1;peer buffers={part}; + pg->recv(buffers,peer,0)->wait(); + write_video(part,out,frame_offset,lo*16); + } + } else { + std::vector buffers={pixels}; + pg->send(buffers,0,0)->wait(); + } + } + // Debug only, and never set by CI: cached once per process so the measured span runs no getenv. + static const bool show_halo_stats=std::getenv("WAN_HALO_STATS")!=nullptr; + if(show_halo_stats) { + const HaloStats& s=halo_stats(); + printf("[halo rank %d/%d] exchanges %lld sent %.1f MB received %.1f MB\n",(int)rank,(int)world, + (long long)s.exchanges,1e-6*s.bytes_sent,1e-6*s.bytes_received); + fflush(stdout); + } +} + +// The default pass: 4 latent frames per GPU. On one GB200 at 704 x 1248 it peaks at 21-22 GiB allocated and takes +// 1.09x the time of 16 frames per GPU, which peaks at 60-61 GiB (8 to 20 s of video). +constexpr int64_t kFramesPerPass=4; + +void run(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out) { + decode_video(z,weights,rank,world,pg,out,kFramesPerPass); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) { + m.def("prepare",&prepare); + m.def("run",&run,py::call_guard()); + m.def("run_with",&decode_video,py::call_guard()); +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/comm.cu b/fastvideo-kernel/wan22_vae_decode/solution/comm.cu new file mode 100644 index 000000000..7abe62072 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/comm.cu @@ -0,0 +1,213 @@ +// Halo edge transport: neighbour-only point to point and the seed's ring all-gather, over an +// eight-channels-per-thread edge pack/scatter. Which arm runs is use_allgather's world rule below -- +// the ring only at world 2, neighbour-only at world 4 and at every world above it -- which WAN_HALO +// can force either way for a measurement. +// THIS POLICY IS STATED IN THREE PLACES AND THEY MUST AGREE: this file header, the block comment on +// use_allgather, and the WAN_HALO contract in comm.h. An earlier revision of this change corrected one +// of the three and was rightly blocked on the other two, so the cross-reference is the fix for the class +// and not for the instance. +// +// Why the point to point ops go in one coalesced group: ProcessGroupNCCL launches an unbatched p2p op as a +// group of its own on the group's stream, so a rank's p2p ops are stream-ordered against each other, and two +// ranks each waiting to send into the other's not-yet-issued receive is a cycle. (NCCL documents that +// serialization; w24's M1387 reaches the same deadlock from a direct all-to-all of unbatched p2p ops.) +// ProcessGroup::startCoalescing/endCoalescing puts all of this rank's neighbour ops into one NCCL group, which +// removes that ordering: the exchange is one hop, posted and closed with no host block. +#include "comm.h" +#include +#include +#include +#include +#include + +namespace { + +cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); } + +// One uint4 is eight bf16 channels; C is a multiple of 8 for every layer that sharded space (256/512/1024). +__global__ void pack_edge8(const uint4* __restrict__ x, uint4* __restrict__ y, int64_t n8, int CV, int T, int H, + int W) { + for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) { + int64_t row = i / (W * CV), rv = i % (W * CV); + int t = (int)(row % T), edge = (int)(row / T); + y[i] = x[((int64_t)t * H + (edge ? H - 2 : 1)) * W * CV + rv]; + } +} + +// Scatters the two received edges into the pad rows, eight channels per thread. +__global__ void unpack_edge8(uint4* __restrict__ x, const uint4* __restrict__ e, int64_t n8, int CV, int T, int H, + int W, int rank, int world) { + for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) { + int64_t row = i / (W * CV), rv = i % (W * CV); + int t = (int)(row % T), edge = (int)(row / T); + int peer = rank + (edge ? 1 : -1); // pad row 0 comes from rank-1, pad row H-1 from rank+1 + uint4 v = make_uint4(0u, 0u, 0u, 0u); + if (peer >= 0 && peer < world) v = e[((int64_t)peer * 2 * T + (int64_t)(1 - edge) * T + t) * W * CV + rv]; + x[((int64_t)t * H + (edge ? H - 1 : 0)) * W * CV + rv] = v; + } +} + +int grid_for(int64_t work) { + int64_t blocks = (work + 255) / 256; + return (int)(blocks > 2048 ? 2048 : (blocks < 1 ? 1 : blocks)); +} + +// Where each transport is used, when WAN_HALO does not force one. With P the bytes of one edge (half the packed +// buffer): the ring sends (world-1)*2P and receives world*2P, of which this exchange reads 2P (one edge from each +// neighbour); the neighbour path sends and receives one edge per neighbour -- 2P in and out for an interior rank, +// P at the ends. So it is 6P -> 2P each way at world 4 and 2P/4P -> P/P at world 2, i.e. the received bytes fall +// by 3x and 4x. So the policy is the ring only where it has been measured to cost nothing, and the +// neighbour path everywhere else: +// world 2 the two arms measured the same, and the copy counts are equal there too. J817, three +// alternating rounds in one job at f121x2, minima 698.6 ms (ring) against 696.5 ms +// (neighbour), 0.3% apart, J570's 2.3% the other way sitting inside its own run-to-run +// spread; and J3057 ran none / ag / p2p at f481x2 in one job, 2305.6 / 2309.7 / 2293.2. +// world 4 neighbour. J2211's own `for A in ag p2p ag p2p` at f481x4: ag 2462.9 / 2244.5 against +// p2p 1010.6 / 1017.5, the p2p arm's A/A 0.68 %. The ring moves (world-1) copies each +// way, the neighbour path one. +// above 4 neighbour, and this rung is chosen on COPY COUNTS rather than on a measurement: the +// ring moves (world-1) copies each way against the neighbour path's 1, so the ratio grows +// with the world size, and no world above 4 has been measured on either arm here -- a +// cross-tray one has never run. The world-4 gap (2.32x) is the only scaled evidence and +// it is the smallest world where the two differ. Whoever first gets eight GPUs should +// treat this rung as the thing to measure, not as settled. +// What that leaves: every world the harness runs is unchanged -- 1 by the early return, 2 true, 4 false +// -- and only world 3 and world 5 and up change. +// WAN_HALO forces either arm for a measurement. +bool use_allgather(int world) { + if (world < 2) return true; // one rank has no neighbour to talk to + const int forced = halo_transport(); + if (forced == kAllGather) return true; + if (forced == kP2P) return false; + return world == 2; +} + +} // namespace + +HaloStats& halo_stats() { + static HaloStats stats; + return stats; +} + +int halo_transport() { + static const int mode = [] { + const char* s = std::getenv("WAN_HALO"); + if (!s) return (int)kPolicy; + if (std::strcmp(s, "ag") == 0 || std::strcmp(s, "allgather") == 0) return (int)kAllGather; + if (std::strcmp(s, "p2p") == 0 || std::strcmp(s, "neighbour") == 0) return (int)kP2P; + return (int)kPolicy; + }(); + return mode; +} + +at::Tensor halo_pack(const at::Tensor& padded) { + const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4); + TORCH_CHECK(C % 8 == 0, "halo_pack needs C a multiple of 8"); + at::Tensor e = at::empty({2, T, W, C}, padded.options()); + const int CV = C / 8; + int64_t n8 = (int64_t)2 * T * W * CV; + pack_edge8<<>>((const uint4*)padded.data_ptr(), (uint4*)e.data_ptr(), n8, CV, T, H, + W); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return e; +} + +void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world) { + const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4); + TORCH_CHECK(C % 8 == 0, "halo_unpack needs C a multiple of 8"); + const int CV = C / 8; + int64_t n8 = (int64_t)2 * T * W * CV; + unpack_edge8<<>>((uint4*)padded.data_ptr(), (const uint4*)edges.data_ptr(), n8, CV, + T, H, W, rank, world); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world, + cudaStream_t comm) { + HaloHandle h; + h.started = true; + // Both buffers are flat here: the p2p slices are element offsets (narrowing the packed [2,T,W,C] on dim 0 would + // index a dimension of size 2, which is what round 1 of this PR did). + TORCH_CHECK(send.is_contiguous(), "halo edges must be contiguous"); + TORCH_CHECK(recv.is_contiguous(), "the halo receive buffer must be contiguous"); + h.send = send.reshape({-1}); + h.recv = recv.reshape({-1}); + HaloStats& s = halo_stats(); + ++s.exchanges; + const int64_t bytes = send.numel() * (int64_t)send.element_size(); + // The transport is ordered on `comm`: ProcessGroupNCCL makes its own stream wait for the current stream and + // the current stream wait for its end event, so this guard is what puts the collective after the pack. + c10::cuda::CUDAStreamGuard guard(c10::cuda::getStreamFromExternal(comm, c10::cuda::current_device())); + if (use_allgather(world)) { + // A ring all-gather: each rank's buffer goes round the ring, (world-1) copies each way. + at::Tensor in = send; + h.keep.push_back(in); + h.keep.push_back(recv); + h.works.push_back(pg->_allgather_base(recv, in)); + s.bytes_sent += bytes * (world - 1); + s.bytes_received += bytes * (world - 1); + return h; + } + // The packed edges are [2,T,W,C]; the p2p slices below are element offsets, so they are taken flat (above). + h.keep.push_back(send); + h.keep.push_back(h.send); + h.keep.push_back(recv); + // Split on rank 0 = rows [0, H/world), so rank+1 is the band BELOW me and rank-1 the band above. Each rank + // sends its first real row up to rank-1 and its last real row down to rank+1 -- one edge per neighbour, not + // both to each -- and reads back only the two neighbour slots of `recv`; its other world-2 slots are never + // written and stay uninitialised, which halo_unpack is careful not to read (it touches slot rank-1's edge 1 + // and slot rank+1's edge 0 and nothing else). On the all-gather path every slot is filled. Both leave the + // same two edges in the same places. + const int64_t half = send.numel() / 2; // one edge: [T,W,C]; edge 0 first, edge 1 second + const int neighbours = (rank > 0 ? 1 : 0) + (rank + 1 < world ? 1 : 0); + s.bytes_sent += (int64_t)neighbours * half * send.element_size(); + s.bytes_received += (int64_t)neighbours * half * send.element_size(); + // Every tensor operation happens before the group opens: a throw between startCoalescing and endCoalescing + // would leave the NCCL group open and hang every rank instead of failing this one. + const bool hasLower = rank + 1 < world, hasUpper = rank > 0; + std::vector recvLower, sendUpper, recvUpper, sendLower; + if (hasLower) { + at::Tensor v = h.recv.narrow(0, (int64_t)(rank + 1) * 2 * half, half); // below: its edge 0 + h.keep.push_back(v); + recvLower.push_back(v); + } + if (hasUpper) { + at::Tensor v = h.send.narrow(0, 0, half); // above: my first real row + h.keep.push_back(v); + sendUpper.push_back(v); + } + if (hasUpper) { + at::Tensor v = h.recv.narrow(0, (int64_t)(rank - 1) * 2 * half + half, half); // above: its edge 1 + h.keep.push_back(v); + recvUpper.push_back(v); + } + if (hasLower) { + at::Tensor v = h.send.narrow(0, half, half); // below: my last real row + h.keep.push_back(v); + sendLower.push_back(v); + } + pg->startCoalescing(c10::DeviceType::CUDA); + if (hasLower) pg->recv(recvLower, rank + 1, 0); + if (hasUpper) pg->send(sendUpper, rank - 1, 0); + if (hasUpper) pg->recv(recvUpper, rank - 1, 0); + if (hasLower) pg->send(sendLower, rank + 1, 0); + // The ops inside a coalesced group carry no work of their own: the group's work is the one to wait on. + h.works.push_back(pg->endCoalescing(c10::DeviceType::CUDA)); + return h; +} + +void wait_edges(HaloHandle& h) { + if (!h.started) return; + // Each wait installs the collective's completion as an event wait on the current stream and does not block + // the host, so the caller's next kernel on that stream is ordered after the transfer. + for (auto& w : h.works) w->wait(); + h.works.clear(); + h.started = false; +} + +at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world) { + at::Tensor recv = at::empty({world * send.numel()}, send.options()); + HaloHandle h = post_edges(pg, send, recv, rank, world, at::cuda::getCurrentCUDAStream()); + wait_edges(h); + return recv; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/comm.h b/fastvideo-kernel/wan22_vae_decode/solution/comm.h new file mode 100644 index 000000000..d1824eb60 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/comm.h @@ -0,0 +1,83 @@ +// Halo edge transport for the row-band split: the two boundary rows a 3x3 convolution needs from its +// neighbours, moved with the cheapest mechanism available. Representation-agnostic by design: it carries +// whatever `halo_pack` packs (today the normalized NDHWC bf16 edge rows; the fusion work may change that), +// and it leaves the packed layout and `halo_unpack` untouched, so the decoded video is bit-identical to the +// all-gather it replaces. +// +// Two ways in: +// halo_exchange(...) one call: allocates the recv buffer, transports and closes. What +// Decoder::conv uses today. +// post_edges(...) / wait_edges(...) the split seam for overlapping the transport with interior compute: +// post enqueues the exchange and returns a handle without a host block, +// wait closes it. `halo_pack` may run on one stream and the transport on +// another: post_edges takes the stream the transport is ordered on. +#pragma once +#include +#include +#include +#include +#include +#include +#include + +using PG = c10::intrusive_ptr; + +// Which transport the halo exchange uses, from WAN_HALO once per process. With nothing set the world-size rule +// in comm.cu decides; the variable is for forcing an arm in a measurement (both must be the same on every rank +// of a world: one rank in a collective and one in a p2p group would never meet). +// unset: the rule -- the ring all-gather only at world 2, where both arms have been measured equal, and the +// neighbour path at world 4 and at every world above it. (World 1 has no exchange at all: ph is 0 +// there, so nothing calls this.) +// "ag" or "allgather": the ring collective, whatever the world size. +// "p2p" or "neighbour": the neighbour path, whatever the world size. This is what the rule selects above +// world 4, and the neighbour path across trays has never been measured: nothing at that world size +// has, on either arm. +enum HaloTransport { kPolicy = 0, kAllGather = 1, kP2P = 2 }; +int halo_transport(); + +// The exchange sites and the bytes they moved. Debug instrumentation only: nothing on the run path reads them, +// and with WAN_HALO_STATS unset they cost one increment per exchange. +struct HaloStats { + int64_t exchanges = 0, bytes_sent = 0, bytes_received = 0; +}; +HaloStats& halo_stats(); + +// Gather the two edges a rank must send: the first and the last real row of the padded NDHWC tensor +// [1,C,T,H,W] (H includes exactly one pad row per side), as [2,T,W,C] contiguous, edge-major: edge 0 is the +// first real row (H index 1), edge 1 the last (H index H-2). Same layout as the seed's pack_edges, eight +// channels per thread. +at::Tensor halo_pack(const at::Tensor& padded); + +// Scatter the received edges into the pad rows: row 0 takes rank-1's edge 1, row H-1 takes rank+1's edge 0, +// and a rank with no neighbour on that side gets zeros. `edges` is the transport's `recv` buffer, and rank+1 +// is the band *below* (higher rows), since the split gives rank 0 the rows [0, H/world). +void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world); + +// The transport, split in two. `send` is halo_pack's [2,T,W,C]; `recv` is the caller's buffer of +// world*2*T*W*C elements in the all-gather's layout (slot r is rank r's two edges), which is what +// halo_unpack reads. The handle owns the works and keeps every send and recv view alive until it is waited. +struct HaloHandle { + std::vector> works; + std::vector keep; + at::Tensor send, recv; // flat views the p2p slices index + bool started = false; + HaloHandle() = default; + HaloHandle(HaloHandle&&) = default; + HaloHandle& operator=(HaloHandle&&) = default; + HaloHandle(const HaloHandle&) = delete; + HaloHandle& operator=(const HaloHandle&) = delete; +}; + +// Enqueue only: no host block, and nothing added to the compute stream. `comm` is the stream the transport is +// ordered on; it must already carry whatever ordering the caller needs (a wait on the pack, say), and the +// unpack must not run on a different stream before wait_edges has been called on it. +HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world, + cudaStream_t comm); + +// Close it: this orders the CALLER'S CURRENT STREAM after the transfer and returns at once. WorkNCCL::wait is +// an event wait on the stream that is current at the call, not a CPU block. So the next kernel enqueued on +// that stream sees the edges -- and a kernel enqueued on any other stream does not. +void wait_edges(HaloHandle& h); + +// Allocate the recv buffer, post, close and return the edges in halo_unpack's layout. +at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/conv.cu b/fastvideo-kernel/wan22_vae_decode/solution/conv.cu new file mode 100644 index 000000000..874c766e3 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/conv.cu @@ -0,0 +1,78 @@ +// CUTLASS 4.2 Blackwell implicit convolution: 2-CTA 256x256x64 tiles, +// TMA im2col loads, tcgen05 MMA, fp32 accumulation, bf16 epilogue. +// Every included CUTLASS/CuTe dependency is vendored beside this source. +#include +#include +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/convnd_problem_shape.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/conv/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/conv/device/conv_universal_adapter.hpp" +#include "cutlass/conv/kernel/conv_universal.hpp" +using namespace cute; +using ElementBF16=cutlass::bfloat16_t; +using LayoutNDHWC=cutlass::layout::TensorNDHWC; +using ArchTag=cutlass::arch::Sm100; +using OperatorClass=cutlass::arch::OpClassTensorOp; +constexpr auto Fprop=cutlass::conv::Operator::kFprop; + +template +at::Tensor launch(const at::Tensor& x,const at::Tensor& w,int ph,int pw,int phu=-1,int pwu=-1) { + if(phu<0)phu=ph;if(pwu<0)pwu=pw; + using Tile=Shape,Int,Shape<_64>>; + using Cluster=Shape,_1,_1>; + using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag,OperatorClass,Tile,Cluster,cutlass::epilogue::collective::EpilogueTileAuto, + float,float,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + using Mainloop=typename cutlass::conv::collective::CollectiveBuilder< + ArchTag,OperatorClass,Fprop,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,float,Tile,Cluster, + cutlass::conv::collective::StageCountAutoCarveout(sizeof(typename Epilogue::SharedStorage))>, + cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp; + using Problem=cutlass::conv::ConvProblemShape; + using Kernel=cutlass::conv::kernel::ConvUniversal; + using Conv=cutlass::conv::device::ConvUniversalAdapter; + int C=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4); + int K=w.size(0),KT=w.size(2),KH=w.size(3),KW=w.size(4); + auto y=at::empty({1,K,D-KT+1,H+ph+phu-KH+1,W+pw+pwu-KW+1},x.options().memory_format(at::MemoryFormat::ChannelsLast3d)); + Problem problem(cutlass::conv::Mode::kCrossCorrelation, + {1,D,H,W,C},{x.stride(0),x.stride(2),x.stride(3),x.stride(4),x.stride(1)}, + {K,KT,KH,KW,C},{w.stride(0),w.stride(2),w.stride(3),w.stride(4),w.stride(1)}, + {0,ph,pw},{0,phu,pwu},{1,1,1},{1,1,1},1); + typename Kernel::StrideC sc{}; + typename Kernel::StrideD sd{}; + cute::for_each(cute::make_seq(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];}); + cute::for_each(cute::make_seq(sd)>{},[&](auto i){cute::get<0,i>(sd)=problem.stride_C[Problem::RankT-2-i];}); + typename Conv::Arguments args{problem,{(const ElementBF16*)x.data_ptr(),(const ElementBF16*)w.data_ptr()}, + {{1.f,0.f},(const ElementBF16*)y.data_ptr(),sc,(ElementBF16*)y.data_ptr(),sd}}; + Conv conv; + TORCH_CHECK(conv.can_implement(args)==cutlass::Status::kSuccess,"CUTLASS cannot implement shape"); + size_t size=Conv::get_workspace_size(args); + auto workspace=at::empty({(int64_t)size},x.options().dtype(at::kByte)); + auto stream=at::cuda::getCurrentCUDAStream(); + auto status=conv.initialize(args,workspace.data_ptr(),stream); + TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS initialization failed ",int(status)); + status=conv.run(stream); + TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS convolution failed ",int(status)); + return y; +} +at::Tensor custom_conv(const at::Tensor& x,const at::Tensor& w,int64_t ph,int64_t pw,int64_t variant) { + return launch<256,256,2>(x,w,ph,pw); +} + +// The output head's 3x1x1 convolution: N=128 because the head has 112 output channels, and a +// 256-wide tile would spend 144 of its 256 columns on nothing. Measured at one clock (1951-2022 +// MHz, both arms in J1044 and J1343): N=128 is 0.83 PFLOP/s = 34% of peak, N=256 is 0.69 = 28%. +// The time is what decides it and it is not close: 6.30 ms against 15.2 ms. +at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w) { + return launch<256,128,2>(x,w,0,0); +} + +at::Tensor phase_conv(const at::Tensor& x,const at::Tensor& w,int a,int b,bool halo) { + at::Tensor v=halo?x.narrow(3,a,x.size(3)-1):x; + return launch<256,256,2>(v,w,halo?0:1-a,1-b,halo?0:a,b); +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/head.cu b/fastvideo-kernel/wan22_vae_decode/solution/head.cu new file mode 100644 index 000000000..cf9443c31 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/head.cu @@ -0,0 +1,82 @@ +// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to +// [-1,1], from the 112-channel plane conv.cu's `conv_n128` (a 3x1x1 convolution) writes to the +// 12-channel plane write_video consumes. +// +// Why the convolution writes 108 channels and this pass sums 9 of them. conv_out is 256x27 MACs +// and its per-position INPUT WINDOW is 6912*2 B (the weight is 12x that and stays resident), i.e. +// 12 FLOP/byte against the machine's ~375, so any form that reads the 27 taps per position is +// operand-bound. Moving the nine spatial taps into the output channel dimension of a 3x1x1 +// (temporal-only) convolution makes the operand 768*2 B for 108 outputs instead of 13.8 KB for +// 12, and the nine shifted copies are summed here. +// +// The value path, so that the rounding is on the record: the three temporal taps are accumulated +// in the convolution's fp32 accumulator and rounded to bf16 in g; the nine spatial taps are then +// summed here in fp32 from nine bf16 g values, the bias is added in fp32, and the result is +// rounded once to bf16 on the store. Main's path is also two roundings (cuDNN's bf16 conv_out, +// then finish_scalar's bf16 of x+bias); the count is the same and what differs is WHERE the first +// one sits - after three taps instead of after twenty-seven - which is why max|d| moves and rms +// does not. The reference accumulates all 27 taps in one fp32 accumulator. +#include +#include +#include +#include +#include "head.h" + +// One thread per output position. Its 12 accumulators take the 12 channels of the same (i,j) +// tap from nine neighbouring positions; the nine shifted reads of each 216-byte channel vector +// are served from L1, so the pass moves each element of g about once. +static __global__ void head_stencil_kernel(const __nv_bfloat16* __restrict__ g, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, + int D,int rows,int W,int H,int g_off,int lo) { + int idx=blockIdx.x*blockDim.x+threadIdx.x; + if(idx>=D*rows*W)return; + int w=idx%W, t=idx/W, r=t%rows, d=t/rows; + int gh0=lo+r; + const int gH=rows+2*g_off; + float acc[12]; +#pragma unroll + for(int c=0;c<12;++c)acc[c]=__bfloat162float(bias[c]); + for(int i=0;i<3;++i) { + int gh=gh0+i-1; + if(gh<0||gh>=H)continue; // conv3d's zero padding, in the plane + int gr=r+i-1+g_off; + if(gr<0||gr>=gH)continue; + for(int j=0;j<3;++j) { + int ww=w+j-1; + if(ww<0||ww>=W)continue; // conv3d's zero padding, across the row + const __nv_bfloat162* p=(const __nv_bfloat162*)(g+((((size_t)d*gH+gr)*W)+ww)*112+(i*3+j)*12); +#pragma unroll + for(int k=0;k<6;++k) { + float2 f=__bfloat1622float2(p[k]); + acc[2*k]+=f.x; + acc[2*k+1]+=f.y; + } + } + } + __nv_bfloat162* o=(__nv_bfloat162*)(out+((((size_t)d*rows+r)*W)+w)*12); +#pragma unroll + for(int k=0;k<6;++k) + o[k]=__floats2bfloat162_rn(fminf(1.f,fmaxf(-1.f,acc[2*k])),fminf(1.f,fmaxf(-1.f,acc[2*k+1]))); +} + +void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out, + int64_t H_global,int64_t g_off,int64_t lo) { + int D=out.size(2),rows=out.size(3),W=out.size(4); + TORCH_CHECK(g.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&g.size(1)==112, + "head_stencil: g must be [1,112,D,rows+2*g_off,W] channels-last"); + TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&out.size(1)==12, + "head_stencil: out must be [1,12,D,rows,W] channels-last"); + // The kernel strides g by gH = rows+2*g_off taken from `out` and reads D*rows*W positions of it, + // so a g of the wrong depth or height reads out of bounds rather than failing here. + TORCH_CHECK(g.size(2)==out.size(2)&&g.size(3)==out.size(3)+2*g_off&&g.size(4)==out.size(4), + "head_stencil: g must be [1,112,D,rows+2*g_off,W] of out's [1,12,D,rows,W]"); + c10::cuda::CUDAGuard guard(g.device()); + int64_t total=(int64_t)D*rows*W; + int threads=256; + int blocks=(int)((total+threads-1)/threads); + head_stencil_kernel<<>>( + (const __nv_bfloat16*)g.data_ptr(),(const __nv_bfloat16*)bias.data_ptr(), + (__nv_bfloat16*)out.data_ptr(),D,rows,W,(int)H_global,(int)g_off,(int)lo); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/head.h b/fastvideo-kernel/wan22_vae_decode/solution/head.h new file mode 100644 index 000000000..0a76b4efb --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/head.h @@ -0,0 +1,19 @@ +#pragma once +#include + +// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to +// [-1,1], from the 112-channel plane the 3x1x1 convolution writes (conv.cu's `conv_n128`) to the +// 12-channel plane write_video consumes. +// g: [1,112,D,rows+2*g_off,W] channels-last-3d bf16, n = (i*3+j)*12 + c for the tap (i,j); +// channels 108..111 are the zero-weight lanes and are never read. +// out: [1,12,D,rows,W] channels-last-3d bf16, written whole. +// H_global is the height of the plane this stencil filters (352 = 44 latent rows x 8; the pixel +// plane is 704 tall and run()'s video row is twice this one). lo is the first row of this rank's +// band IN THAT PLANE, 8*(44*rank/world), and it must stay that: run() passes lo*16 to write_video, +// the same row on the 704-tall video grid, and the two derivations have to agree. +void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out, + int64_t H_global,int64_t g_off,int64_t lo); + +// conv.cu: the head's 3x1x1 convolution, at the tile its 112 output channels want. Declared here +// so the two-sided declaration of a CUTLASS tile cannot drift from its definition in conv.cu. +at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu b/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu new file mode 100644 index 000000000..684ea67a6 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu @@ -0,0 +1,99 @@ +#include "kernels.h" +#include +#include +#include +#include +using B = __nv_bfloat16; +union Pack { uint4 v; B b[8]; __nv_bfloat162 h[4]; }; +static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); } +static const B* ptr(const Tensor& t) { return t.defined() ? (const B*)t.data_ptr() : nullptr; } +static B* ptr(Tensor& t) { return (B*)t.data_ptr(); } +Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like) { + return at::empty({1,c,t,h,w},like.options().memory_format(at::MemoryFormat::ChannelsLast3d)); +} +Tensor last_frames(const Tensor& x,int n,int ph) { + int C=x.size(1),T=x.size(2),H=x.size(3)-2*ph,W=x.size(4); + Tensor y=nhwc_empty(C,n,H,W,x); + size_t row=(size_t)H*W*C*sizeof(B),src_pitch=(size_t)(H+2*ph)*W*C*sizeof(B); + const B* source=ptr(x)+((int64_t)(T-n)*(H+2*ph)+ph)*W*C; + C10_CUDA_CHECK(cudaMemcpy2DAsync(ptr(y),row,source,src_pitch,row,n,cudaMemcpyDeviceToDevice,stream())); + return y; +} +// The halo edge pack/scatter lives in comm.cu (halo_pack/halo_unpack, vectorised), with the transport. +__global__ void finish_vector(B* x,const B* b,const B* s,int64_t n,int C) { + for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i>>(ptr(x),ptr(b),ptr(s),n,C); + else finish_scalar<<>>(ptr(x),ptr(b),ptr(s),n,C); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +__global__ void write_kernel(const B* x,B* out,int64_t n,int T,int H,int W,int fullT,int fullH,int offsetT,int offsetH) { + int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x; + if(i>>(ptr(x),ptr(out),x.numel(),x.size(2),x.size(3),x.size(4),out.size(2),out.size(3),frame_offset,row_offset); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +// phases_kernel used to sit here: PR17 replaces it, together with finish_phases below, with one +// bandwidth-bound pass in resample.cu that writes the interleaved tensor directly. +// F(4x4,3x3), with fp32 spatial transforms and bf16 transformed operands. +__device__ __forceinline__ void bt6(float v0,float v1,float v2,float v3,float v4,float v5,float* y) { + float a=v4-4*v2,b=v3-4*v1,c=v4-v2,d=2*(v3-v1); + y[0]=4*v0-5*v2+v4;y[1]=a+b;y[2]=a-b;y[3]=c+d;y[4]=c-d;y[5]=4*v1-5*v3+v5; +} +__global__ void winograd_input_kernel(const B* x,B* v,int C,int D,int H,int W,int PH,int TH,int TW) { + int HP=H+2*PH,CV=C/2;int64_t n=(int64_t)D*TH*TW*CV; + for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i=0&&y=0&&xx>>(ptr(x),ptr(v),C,D,H,W,ph,TH,TW); + C10_CUDA_KERNEL_LAUNCH_CHECK();return v; +} +// (WinoPtrs, at6, winograd_output_kernel and winograd_output are deleted: with the fused inverse +// in wino_epilogue.cu nothing calls them. That file's a6v reads the same A^T out of shared memory +// instead of 36 separate global pointers, and it is the only copy of the transform left there.) + diff --git a/fastvideo-kernel/wan22_vae_decode/solution/kernels.h b/fastvideo-kernel/wan22_vae_decode/solution/kernels.h new file mode 100644 index 000000000..ff1efda7b --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/kernels.h @@ -0,0 +1,21 @@ +#pragma once +#include +using at::Tensor; +Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like); +Tensor norm_pad(const Tensor& x, const Tensor& gamma, const Tensor& history, int pt, int ph, bool silu); +Tensor last_frames(const Tensor& padded, int n, int ph); +void finish(Tensor& x, const Tensor& bias, const Tensor& residual); +void write_video(const Tensor& x, Tensor& out, int frame_offset, int row_offset); + +Tensor winograd_input(const Tensor& x,int ph); +Tensor winograd_input_fast(const Tensor& x,int ph); +// winograd_output / winograd_output_kernel are deleted: the fused inverse in wino_epilogue.cu +// replaced their only caller, and they had no live user left (PR172 round-1 review). + +// Fused inverse transform + bias/residual + RMS norm/SiLU. y = A^T M A + bias + residual; +// n (when gamma is given) = silu(rmsnorm(y)*sqrt(C)*gain) in the (D+pt, H+2ph, W, C) padded +// layout the next convolution consumes, with its first pt frames taken from `hist`. +struct WinoOut { Tensor y; Tensor n; }; +WinoOut wino_epilogue(const std::vector& m,int H,int W,const Tensor& bias, + const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph, + bool write_y=true); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/norm.cu b/fastvideo-kernel/wan22_vae_decode/solution/norm.cu new file mode 100644 index 000000000..6e0bbd124 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/norm.cu @@ -0,0 +1,198 @@ +// norm.cu — RMS norm + SiLU + pad for the Wan2.2 VAE decoder, one warp per position. +// +// This replaces the seed's norm_vector / pad_vector pair in kernels.cu with the same +// call contract (norm_pad(x, gamma, hist, pt, ph, silu)) and the same layout (NDHWC, +// [T+pt][H+2ph][W][C], history frames copied through, pad rows zeroed). +// +// Why it is faster (measured in probe/norm_probe2.cu and norm_probe3.cu, one GB200): +// - the seed is XU-pipe bound (w08's and w14's ncu: 78% SM busy, 16% DRAM, XU 75%): +// per element it pays cvt bf16->f32 for the sum, five cvts for the reference's three +// bf16 roundings, __expf and an IEEE divide. This does +// * bf16 -> f32 by an integer shift (ALU pipe, exact, no XU op), +// * one rounding for the whole normalize*sqrt(C)*gain product instead of three, +// which is also closer to the fp32 reference the decode is graded against, +// * SiLU as u + u*tanh(u) with u = 0.5*pre-activation (one tanh.approx.f32), +// * bf16 stores two elements at a time (cvt.rn.bf16x2.f32), +// i.e. ~1.5 XU ops per element instead of ~11. +// - rows are blocked so the per-position integer divisions of the seed's flat-pixel loop +// happen once per row, not once per position (at C=256 a position is only 8 elements). +// - P positions per warp in flight, the number that measured best per channel count. +// +// Numerics: one bf16 rounding for the whole normalize*gain product instead of three, and +// SiLU in fp32 through tanh.approx.f32. That approximation's documented worst case is an +// absolute error of about 2^-11 (5e-4) in tanh, which is harmless here because it enters +// the result as u*2^-11 and the stored value is bf16, so the perturbation stays below the +// bf16 quantum of u + u*tanh(u) over the range these pre-activations take. Measured, and +// the place to look if that ever changes (bench.py, f121x1 vs the fp32 reference): +// seed max 0.1709 (0.34x budget) rms 0.002242 (0.38x budget) +// this max 0.1992 (0.40x budget) rms 0.002187 (0.37x budget) +// i.e. the max moves up a little, the rms moves down, both far inside the budget. +#include +#include +#include +#include +#include +#include +#include + +using B = __nv_bfloat16; + +namespace { + +__device__ __forceinline__ void split2(uint32_t w, float& lo, float& hi) { + lo = __uint_as_float(w << 16); // exact bf16 -> f32, ALU pipe + hi = __uint_as_float(w & 0xffff0000u); +} +__device__ __forceinline__ float tanh_approx(float x) { + float r; asm("tanh.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); return r; +} +__device__ __forceinline__ uint32_t pack2(float a, float b) { + uint32_t r; asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(b), "f"(a)); return r; +} + +// MODE 0: RMS norm * sqrt(C) * gain, then SiLU. MODE 1: the same, no SiLU. MODE 2: copy only. +template +__global__ void norm_pad_kernel(const B* __restrict__ x, const B* __restrict__ gain, + const B* __restrict__ hist, B* __restrict__ y, + int T, int H, int W, int PT, int PH) { + constexpr int CV = C / 8; // uint4 per position + constexpr int NP = C / 256; // uint4 per lane per position + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5; + const int HP = H + 2 * PH; + const int64_t rows = (int64_t)(T + PT) * HP; + const float sqrtC = sqrtf((float)C); + for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) { + const int h = (int)(row % HP), t = (int)(row / HP); + const bool interior = (h >= PH && h < H + PH); + const bool current = (t >= PT); // a history frame is copied through, never re-normalized + const int ph = h - PH; + const uint4* src = nullptr; + if (interior) { + if (current) src = (const uint4*)x + ((int64_t)(t - PT) * H + ph) * W * CV; + else if (hist) src = (const uint4*)hist + ((int64_t)t * H + ph) * W * CV; + } + for (int w0 = warp * P; w0 < W; w0 += warps * P) { + uint4 v[P][NP]; + #pragma unroll + for (int q = 0; q < P; q++) { + const int w = w0 + q; + #pragma unroll + for (int j = 0; j < NP; j++) { + v[q][j] = make_uint4(0, 0, 0, 0); + if (src && w < W) v[q][j] = src[(int64_t)w * CV + lane + j * 32]; + } + } + if (MODE < 2 && interior && current) { + float s[P][2]; + #pragma unroll + for (int q = 0; q < P; q++) { s[q][0] = 0.f; s[q][1] = 0.f; } + #pragma unroll + for (int q = 0; q < P; q++) { + #pragma unroll + for (int j = 0; j < NP; j++) { + float f[8]; + split2(v[q][j].x, f[0], f[1]); split2(v[q][j].y, f[2], f[3]); + split2(v[q][j].z, f[4], f[5]); split2(v[q][j].w, f[6], f[7]); + #pragma unroll + for (int k = 0; k < 8; k += 2) { + s[q][0] = fmaf(f[k], f[k], s[q][0]); + s[q][1] = fmaf(f[k + 1], f[k + 1], s[q][1]); + } + } + } + #pragma unroll + for (int q = 0; q < P; q++) { + float t0 = s[q][0] + s[q][1]; + #pragma unroll + for (int d = 16; d; d >>= 1) t0 += __shfl_xor_sync(0xffffffff, t0, d); + // 0.5 * sqrt(C) / max(sqrt(sum), 1e-12): one multiply per element, u = 0.5*pre-activation + const float hs = 0.5f * sqrtC * rsqrtf(fmaxf(t0, 1e-24f)); + #pragma unroll + for (int j = 0; j < NP; j++) { + const uint4 gn = ((const uint4*)gain)[lane + j * 32]; + float g[8], xf[8], yf[8]; + split2(gn.x, g[0], g[1]); split2(gn.y, g[2], g[3]); + split2(gn.z, g[4], g[5]); split2(gn.w, g[6], g[7]); + split2(v[q][j].x, xf[0], xf[1]); split2(v[q][j].y, xf[2], xf[3]); + split2(v[q][j].z, xf[4], xf[5]); split2(v[q][j].w, xf[6], xf[7]); + #pragma unroll + for (int k = 0; k < 8; k++) { + const float u = xf[k] * hs * g[k]; + yf[k] = (MODE == 0) ? fmaf(u, tanh_approx(u), u) : (2.f * u); + } + v[q][j].x = pack2(yf[0], yf[1]); v[q][j].y = pack2(yf[2], yf[3]); + v[q][j].z = pack2(yf[4], yf[5]); v[q][j].w = pack2(yf[6], yf[7]); + } + } + } + #pragma unroll + for (int q = 0; q < P; q++) { + const int w = w0 + q; + if (w < W) { + uint4* dst = (uint4*)y + row * W * CV + (int64_t)w * CV + lane; + #pragma unroll + for (int j = 0; j < NP; j++) dst[j * 32] = v[q][j]; + } + } + } + } +} + +// Copy + pad for any C divisible by 8: the calls with no gain (post_quant_conv's input, the temporal convs). +__global__ void copy_row_kernel(const uint4* __restrict__ x, const uint4* __restrict__ hist, + uint4* __restrict__ y, int T, int H, int W, int PT, int PH, int CV) { + const int HP = H + 2 * PH; + const int64_t rows = (int64_t)(T + PT) * HP; + const int64_t nrow = (int64_t)W * CV; + for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) { + const int h = (int)(row % HP), t = (int)(row / HP); + const bool interior = (h >= PH && h < H + PH); + const uint4* src = nullptr; + if (interior) { + const int tt = t - PT, ph = h - PH; + if (tt >= 0) src = x + ((int64_t)tt * H + ph) * nrow; + else if (hist) src = hist + ((int64_t)t * H + ph) * nrow; + } + uint4* dst = y + row * nrow; + const uint4 zz = make_uint4(0, 0, 0, 0); + for (int64_t i = threadIdx.x; i < nrow; i += blockDim.x) dst[i] = src ? src[i] : zz; + } +} + +} // namespace + +at::Tensor norm_pad(const at::Tensor& x, const at::Tensor& gamma, const at::Tensor& hist, + int pt, int ph, bool silu) { + const int C = (int)x.size(1), T = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4); + auto y = at::empty({1, C, T + pt, H + ph * 2, W}, + x.options().memory_format(at::MemoryFormat::ChannelsLast3d)); + const B* xp = (const B*)x.data_ptr(); + const B* gp = gamma.defined() ? (const B*)gamma.data_ptr() : nullptr; + const B* hp = hist.defined() ? (const B*)hist.data_ptr() : nullptr; + B* yp = (B*)y.data_ptr(); + const int64_t rows = (int64_t)(T + pt) * (H + ph * 2); + auto stream = at::cuda::getCurrentCUDAStream(); + const int blocks = (int)std::min(rows, 1024); + const int threads = 256; + auto go = [&](auto tag) { + constexpr int CC = decltype(tag)::value; + // P positions in flight per warp: measured best at 4 for C<=512, 2 at C=1024 + constexpr int P = (CC == 1024) ? 2 : 4; + if (!gp) norm_pad_kernel<<>>(xp, gp, hp, yp, T, H, W, pt, ph); + else if (silu) norm_pad_kernel<<>>(xp, gp, hp, yp, T, H, W, pt, ph); + else norm_pad_kernel<<>>(xp, gp, hp, yp, T, H, W, pt, ph); + }; + const bool fast = (C == 1024 || C == 512 || C == 256); + TORCH_CHECK(!gp || fast, "Unsupported RMS width ", C); + if (fast) { + if (C == 1024) go(std::integral_constant{}); + else if (C == 512) go(std::integral_constant{}); + else go(std::integral_constant{}); + } else { + TORCH_CHECK(C % 8 == 0, "norm_pad copy needs C divisible by 8, got ", C); + copy_row_kernel<<>>((const uint4*)xp, (const uint4*)hp, (uint4*)yp, T, H, W, + pt, ph, C / 8); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/prefix.cu b/fastvideo-kernel/wan22_vae_decode/solution/prefix.cu new file mode 100644 index 000000000..00a9b9630 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/prefix.cu @@ -0,0 +1,56 @@ +// Prefix ownership transpose: no arithmetic, only bf16 NDHWC packets. +#include "prefix.h" +#include "kernels.h" +#include +#include +#include + +namespace { +// Blocks for destination bands are disjoint; each packet is [T,H_peer,W,C]. +__global__ void pack_prefix_rows(const uint4* x,uint4* send,int T,int H,int WC8,int world) { + const int peer=blockIdx.y; + const int lo=H*peer/world, hi=H*(peer+1)/world; + const int64_t n=(int64_t)T*(hi-lo)*WC8; + const int64_t offset=(int64_t)T*lo*WC8; + for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i& pg) { + const int C=local.size(1),T=local.size(2),H=local.size(3),W=local.size(4); + TORCH_CHECK(C%8==0 && local.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix needs aligned NDHWC"); + TORCH_CHECK(T==total_frames*(rank+1)/world-total_frames*rank/world,"Prefix ownership mismatch"); + auto send=at::empty({local.numel()},local.options()); + const int band_height=H*(rank+1)/world-H*rank/world; + auto output=nhwc_empty(C,total_frames,band_height,W,local); + // ProcessGroup splits count bf16 elements, not bytes. Receive into the + // physical NDHWC order: the source ranks own chronological frame intervals. + auto physical=output.permute({0,2,3,4,1}); + TORCH_CHECK(physical.is_contiguous(),"Prefix receive must be dense NDHWC"); + auto recv=physical.view({-1}); + if(T) { + dim3 grid(256,world); + pack_prefix_rows<<>>( + (const uint4*)local.data_ptr(),(uint4*)send.data_ptr(),T,H,W*(C/8),world); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + std::vector in_split(world),out_split(world); + int64_t sent=0,received=0; + for(int p=0;palltoall_base(recv,send,out_split,in_split)->wait(); + return output; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/prefix.h b/fastvideo-kernel/wan22_vae_decode/solution/prefix.h new file mode 100644 index 000000000..5815a99f1 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/prefix.h @@ -0,0 +1,7 @@ +#pragma once +#include +#include +// Exchange chronological full-frame prefix slices for chronological row bands. +// Source rank r owns [floor(T*r/world), floor(T*(r+1)/world)). +at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames, + int rank,int world,const c10::intrusive_ptr& pg); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/resample.cu b/fastvideo-kernel/wan22_vae_decode/solution/resample.cu new file mode 100644 index 000000000..896b99003 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/resample.cu @@ -0,0 +1,96 @@ +// Resampler stage: the four-phase spatial upsamplers of up_blocks 0,1,2 of the Wan2.2 VAE decoder. +// See resample.h for why the phase convolutions stay direct and what this pass replaces. +#include "resample.h" +#include +#include +#include +using B = __nv_bfloat16; +// conv.cu: a 2x2 low-resolution convolution of one output parity, padded per phase; `halo` is set when +// the input carries a one-row halo from the neighbour rank, so the input is narrowed instead of padded. +at::Tensor phase_conv(const at::Tensor&, const at::Tensor&, int, int, bool); + +static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); } +static const B* pc(const Tensor& t) { return (const B*)t.data_ptr(); } +static B* pm(Tensor& t) { return (B*)t.data_ptr(); } + +union P8 { uint4 v; B b[8]; __nv_bfloat162 h[4]; }; + +// One thread: eight channels of one 2x2 low-resolution tile, all four phases. The grid carries the +// coordinates (tx, t, ty) so there is no integer division on the hot path, and every access is a +// 16-byte vector: four phase outputs and the DupUp3D source in, four interleaved positions out. +__global__ void upsampler_kernel(const B* __restrict__ p0, const B* __restrict__ p1, + const B* __restrict__ p2, const B* __restrict__ p3, + B* __restrict__ y, const B* __restrict__ bias, + const B* __restrict__ skip, int C, int HL, int WL, int TC, + bool temporal, bool first) { + const int cg = threadIdx.x; // 8-channel group + const int tx = blockIdx.x, t = blockIdx.y, ty = blockIdx.z; + const int CP = C / 8, SP = TC / 8; + const P8 b = {((const uint4*)bias)[cg]}; + const int st = temporal ? (first ? (t + 1) / 2 : t / 2) : t; + const int64_t wC = 2 * (int64_t)WL * C; // one row of the interleaved output + const int64_t tC = 2 * (int64_t)HL * wC; // one frame of it +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int yy = 2 * ty + i; + if (yy >= HL) continue; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int xx = 2 * tx + j; + if (xx >= WL) continue; + const int64_t q0 = ((int64_t)t * HL + yy) * WL + xx; + P8 ph[4]; + ph[0].v = ((const uint4*)p0)[q0 * CP + cg]; + ph[1].v = ((const uint4*)p1)[q0 * CP + cg]; + ph[2].v = ((const uint4*)p2)[q0 * CP + cg]; + ph[3].v = ((const uint4*)p3)[q0 * CP + cg]; + const int64_t s = ((int64_t)st * HL + yy) * WL + xx; + P8 s0, s1; + if (temporal) { + s0.v = ((const uint4*)skip)[s * SP + cg]; + s1.v = s0.v; // nearest x2 in time: both row parities alike + } else { + // up block 2: output channel J takes source channel 2*J + (row parity), so read both parities. + s0.v = ((const uint4*)skip)[s * SP + 2 * cg]; + s1.v = ((const uint4*)skip)[s * SP + 2 * cg + 1]; + } + B* out = y + (int64_t)t * tC + (int64_t)(2 * yy) * wC + (int64_t)(2 * xx) * C + 8 * cg; + if (!temporal) { + // channel J of the output takes source channel 2J + a: even channels from s0, odd from s1. + P8 e, o; +#pragma unroll + for (int k = 0; k < 4; ++k) { + e.b[k] = s0.b[2 * k]; + e.b[k + 4] = s1.b[2 * k]; + o.b[k] = s0.b[2 * k + 1]; + o.b[k + 4] = s1.b[2 * k + 1]; + } + s0.v = e.v; + s1.v = o.v; + } + P8 r; +#pragma unroll + for (int p = 0; p < 4; ++p) { + const P8& sv = (p & 2) ? s1 : s0; // row parity of this phase +#pragma unroll + for (int k = 0; k < 4; ++k) r.h[k] = __hadd2(__hadd2(ph[p].h[k], b.h[k]), sv.h[k]); + const int a = p >> 1, bcol = p & 1; + ((uint4*)(out + (int64_t)a * wC + (int64_t)bcol * C))[0] = r.v; + } + } + } +} + +Tensor resample_upsample(const Tensor& x, const std::vector& pw, const Tensor& bias, + const Tensor& skip, bool halo, bool temporal, bool first) { + TORCH_CHECK(pw.size() == 4, "resample_upsample needs the four phase kernels"); + std::vector q(4); + for (int a = 0; a < 2; ++a) for (int b = 0; b < 2; ++b) q[2 * a + b] = phase_conv(x, pw[2 * a + b], a, b, halo); + const int C = q[0].size(1), D = q[0].size(2), HL = q[0].size(3), WL = q[0].size(4); + Tensor y = at::empty({1, C, D, 2 * HL, 2 * WL}, x.options().memory_format(at::MemoryFormat::ChannelsLast3d)); + dim3 block(C / 8, 1), grid((WL + 1) / 2, D, (HL + 1) / 2); + upsampler_kernel<<>>(pc(q[0]), pc(q[1]), pc(q[2]), pc(q[3]), pm(y), pc(bias), + pc(skip), C, HL, WL, (int)skip.size(1), temporal, first); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/resample.h b/fastvideo-kernel/wan22_vae_decode/solution/resample.h new file mode 100644 index 000000000..4f9be646c --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/resample.h @@ -0,0 +1,51 @@ +#pragma once +// Resampler stage of the Wan2.2 VAE decoder: the four-phase spatial upsamplers of up_blocks 0,1,2. +// +// A resampler is "nearest x2 in space, then a 3x3x3 convolution", which equals four 2x2 low-resolution +// convolutions, one per output parity (their repeated taps already summed into per-phase kernels by +// prepare()). Those convolutions stay direct (4 taps, no transform: F(2x2,2x2) would cut the MACs 1.78x +// but materialise 4 x 9/4 = 9 x the low-resolution tensor of transformed input and 2.25x the output of +// products, which prices out as a wash - measured and posted). +// +// What is fixed here is the pass after them: the seed's interleave kernel walked the four phase tensors +// and the DupUp3D source with four 64-bit division/modulo chains per 16 bytes and ran at 3.9 TB/s, half +// of HBM. This one walks tiles and channel groups with the grid carrying the coordinates, so the pass +// is bandwidth bound: it reads the four phase outputs and the DupUp3D source once and writes the +// interleaved high-resolution tensor once, with bias and the shortcut added in bf16 at the same +// rounding points the seed used. +#include +#include +using at::Tensor; + +// One resampler over one latent chunk. x is the low-resolution input, already padded by one row top +// and bottom ([1,C,D,H+2,W] channels-last) with the halo installed at world > 1 (`halo`), and by the two +// causal history frames in time; pw are the four per-phase 2x2 kernels ([C_out,C_in,1,2,2], phase = +// 2*row_parity + col_parity), bias [C_out], and skip the DupUp3D source ([1,TC,Ds,HL,WL]). `temporal` +// is true for up_blocks 0 and 1 (the shortcut and the phase kernels are nearest-x2 in time as well), +// false for up block 2, whose shortcut is the channel-pair-to-row rearrangement. `first` marks the +// chunk that holds latent frame 0. +// Returns the interleaved high-resolution tensor [1,C_out,D,2*HL,2*WL] channels-last. +// +// The pass moves 100.5 GB of logical traffic for f121x1 (the four phase tensors read 42.5, the DupUp3D +// source 15.5, the output written 42.5) in 14.8 ms = 6.8 TB/s. +// +// THAT 6.8 IS NOT A BOUND AND MUST NOT BE TRACED TO ONE. It is solved from 100.5 / 14.8, so quoting it +// against a ceiling is an identity - w06 M18769/M18811, w22 M18463. The bounds on this board that could +// apply are: the MIX ceiling for a read+write stage, 5.07-5.59 TB/s at this stage's own f_r (w08's +// f_r-form, J2654/J2944); w14's per-thread store-loop write arm, 4.86-5.15 (J3082); and w12's 4-byte- +// per-lane burst arm, 7.35-7.50 (J3029), which is BURST-FLATTERED at 2 KB against the tree's 128 B and +// which w12 withdrew. 6.9 TB/s was a deflation of that withdrawn arm, so it inherits the arm's defect: +// it has no f_r, it belongs to a limb with no read, and this is a read+write stage. +// +// So the honest statement is that the stage moves 100.5 GB at 6.8 TB/s and that NO CEILING FOR ITS OWN +// STREAM HAS BEEN MEASURED. The counters on this pass would supply one; that run has not been made. +// Do not write "at its byte bound", and do not write "the model is conservative" (106.7 GB / 14.8 ms is +// 7.21 TB/s, which is ABOVE the mixed ceiling - the same fact as "the model is 6 % low", stated with the +// sign that does not flatter it). +// +// J2574 also counts 1.37x the sectors of its logical read. That is sector waste on the access pattern +// and not work the algorithm needs: this kernel reads each phase tensor and the shortcut at the +// position's own address, with no neighbour re-read and no stencil anywhere in it. + +Tensor resample_upsample(const Tensor& x, const std::vector& pw, const Tensor& bias, + const Tensor& skip, bool halo, bool temporal, bool first); diff --git a/fastvideo-kernel/wan22_vae_decode/solution/solution.json b/fastvideo-kernel/wan22_vae_decode/solution/solution.json new file mode 100644 index 000000000..9f1850549 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/solution.json @@ -0,0 +1 @@ +{"languages":["cuda_cpp"],"entry_point":"binding.cpp::run","destination_passing_style":true,"compile_options":{"cflags":["-O3","-std=c++17"],"cuda_cflags":["-O3","-std=c++17","-lineinfo"],"ld_flags":["-lcuda"]},"dependencies":[]} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu b/fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu new file mode 100644 index 000000000..cb0f2ece8 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu @@ -0,0 +1,122 @@ +// w11: a CUTLASS implicit convolution whose epilogue stores into a caller-given destination with arbitrary +// (N,D,H,W) strides, and can accumulate on what the destination already holds (beta=1, source == destination). +// One user: the temporal upsampler, which writes the frame-interleaved video straight out (stride 2 in D). +// The resampler's phase tensors come from conv.cu's phase_conv, not from this file. +// Nothing else changes: same tile (256x256x64), same 2-CTA cluster, same TMA im2col, same bf16 epilogue. +#include +#include +#include +#include +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/convnd_problem_shape.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/conv/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/conv/device/conv_universal_adapter.hpp" +#include "cutlass/conv/kernel/conv_universal.hpp" +at::Tensor nhwc_empty(int64_t c, int64_t t, int64_t h, int64_t w, const at::Tensor& like); +using namespace cute; +using B = __nv_bfloat16; +using ElementBF16 = cutlass::bfloat16_t; +using LayoutNDHWC = cutlass::layout::TensorNDHWC; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassTensorOp; +constexpr auto Fprop = cutlass::conv::Operator::kFprop; + +// x: activation [1,C,D,H,W] channels-last; w: [K,C,KT,KH,KW] channels-last; y: the destination, with its own strides. +template +static void launch_into(const at::Tensor& x, const at::Tensor& w, at::Tensor& y, int ph, int pw, float beta, + int phu = -1, int pwu = -1) { + if (phu < 0) phu = ph; + if (pwu < 0) pwu = pw; + using Tile = Shape, Int, Shape<_64>>; + using Cluster = Shape, _1, _1>; + using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto, float, float, + ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + using Mainloop = typename cutlass::conv::collective::CollectiveBuilder< + ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float, Tile, Cluster, + cutlass::conv::collective::StageCountAutoCarveout(sizeof(typename Epilogue::SharedStorage))>, + cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp; + using Problem = cutlass::conv::ConvProblemShape; + using Kernel = cutlass::conv::kernel::ConvUniversal; + using Conv = cutlass::conv::device::ConvUniversalAdapter; + int C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4); + int K = w.size(0), KT = w.size(2), KH = w.size(3), KW = w.size(4); + int OD = y.size(2), OH = y.size(3), OW = y.size(4); + TORCH_CHECK(OD == D - KT + 1 && OH == H + ph + phu - KH + 1 && OW == W + pw + pwu - KW + 1, + "launch_into: destination shape does not match the convolution"); + TORCH_CHECK(y.size(1) == K && y.stride(1) == 1, "launch_into: destination must be channel-innermost"); + Problem problem(cutlass::conv::Mode::kCrossCorrelation, + {1, D, H, W, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)}, + {K, KT, KH, KW, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)}, + {0, ph, pw}, {0, phu, pwu}, {1, 1, 1}, {1, 1, 1}, 1); + // The output's own strides: index 0 N, 1 D, 2 H, 3 W, 4 C (the layout the C strides already use). + problem.stride_C[0] = y.stride(0); + problem.stride_C[1] = y.stride(2); + problem.stride_C[2] = y.stride(3); + problem.stride_C[3] = y.stride(4); + problem.stride_C[4] = 1; + typename Kernel::StrideC sc{}; + typename Kernel::StrideD sd{}; + cute::for_each(cute::make_seq(sc)>{}, + [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; }); + cute::for_each(cute::make_seq(sd)>{}, + [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; }); + const ElementBF16* dst = (const ElementBF16*)y.data_ptr(); + typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()}, + {{1.f, beta}, dst, sc, (ElementBF16*)y.data_ptr(), sd}}; + Conv conv; + TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess, "CUTLASS cannot implement shape"); + size_t size = Conv::get_workspace_size(args); + auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte)); + auto stream = at::cuda::getCurrentCUDAStream(); + auto status = conv.initialize(args, workspace.data_ptr(), stream); + TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialization failed ", int(status)); + status = conv.run(stream); + TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS convolution failed ", int(status)); +} + +// The temporal upsampler: y[t] = w[:,:,2] x[t-2] + w[:,:,1] x[t-1] + w[:,:,0] x[t] + bias, with frame 0 passed +// through unchanged when `first`, the two temporal halves of the 2C channels going to consecutive frames. +// p is the activation already padded with its two causal history frames ([1,C,nf+2,H,W], channels-last). +// INVARIANT: every element of y is written by head_bias_kernel, and the two parity arms below cover the +// non-head frames disjointly, each exactly once -- which is what makes beta=1 onto an at::empty safe. +static __global__ void head_bias_kernel(const B* head, const B* bias, B* y, int C, int T2, int64_t HW, int off) { + int64_t n = (int64_t)T2 * HW * (C / 8); + for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t)gridDim.x * blockDim.x) { + int cg = (int)(i % (C / 8)); int64_t r = i / (C / 8); + int64_t pos = r % HW; int t = (int)(r / HW); + uint4 v; + if (t < off) v = ((const uint4*)head)[pos * (C / 8) + cg]; + else v = ((const uint4*)(bias + ((t - off) & 1) * C))[cg]; + ((uint4*)y)[i] = v; + } +} + +at::Tensor time_upsample(const at::Tensor& p, const at::Tensor& head, const at::Tensor& w, const at::Tensor& bias, + int64_t T, int64_t nf, bool first) { + const at::Tensor& ref = p.defined() ? p : head; + int C = bias.size(0) / 2, H = ref.size(3), W = ref.size(4); + int T2 = (int)(2 * T - (first ? 1 : 0)); + at::Tensor y = nhwc_empty(C, T2, H, W, ref); + int off = first ? 1 : 0; + head_bias_kernel<<<1024, 256, 0, at::cuda::getCurrentCUDAStream()>>>( + first ? (const B*)head.data_ptr() : nullptr, (const B*)bias.data_ptr(), (B*)y.data_ptr(), C, T2, + (int64_t)H * W, off); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (nf > 0) { + int64_t sN = y.stride(0), sD = 2 * y.stride(2), sH = y.stride(3), sW = y.stride(4); + for (int par = 0; par < 2; par++) { + auto yv = y.as_strided({1, (int64_t)C, nf, (int64_t)H, (int64_t)W}, {sN, 1, sD, sH, sW}, + (int64_t)(off + par) * y.stride(2)); + launch_into<256, 256, 2>(p, w.narrow(0, (int64_t)par * C, C), yv, 0, 0, 1.f); + } + } + return y; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu b/fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu new file mode 100644 index 000000000..c94ce25fb --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu @@ -0,0 +1,400 @@ +// Fused F(4x4,3x3) inverse transform for the causal 3x3x3 Winograd convolutions. +// One pass over the 36 component tensors M produces BOTH +// y = A^T M A + bias + residual (raw conv output: the resnet's residual operand) +// n = silu(rmsnorm(y) * sqrt(C) * gain) (the next convolution's normalized, activated, +// temporal+spatial padded input) +// so the normalized tensor is never written and read again by a separate norm pass. +// The RMS statistic needs every output channel of a position, so one block holds all of C +// for one 4x4 output tile: blockDim.x == C/2, one channel pair per thread. +// +// M layout: 36 separate allocations, each [D][TH][TW][C] bf16, C innermost, so a plane's +// position stride is 1 in C and its plane offset is (((t*TH+ty)*TW+tx)*C + ch). +// y is [D][H][W][C]; n is [D+PT][H+2PH][W][C]. +// +// Two lessons from the guild, both measured elsewhere and both applied here: +// - w01's J576 (input transform): the head carries no divisions, no predicates and no 64-bit +// address chain. 3-D grid (tx,ty,t), a separate EXACT instantiation with the bounds test +// removed, and a 32-bit plane byte offset. EXACT is dispatched per call: it holds at +// 88x156, 176x312 and 352x624, whose widths are multiples of 4 -- NOT at 44x78, where +// W=78 leaves the last tile two columns short and the checked path is taken. +// - w05's J627 (inverse transform, 5.5 TB/s at 64 registers): no thread loads 4 B of one plane. +// The block stages the 36 components of its tile into shared memory cooperatively -- one +// plane's C channels is C*2 contiguous bytes, so a warp pass is 16 B per lane and one fully +// coalesced 512 B request -- and the consumer reads smem[comp][2c] as a 4-byte pair, which is +// 128 contiguous bytes per warp and therefore conflict-free. No transpose on either side. +// +// Per-element work follows w02's recipe (M137, PR15): bf16->f32 by an integer shift so the +// conversions leave the XU pipe, one bf16 rounding for the whole normalize*sqrt(C)*gain product +// instead of the reference's three, SiLU as u + u*tanh(u) with u = pre-activation/2 and +// tanh.approx.f32, and 0.5*sqrt(C)*rsqrt(sum) as one per-position scalar. +#include "kernels.h" +#include +#include +#include +#include +using B=__nv_bfloat16; +static cudaStream_t e_stream(){return at::cuda::getCurrentCUDAStream();} +static const B* e_ptr(const Tensor& t){return t.defined()?(const B*)t.data_ptr():nullptr;} +static B* e_mut(Tensor& t){return t.defined()?(B*)t.data_ptr():nullptr;} + +struct EpiPtrs{const B* p[36];}; + +// A of F(4x4,3x3) applied along one axis: out[i] = sum_v A[i][v] v[v], two fp32 lanes per +// instruction (a6v below is the only form this file uses). add/sub/fma.f32x2 need sm_100+; this +// tree is built for sm_100a only (solution.json), so no scalar fallback is kept beside it -- a +// scalar copy would be dead source here, which is why the scalar `a6` was removed rather than left. +// f32x2: PTX 8.6 gives add/sub/mul/fma on a .b64 operand holding two f32 lanes (sm_100a). +// This kernel already carries two independent fp32 lanes per thread -- the channel pair a thread +// owns -- so `ch` IS the SIMD lane and nothing crosses a lane. Each instruction performs exactly +// the arithmetic the scalar instruction performs, in the same association. +struct V2{ + float x,y; + __device__ V2(){} __device__ V2(float a,float b):x(a),y(b){} +}; +__device__ __forceinline__ unsigned long long pk(V2 a){ + return (unsigned long long)__float_as_uint(a.y)<<32|__float_as_uint(a.x); +} +__device__ __forceinline__ V2 upk(unsigned long long u){ + V2 r; r.x=__uint_as_float((unsigned)u); r.y=__uint_as_float((unsigned)(u>>32)); return r; +} +__device__ __forceinline__ V2 vadd(V2 a,V2 b){ + unsigned long long r; asm("add.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r); +} +__device__ __forceinline__ V2 vsub(V2 a,V2 b){ + unsigned long long r; asm("sub.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r); +} +__device__ __forceinline__ V2 vfmas(V2 a,float c,V2 acc){ + V2 k; k.x=k.y=c; unsigned long long r; + asm("fma.rn.f32x2 %0,%1,%2,%3;":"=l"(r):"l"(pk(a)),"l"(pk(k)),"l"(pk(acc))); return upk(r); +} +__device__ __forceinline__ V2 b2v(uint32_t w){ + V2 r; r.x=__uint_as_float(w<<16); r.y=__uint_as_float(w&0xffff0000u); return r; +} +// The transform of A above, two lanes per instruction. +__device__ __forceinline__ void a6v(V2 v0,V2 v1,V2 v2,V2 v3,V2 v4,V2 v5,V2* y){ + V2 a=vadd(v1,v2), b=vsub(v1,v2), c=vadd(v3,v4), d=vsub(v3,v4); + y[0]=vadd(vadd(v0,a),c); y[1]=vfmas(d,2.f,b); y[2]=vfmas(c,4.f,a); y[3]=vadd(vfmas(d,8.f,b),v5); +} +// Exact bf16 pair -> two floats without the XU pipe: a bf16 is the top half of an fp32. +__device__ __forceinline__ float2 b2f(uint32_t w){ + return make_float2(__uint_as_float(w<<16),__uint_as_float(w&0xffff0000u)); +} +__device__ __forceinline__ float tanh_ap(float x){float r;asm("tanh.approx.f32 %0, %1;":"=f"(r):"f"(x));return r;} +// silu(x) with u = x/2: x*sigmoid(x) = u + u*tanh(u). +__device__ __forceinline__ float silu2(float u){return __fmaf_rn(u,tanh_ap(u),u);} +__device__ __forceinline__ uint32_t add2(uint32_t a,uint32_t b){ + __nv_bfloat162 x=*(const __nv_bfloat162*)&a,y=*(const __nv_bfloat162*)&b; + __nv_bfloat162 r=__hadd2(x,y); return *(const uint32_t*)&r; +} +__device__ __forceinline__ uint32_t pack2(float x,float y){ + __nv_bfloat162 r=__floats2bfloat162_rn(x,y); return *(const uint32_t*)&r; +} +__device__ __forceinline__ float sq8(const uint4& v){ + float2 a=b2f(v.x),b=b2f(v.y),c=b2f(v.z),d=b2f(v.w); + return (__fmaf_rn(a.x,a.x,a.y*a.y)+__fmaf_rn(b.x,b.x,b.y*b.y)) + +(__fmaf_rn(c.x,c.x,c.y*c.y)+__fmaf_rn(d.x,d.x,d.y*d.y)); +} + +// Stage the 36 components of one 4x4 tile into shared memory: 36*C/8 tasks of 16 bytes, +// 9 per thread at any C. Each warp's 32 lanes take 32 consecutive 16-byte slots, i.e. 512 +// contiguous bytes of one plane. +// +// The copy is `cp.async.cg` (16 B, global->shared, L2-cached, NO register round trip) rather than +// the `uint4` load-then-store it replaces. Measured reason (J3871, wino_epi_kernel<1024,1>, ncu, +// 3 passes): the binding FIXABLE stall is `mio_throttle 2.11` cycles per issued instruction -- +// second only to `not_selected 2.73`, which is not a defect (it means the SM already has eligible +// warps to pass over) -- and `l1tex__throughput` is the highest pipe at 62.7-64.7 %. The old form +// is a global load AND a shared store per 16 B, so it spends two MIO slots and a writeback +// register per slot; this form spends one slot and no register. Per block at C=1024 that is +// 9*512 = 4,608 pairs, i.e. 9,216 MIO instructions becoming 4,608. +// +// Bytes moved, alignment and the resulting V are IDENTICAL: same comp/off decomposition, same +// 16-byte slot, same plane offset inside smem, so the transform reads what it read before and the +// output is bit-for-bit the old output. This is a scheduling change, not a numerical one. +// +// The commit/wait pair sits here rather than at the call site because the transform consumes the +// staging immediately -- there is no independent work to overlap the transfer with inside one +// tile. Overlap across tiles needs a persistent loop and is a separate change, not this one. +// The form above is measured, not preferred: `cp.async.cg` beats the `uint4` load+store it +// replaced by 12.7 ms at f121x1 (J4745, one job, one clock, A/A 0.2 ms). +template +__device__ __forceinline__ void stage36(B* sm,const EpiPtrs& mp,unsigned src_tile){ + constexpr int CV=C/2; + constexpr int PER=C/8; // 16-byte slots per plane + const unsigned smbase=(unsigned)__cvta_generic_to_shared(sm); + #pragma unroll + for(int i=0;i<9;i++){ + const int task=threadIdx.x+i*CV; + const int comp=task/PER, off=task%PER; + const unsigned dst=smbase+(unsigned)((comp*C+off*8)*2); + const char* src=(const char*)mp.p[comp]+src_tile+off*16; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"::"r"(dst),"l"(src)); + } + asm volatile("cp.async.commit_group;\n"::); + asm volatile("cp.async.wait_group 0;\n"::); +} + +template +__global__ __launch_bounds__(C/2) void wino_epi_kernel( + EpiPtrs mp,B* __restrict__ y,B* __restrict__ n, + const B* __restrict__ bias,const B* __restrict__ residual,const B* __restrict__ gain, + int H,int W,int TH,int TW,int PT,int PH,bool wy){ + // `wy` is the CALLER's assertion that nobody reads this launch's y. It is only ever false where + // the very next conv has idx == fnext and therefore takes the fused `prepared` edge (see + // binding.cpp: resnet's first conv, whose consumer is its own second conv). The n tensor is + // written either way, so this can never be false without a n to carry the activation. + constexpr int CV=C/2; // channel pairs per position + constexpr int WP=CV/32; // warps per block + const int c=threadIdx.x; + const int tx=blockIdx.x, ty=blockIdx.y, t=blockIdx.z; + extern __shared__ __align__(16) B smem[]; + B* sm=smem; // 36 planes x C, this tile + // 32-bit plane byte offset: D*TH*TW*C*2 bytes, where D is the PIXEL frame count at this level + // (4*chunk-3 at up block 1, not the latent chunk size -- the earlier "34" was the latent count). + // chunk_size is NOW min(16*world, 61) -- this PR is what moved the rule off `8*world`, so chunk + // 32 at world 4 no longer exists and the worst case moves with it. Over the shipping configs the + // largest src_tile is world 2 / chunk 32 = 4.39e8 bytes and world 1 / chunk 16 = 4.29e8, both + // ~10x inside UINT32_MAX and both enforced at run time by the guard further down this file. + // The `8.8e8` that stood here was 4x too large and it described a configuration this PR retires. + // A chunk size that more than doubles again, or a wider clip, must widen this to unsigned long long. + const unsigned src_tile=(unsigned)((((t*TH+ty)*TW+tx)*C)*2u); + + stage36(sm,mp,src_tile); + __syncthreads(); + + V2 acc[4][4]; + #pragma unroll + for(int i=0;i<4;i++) + #pragma unroll + for(int b=0;b<4;b++) acc[i][b]=V2(0.f,0.f); + + const B* q0=sm+2*c; // this thread's channel pair, in each plane + #pragma unroll + for(int u=0;u<6;u++){ + V2 f[6]; + #pragma unroll + for(int v=0;v<6;v++) f[v]=b2v(*(const uint32_t*)(q0+(6*u+v)*C)); + V2 q[4]; + a6v(f[0],f[1],f[2],f[3],f[4],f[5],q); + #pragma unroll + for(int b=0;b<4;b++){ + const V2 v=q[b]; + if(u==0) acc[0][b]=vadd(acc[0][b],v); + else if(u==1){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vadd(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vadd(acc[3][b],v);} + else if(u==2){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vsub(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vsub(acc[3][b],v);} + else if(u==3){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,8.f,acc[3][b]);} + else if(u==4){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,-2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,-8.f,acc[3][b]);} + else {acc[3][b]=vadd(acc[3][b],v);} + } + } + + const int Hp=H+2*PH; + float2 biasv=make_float2(0.f,0.f),gainv=make_float2(0.f,0.f); + if(bias) biasv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)bias)[c]); + if(gain) gainv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)gain)[c]); + const long yb=(((long)t*H+4*ty)*W+4*tx)*CV+c; + + const __nv_bfloat162 bv=__floats2bfloat162_rn(biasv.x,biasv.y); + if(!n){ + // No prepared tensor: write the raw conv output in the natural per-thread layout. + #pragma unroll + for(int i=0;i<4;i++){ + const int yy=4*ty+i; + #pragma unroll + for(int b=0;b<4;b++){ + const int xx=4*tx+b; + const long off=(long)(i*W+b)*CV; + const bool ok=EXACT?true:(yy>5, ln=c&31; + const float hs=0.5f*sqrtf((float)C); + uint4 gv[NU]; + // Chunk map: lane ln takes chunk (k*32+ln), i.e. channel (k*32+ln)*8, so the 32 lanes of one + // warp instruction cover 32 consecutive chunks = 512 contiguous bytes. The previous map gave + // lane ln the block [ln*CW, ln*CW+CW), which puts neighbours CW*2 bytes apart: contiguous at + // C=256 (16 B), half a line at C=512 (32 B), a quarter at C=1024 (64 B) -- and at C=1024 that + // is a 4x sector over-fetch on the residual read and on both stores. w21's J1096 gives it in + // counters (L1 hit 5.32%, L2 hit 0.20%, DRAM 65.6%): no reuse, so every over-fetched sector + // is paid at DRAM. Measured at D=8, min of three rounds: C=256 1.000x (identical expression at + // NU=1), C=512 1.244x, C=1024 1.431x. It permutes which channel each k reads, so the 4-term + // sum of squares accumulates in a different order and s rounds differently in its last bits -- + // f121x1 max|d| 0.1605 -> 0.147, rms unchanged to five digits. + #pragma unroll + for(int k=0;k>2), xx=4*tx+(p&3); + const bool ok=EXACT?true:(yy>=1) sq+=__shfl_xor_sync(0xffffffffu,sq,d); + if(ok){ + const float s=hs*rsqrtf(fmaxf(sq,1e-24f)); + const long ybase=((((long)t*H+yy)*W+xx)*C)*2; + const long nbase=(((((long)(t+PT)*Hp+PH+yy)*W+xx)*C))*2; + #pragma unroll + for(int k=0;k +static void launch_epi(const dim3& grid,const dim3& block,size_t smem, + const EpiPtrs& p,B* y,B* n, + const B* b,const B* r,const B* g, + int H,int W,int TH,int TW,int PT,int PH,bool wy,cudaStream_t s){ + auto k=wino_epi_kernel; + if(smem>48*1024){ + static bool done[2]={false,false}; + if(!done[EXACT?1:0]){ + C10_CUDA_CHECK(cudaFuncSetAttribute(k,cudaFuncAttributeMaxDynamicSharedMemorySize,36*C*2)); + done[EXACT?1:0]=true; + } + } + k<<>>(p,y,n,b,r,g,H,W,TH,TW,PT,PH,wy); +} + +WinoOut wino_epilogue(const std::vector& m,int H,int W,const Tensor& bias, + const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph, + bool write_y){ + // `write_y=false` is admissible ONLY when the caller's next conv has idx == the `fnext` this call + // prepares for, so that next conv takes the `prepared` edge and never reads the returned y. Out.y + // is still a distinct allocation: its unread contents are undefined, not zero. + TORCH_CHECK(m.size()==36&&m[0].dim()==5,"wino_epilogue: M list ",m.size()," rank ",m[0].dim()); + const int C=m[0].size(1),D=m[0].size(2),TH=m[0].size(3),TW=m[0].size(4); + TORCH_CHECK(C==256||C==512||C==1024,"wino_epilogue: unsupported channel count"); + auto s=e_stream(); + EpiPtrs p; for(int i=0;i<36;i++) p.p[i]=(const B*)m[i].data_ptr(); + WinoOut out; + out.y=nhwc_empty(C,D,H,W,m[0]); + const int Hp=H+2*ph; + if(gamma.defined()){ + out.n=nhwc_empty(C,D+pt,Hp,W,m[0]); + B* np=e_mut(out.n); + const long pairs=(long)(D+pt)*Hp*W*(C/2); + if(pt>0&&!hist.defined()) + epi_zero<<<(int)std::min((long)2048,((long)pt*Hp*W*(C/2)+4095)/4096),256,0,s>>>(np,(long)pt*Hp*W*(C/2)); + if(ph>0) + epi_border<<<(int)std::min((long)2048,(pairs+4095)/4096),256,0,s>>>(np,C,D+pt,Hp,W); + if(pt>0&&hist.defined()){ + // `hist` is a `last_frames` output: a COMPACT [C,pt,H,W] channels-last-3d tensor, so its + // per-frame SOURCE pitch is H*W*C -- unpadded. Only the DESTINATION is padded: out.n is + // [C,D+pt,Hp,W] with pitch Hp*W*C, and these frames land in its interior rows [ph,ph+H). + // The two pitches are one of each and neither is a typo: passing Hp*W*C as the source pitch + // would read pt frames of garbage, and passing H*W*C as the destination pitch would write + // the pad rows away. If `last_frames` changes its shape or its pitch, this copy changes. + const size_t row=(size_t)H*W*C*sizeof(B); + C10_CUDA_CHECK(cudaMemcpy2DAsync(np+(long)ph*W*C,(size_t)Hp*W*C*sizeof(B), + e_ptr(hist),row,row,pt,cudaMemcpyDeviceToDevice,s)); + } + } + // src_tile in the kernel is the only 32-bit index in this file, so assert here what it relies on + // rather than trusting the kernel's comment. (The y and n offsets in the kernel body are formed in + // 64-bit and are not bounded by this.) Checked against UINT32_MAX, not INT32_MAX: the index is + // unsigned, and a check against the signed bound would refuse shapes that are in fact safe. + { const unsigned long long off=(((unsigned long long)(D-1)*TH+(TH-1))*TW+(TW-1))*(unsigned long long)C*2ull; + TORCH_CHECK(off<=0xFFFFFFFFull,"wino_epilogue: 32-bit tile offset ",off, + " exceeds UINT32_MAX; widen src_tile in wino_epilogue.cu"); } + const bool exact=(TH*4==H&&TW*4==W); + const dim3 grid(TW,TH,D); + const dim3 block(C/2); + const size_t shmem=(size_t)36*C*2; + B* yp=e_mut(out.y); B* np=e_mut(out.n); + const B* bp=e_ptr(bias); const B* rp=e_ptr(residual); const B* gp=e_ptr(gamma); + if(C==256){ if(exact) launch_epi<256,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); + else launch_epi<256,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); } + else if(C==512){ if(exact) launch_epi<512,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); + else launch_epi<512,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); } + else { if(exact) launch_epi<1024,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); + else launch_epi<1024,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu b/fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu new file mode 100644 index 000000000..0b73fc350 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu @@ -0,0 +1,247 @@ +// F(4x4,3x3) spatial input transform of the Winograd path of the Wan2.2 VAE decoder. +// +// ONE PASS: `winograd_input_fast` replaces `kernels.cu`'s `winograd_input` for the 18 Winograd +// convolutions (1024/512/256 channels; any other channel count keeps the general kernel), and +// its output V is BIT-IDENTICAL to it -- same patch origin, same block-uniform fast-path +// condition, same B^T-then-B order with the in-place column update, same 6a+b plane index, same +// single bf16 rounding at the store, same plane stride. Nothing else in the tree moves but the +// one call site in binding.cpp and one declaration in kernels.h. +// +// Why it is faster, from the profile of the kernel it replaces (w01's J549): 581.8 M +// warp-instructions in 1.031 ms at 59.5% issue-active, of which 43.5% are ALU/integer against +// 47.6% FMA and 4.9% LSU; DRAM 23.6%, FMA pipe 32.4%. It sits at 82% of its own instruction +// issue bound and the integer address math is the only large removable block: +// 1. a 3-D grid, blockIdx = (tx, ty, t) and threadIdx = the channel pair, so nothing divides; +// 2. one block-uniform fast path for the interior tiles, where the 36 loads carry no +// predicate at all -- only the first and last tile of a row and the first and last tile +// row can touch the time pad or the frame edge -- and the patch row pointer advances by +// the row stride, so the 36 addresses are 36 LSU instructions with folded offsets; +// 3. the 36 plane stores walk one pointer by V's plane stride instead of doing an int64 +// multiply per store. +// +// MORE BLOCKS PER TILE INSTEAD OF ONE: SPLIT = 4 AT C = 1024, 2 AT C = 512 AND 256. +// A tile's C/2 channel pairs used to be covered by one block of C/2 threads, so `blockDim` and +// `C` were the same variable and the grid held exactly one block per tile. A block is now the +// tile's width divided by SPLIT, and SPLIT blocks of the grid cover one tile. At C = 1024 that +// moves the level from one block of 16 warps (32.08 ms) to five blocks of 4 warps (22.42 ms, +// 20 warps resident); at C = 512 and C = 256 the same move is worth 1.05 % and -0.4 % -- i.e. +// nothing -- so those two stay at SPLIT = 2. See the table below for why. +// +// `__launch_bounds__(CP, 512 / CP)` IS LOAD-BEARING AND IS HALF THE MEASUREMENT. It pins the +// register budget at 65536/512 = 128 at EVERY width -- exactly what main's +// `__launch_bounds__(512, 1)` grants at its one wide instantiation. Without it ptxas may spend +// up to the architectural cap of 255 registers per thread (65536/CP is itself 512 at CP = 128, +// so it is the 255 that binds there), floor(65536/(CP * REG)) falls, the resident warps follow, +// and the same arm runs 2.25x SLOWER. The counterfactual is measured at ONE instantiation: +// J3705 ran `wino_in_kernel_st<256,2>` both ways -- `__launch_bounds__(256, 1)` at 59.11 ms +// against `(256, 2)` at 26.26 ms, one integer and one body. At the shipped CP = 128 the no-pin +// case gives floor(65536/(128*255)) = 2 blocks of 4 warps against the pinned 5, i.e. 20 resident +// warps down to 8 -- a 2.5x cut, not a halving, and the reason C = 1024 needs the pin most. +// The declaration and the launch are one change; neither is the arm without the other. +// +// Measured (J3868, one build, one process, A B B B B A over SPLIT = 1 2 4 8 2 1, f121x1 +// --kernels, 1132 MHz on every leg, exit 0). Both control legs are shown so the drift is visible: +// SPLIT=1 SPLIT=2 SPLIT=4 SPLIT=8 2 -> 4 +// C = 1024 (shipped) 32.08/32.08 26.26/26.53 22.42 22.27 1.1775 +// C = 512 (2 kept) 64.06/63.92 54.60/55.44 54.45 55.45 1.0105 +// C = 256 (2 kept) 60.47/60.66 59.56/60.05 60.04 (n/a) 0.9962 +// CALL 864.7 847.6 844.2 845.2 +// A/A 0.3 %, B/B 1.0-1.5 %; max |d| 0.1543 of 0.4984, rms 0.002093 of 0.005847, PSNR 59.6 dB, +// IDENTICAL to every printed digit on all six legs. +// The C = 256 SPLIT = 8 cell is `n/a`, not a decline: J3868's probe only instantiated the CP +// values it needed and SPLIT = 8 at C = 256 would ask for CP = 16, which that build did not +// contain. The head's launcher is an exhaustive switch with no fall-back path of any kind. +// +// ONLY THE 2 -> 4 MOVE IS A C = 1024-ONLY GAIN, and that is why the other two keep SPLIT = 2: +// the shipped SPLIT = 2 already takes the 1 -> 2 step on BOTH of them -- 64.06 -> 54.60 at +// C = 512 is 1.1733x and 60.47 -> 59.56 at C = 256 is 1.0183x, the +15.5 % and +1.5 % below -- +// so neither is a rung that gains nothing. What is declined on them is the FOURTH block: +// C = 512's 1.0105 is inside the B/B drift (its own legs read 54.60 and 55.44, a 1.5 % spread); +// C = 256's 0.9962 REVERSES -- its legs read 59.56, 60.05 and 60.66, a 1.8 % spread, and 60.04 +// sits inside it, so the option is declined rather than merely unpriced; +// SPLIT = 8 at C = 1024 is 22.27 against 22.42 -- 0.7 % at a 1.0 % drift, a tie, not a win. +// +// WHY C = 1024 AND NOT THE OTHERS. With REG 96 the realized resident warps are +// floor(65536/(96*blockDim)) * blockDim/32, i.e. 16 at blockDim 512 or 256, 20 at 128 or 64, and +// 21 at 32. Main's C = 1024 case used blockDim 512, which rounds to ONE block and never reaches +// past 16 warps although 2048/96 = 21.3 is its ceiling; `__launch_bounds__` cannot fix that, +// because 2 * 512 * 96 = 98,304 > 65,536. Only a narrower block can, and narrowing it is SPLIT. +// C = 512 and C = 256 already reach 20 warps at SPLIT = 2, so there is nothing there to buy. +// The three levels' rates were 2.609 / 3.340 / 3.760 TB/s (J3195's counted bytes). DERIVED, not +// re-measured: at fixed byte counts rate_after/rate_before = t_before/t_after, so from J3868's own +// legs the after-rates are 2.609*32.08/22.42 = 3.733, 3.340*64.06/55.44 = 3.859 and +// 3.760*60.47/59.56 = 3.817 TB/s -- a 44.1 % spread closed to 3.4 %. The largest move is the +// slow level's own (+43 %), and the fast two barely move (+15.5 % and +1.5 %), so what happened is +// that the slow level was walked up to where the other two already sat, not that a ceiling moved. +// An earlier draft of this comment stated 3.952 for C = 256, which needs 60.47 -> 57.5 ms and no +// leg of J3868 is below 59.56; it was a number carried over from a withdrawn row and it is wrong. +// +// `J3761`'s rebuild of the round-1 source: BUILD_RC=0, REG:96 STACK:0 SHARED:0 LOCAL:0. +// `J3889`'s build of this source: the same, with `wino_in_kernel<128,4>` at REG:95 (one register +// lower, which changes no block count: 65536/(128*95) = 5.39 -> 5, the same as at 96). +// Correctness at worlds 1 and 4 on all three short clips: J3889 f1x1/f5x1/f17x1 ok at +// x0.41/x0.41/x0.65 and J3908 f1x4/f5x4/f17x4 ok at x0.45/x0.45/x0.66, every leg inside main's +// own margins. +// +// The share of the call is flat in clip length (17.6 % at f121x1, 17.9 % at f481x1) and this +// function takes no rank, no world and no process group, so a world-1 cell is the geomean cell. +// +// Measured by w04 on main's own tree (kernel#1, this change alone, f121x1 at 1849 MHz): the +// three instantiations total 114.4 ms against 221.8 ms of `winograd_input_kernel`, and the call +// 969.9 -> 862.5 ms (-11.1%), with V bit-identical -- f1/f5/f17 keep the seed's max |d| and rms +// to the digit -- and f481x1 ok at x0.38 of the max and x0.41 of the rms allowance. w21 wrote +// the kernel; w08 owns this file. +#include "kernels.h" +#include +#include +#include +#include + +using B = __nv_bfloat16; +static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); } + +__device__ __forceinline__ void wbt6(float v0, float v1, float v2, float v3, float v4, float v5, float* y) { + float a = v4 - 4 * v2, b = v3 - 4 * v1, c = v4 - v2, d = 2 * (v3 - v1); + y[0] = 4 * v0 - 5 * v2 + v4; y[1] = a + b; y[2] = a - b; y[3] = c + d; y[4] = c - d; y[5] = 4 * v1 - 5 * v3 + v5; +} + +// One block per (part-tile, frame); threadIdx.x = the channel pair within this block's share. +// CP = C/2/SPLIT is the block size. The shipped dispatch is SPLIT = 4 at C = 1024 and SPLIT = 2 +// at C = 512 and 256, so blockDim is 128 / 128 / 64 at C = 1024 / 512 / 256 -- but SPLIT is NOT +// uniform and no comment here may claim it is. The invariant that actually guarantees coverage is +// CP * SPLIT == C / 2 (128*4 = 128*2 = 64*2 = C/2) +// enforced by the `static_assert` in `winograd_input_fast` below, together with `grid.x == SPLIT*TW` +// from the same `constexpr`. `(blockIdx.x % SPLIT) * CP + threadIdx.x` then walks [0, SPLIT*CP) = +// [0, C/2) exactly once -- every channel pair of a tile is written once and no pair twice -- and +// `blockIdx.x / SPLIT` over a grid of SPLIT*TW hits each tile exactly SPLIT times, for any SPLIT: +// the remainder cycles 0..SPLIT-1 within each group of SPLIT consecutive blockIdx, so TW need not be +// a multiple of SPLIT. +// +// SPLIT = 1 with CP = C/2 gives back main's one-block-per-tile BODY, but not its attribute (this +// one is `__launch_bounds__(CP, 512/CP)`, which equals main's `(CP, 1)` only at CP = 512) and not +// its grid shape, so do not use it as an exact revert; revert by checking out the parent commit. +// +// The launched blockDim must equal CP. It does at all three call sites below, from the same +// template value, so the two cannot drift apart. +template +__global__ void __launch_bounds__(CP, 512 / CP) +wino_in_kernel(const B* __restrict__ x, B* __restrict__ v, int C, int D, int H, int W, int PH, int TH, int TW) { + const int cp = (blockIdx.x % SPLIT) * CP + threadIdx.x; + const int tx = blockIdx.x / SPLIT; + const int ty = blockIdx.y, t = blockIdx.z; + const int HP = H + 2 * PH; + const int y0 = 4 * ty - 1 + PH, x0 = 4 * tx - 1; + const int64_t rowStep = (int64_t)W * C; + const B* base = x + (int64_t)t * HP * rowStep + (int64_t)y0 * rowStep + (int64_t)x0 * C + 2 * cp; + float d[2][6][6]; + // Uniform across the block, so it is one branch and not 36 predicates. + if ((y0 >= 0) && (y0 + 5 < HP) && (x0 >= 0) && (x0 + 5 < W)) { + const B* prow = base; +#pragma unroll + for (int a = 0; a < 6; a++) { +#pragma unroll + for (int b = 0; b < 6; b++) { + float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C)); + d[0][a][b] = f.x; d[1][a][b] = f.y; + } + prow += rowStep; + } + } else { + const B* prow = base; +#pragma unroll + for (int a = 0; a < 6; a++) { + const int y = y0 + a; + const bool ok = (y >= 0) && (y < HP); +#pragma unroll + for (int b = 0; b < 6; b++) { + const int xx = x0 + b; + if (ok && xx >= 0 && xx < W) { + float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C)); + d[0][a][b] = f.x; d[1][a][b] = f.y; + } else { + d[0][a][b] = 0.f; d[1][a][b] = 0.f; + } + } + prow += rowStep; + } + } + // B^T d down the patch's 6 columns, then B along its 6 rows, one bf16 rounding at the end. +#pragma unroll + for (int k = 0; k < 2; k++) +#pragma unroll + for (int b = 0; b < 6; b++) { + float q[6]; + wbt6(d[k][0][b], d[k][1][b], d[k][2][b], d[k][3][b], d[k][4][b], d[k][5][b], q); +#pragma unroll + for (int a = 0; a < 6; a++) d[k][a][b] = q[a]; + } + B* dst = v + (int64_t)(((t * TH + ty) * TW + tx) * C + 2 * cp); + const int64_t plane = (int64_t)D * TH * TW * C; +#pragma unroll + for (int a = 0; a < 6; a++) { + float q[2][6]; +#pragma unroll + for (int k = 0; k < 2; k++) + wbt6(d[k][a][0], d[k][a][1], d[k][a][2], d[k][a][3], d[k][a][4], d[k][a][5], q[k]); +#pragma unroll + for (int b = 0; b < 6; b++, dst += plane) *(__nv_bfloat162*)dst = __floats2bfloat162_rn(q[0][b], q[1][b]); + } +} + +Tensor winograd_input_fast(const Tensor& x, int ph) { + const int C = x.size(1), D = x.size(2), H = x.size(3) - 2 * ph, W = x.size(4); + const int TH = (H + 3) / 4, TW = (W + 3) / 4; + // The 18 Winograd convolutions are the 1024-, 512- and 256-channel ones; anything else + // keeps the general kernel. + if (C != 1024 && C != 512 && C != 256) return winograd_input(x, ph); + auto v = at::empty({36, D, TH, TW, C}, x.options()); + // SPLIT = 4 at C = 1024, 2 at C = 512 and 256. CP * SPLIT = C/2 at every case, so a tile's + // channel pairs are covered exactly once; grid.x is SPLIT * TW, a block count and not a tile + // count. blockDim is CP, taken from the same template value, so the two cannot drift. + // + // Why 4 only at C = 1024 (J3868: 2 -> 4 on the rung at each C, one build, one process): + // C=1024 <256,2> 26.26/26.53 -> <128,4> 22.42 1.1775 warps 16 -> 20 + // C= 512 <128,2> 54.60/55.44 -> < 64,4> 54.45 1.0105 warps 20 -> 20 inside the 1.5 % drift + // C= 256 < 64,2> 59.56/60.05 -> < 32,4> 60.04 0.9962 warps 20 -> 21 it REVERSES + // So the resolved gain is C = 1024 alone, and it is the only rung whose split buys warps: + // with REG 96 the realized resident warps are floor(65536/(96*blockDim))*blockDim/32, i.e. + // 16 at blockDim 512 or 256, 20 at 128 or 64, 21 at 32. The step off 16 is the whole of it. + // SPLIT = 8 at C = 1024 measured 22.27 against 22.42 -- 0.7 % at a 1.0 % B/B drift, not + // resolvable, so it is not taken. A tie is not a win. + // + // ONE `constexpr` PER CASE BINDS THE TEMPLATE ARGUMENT, THE BLOCK WIDTH, THE GRID MULTIPLIER AND + // THE COVERAGE INVARIANT. Three literals in three places is a desync waiting to become a bug: + // raise SPLIT without raising grid.x and `tx = blockIdx.x / SPLIT` stays inside TW, the kernel + // silently computes one block's share of every tile and the rest of V is never written; raise + // grid.x without SPLIT and `tx` runs past TW, computing x0 = 4*tx-1 for tiles that do not exist, + // which is an out-of-bounds read AND write rather than a slowdown. The `static_assert` makes the + // first a compile error and the single `constexpr` makes the second unrepresentable. + const B* xp = (const B*)x.data_ptr(); + B* vp = (B*)v.data_ptr(); + auto st = stream(); + switch (C) { + case 1024: { + constexpr int SPLIT = 4, CP = 1024 / 2 / SPLIT; + static_assert(CP * SPLIT == 1024 / 2, "CP * SPLIT must equal C/2 or a tile is not covered"); + const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D); + wino_in_kernel<<>>(xp, vp, C, D, H, W, ph, TH, TW); + break; + } + case 512: { + constexpr int SPLIT = 2, CP = 512 / 2 / SPLIT; + static_assert(CP * SPLIT == 512 / 2, "CP * SPLIT must equal C/2 or a tile is not covered"); + const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D); + wino_in_kernel<<>>(xp, vp, C, D, H, W, ph, TH, TW); + break; + } + case 256: { + constexpr int SPLIT = 2, CP = 256 / 2 / SPLIT; + static_assert(CP * SPLIT == 256 / 2, "CP * SPLIT must equal C/2 or a tile is not covered"); + const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D); + wino_in_kernel<<>>(xp, vp, C, D, H, W, ph, TH, TW); + break; + } + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return v; +} diff --git a/fastvideo-kernel/wan22_vae_decode/tools/audit.py b/fastvideo-kernel/wan22_vae_decode/tools/audit.py new file mode 100644 index 000000000..4d3026938 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/tools/audit.py @@ -0,0 +1,359 @@ +"""The stack rules of a solution, checked on its sources; and who launched each kernel of a call, found at run time. + +The stack is CUDA C++ (with inline PTX), compiled ahead of time into one extension module (see loader.py), with +communication through the ProcessGroup it is given (NCCL) or CUDA peer access and IPC between the ranks. + +Static rules (check): solution.json declares the languages in LANGUAGES, and RULES maps a rule's name to what it +looks for in the C/C++/CUDA sources and why. Kinds of match: + + include the target of an #include + text the source with comments removed + flag solution.json: a compile option + vendored an #include target that must resolve to a file of the solution itself + +A task's spec.py may extend the table: RULES = {name: rule, or None to drop a default rule}. + +Run time (trace): one call runs under CUPTI. Every kernel is attributed to the library that launched it, from the +native call stack of its launch: kernels launched from inside cuBLAS/cuBLASLt, cuDNN or ATen (libtorch) are library +kernels, and so are Triton's (a stack the rules exclude); NCCL's are communication, which the rules allow, counted +apart; kernels launched from the solution's own module or by driver calls are the solution's, whatever their names, +so kernels built from CUTLASS/CuTe templates copied into the solution count as its own. Kernels whose launch cannot be +traced (graph replays, device-side launches) are classified by name. Library kernels are allowed, as placeholders: +their count and their share of the call's GPU time show where the solution still depends on them. Only kernel +launches call back into this module (enabling the whole driver API would call back on NCCL's watchdog thread, which +queries events while it holds a lock the launching thread needs: a deadlock under the GIL). + + python tools/audit.py [solution] the static rules +""" + +import argparse +import collections +import ctypes +import json +import os +import re +import sys + +CPP_FILES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh", ".hh", ".hxx", ".inl", ".inc") +LANGUAGES = ["cuda_cpp"] + +RULES = { + "vendored-cutlass": { + "vendored": r"(cutlass|cute)/.*", + "why": "CUTLASS/CuTe headers are copied into the solution and maintained there (cutlass/..., cute/... next to " + "the entry file), not taken from the installed CUTLASS"}, + "runtime-build": { + "include": r"(.*/)?(nvrtc|nvJitLink|nvPTXCompiler)\.h", + "text": r"\b(nvrtc[A-Z]\w*|nvJitLink[A-Z]\w*|nvPTXCompiler\w*|cuModuleLoad\w*|" + r"cuLink(Create|AddData|AddFile|Complete)\w*|cuLibraryLoad\w*|cudaLibraryLoad\w*)\s*\(", + "flag": r"-l(nvrtc|nvJitLink|nvptxcompiler)\w*", + "why": "nothing is compiled or loaded at run time (NVRTC, nvJitLink, PTX or cubins loaded through the driver): " + "device code, inline PTX included, is compiled by nvcc with the solution"}, + "python-from-cpp": { + "text": r"\b(py(bind11)?::module_?::import|py(bind11)?::(exec|eval)(_file)?|PyImport_\w+|PyRun_\w+)\b", + "why": "run() stays in C++ and runs no Python: Triton, torch.compile, torch.utils.cpp_extension and other " + "kernel toolkits would come in that way"}, +} + + +def rule_table(extra=None): + """The rule table in force: the defaults, changed by a task's RULES (None drops a rule).""" + table = dict(RULES) + for name, rule in (extra or {}).items(): + if rule is None: + table.pop(name, None) + else: + table[name] = rule + return table + + +def _match(pattern, s): + return re.fullmatch(pattern, s) is not None + + +def _strip_cpp_comments(text): + """Removes // and /* */ comments (string literals are not parsed; good enough for these rules).""" + text = re.sub(r"/\*.*?\*/", " ", text, flags=re.S) + return re.sub(r"//[^\n]*", "", text) + + +def _includes(text): + return [(m.group(1), text.count("\n", 0, m.start()) + 1) + for m in re.finditer(r'^\s*#\s*include\s*[<"]([^>"]+)[>"]', text, flags=re.M)] + + +def check(solution_dir, rules, languages=None): + """Violations of `rules` by a solution's sources: a list of "file:line: what (rule: why)" strings. + + languages: the languages the solution may declare in solution.json (default LANGUAGES).""" + languages = LANGUAGES if languages is None else languages + out = [] + try: + with open(os.path.join(solution_dir, "solution.json")) as f: + spec = json.load(f) + except (OSError, ValueError) as e: + return [f"solution.json: unreadable ({e})"] + if not isinstance(spec, dict): + return ["solution.json: must hold an object"] + declared = spec.get("languages") + if not isinstance(declared, list) or not declared or any(lang not in languages for lang in declared): + out.append(f"solution.json: languages {declared} are not allowed here (allowed: {languages})") + options = spec.get("compile_options") if isinstance(spec.get("compile_options"), dict) else {} + flags = [f for v in options.values() if isinstance(v, list) for f in v] + for name, rule in rules.items(): + for flag in flags: + if "flag" in rule and _match(rule["flag"], str(flag)): + out.append(f"solution.json: compile option {flag!r} ({name}: {rule['why']})") + for root, dirs, files in os.walk(solution_dir): + dirs[:] = sorted(d for d in dirs if d != "__pycache__") + for file in sorted(files): + if not file.endswith(CPP_FILES): + continue + path = os.path.join(root, file) + try: + with open(path, encoding="utf-8") as f: + text = f.read() + except (OSError, UnicodeDecodeError): + continue + out += _check_cpp(solution_dir, os.path.relpath(path, solution_dir), text, rules) + return out + + +def _check_cpp(solution_dir, rel, text, rules): + out = [] + code = _strip_cpp_comments(text) + includes = _includes(code) + for name, rule in rules.items(): + for target, line in includes: + if "include" in rule and _match(rule["include"], target): + out.append(f"{rel}:{line}: #include <{target}> ({name}: {rule['why']})") + if "vendored" in rule and _match(rule["vendored"], target): + here = os.path.dirname(os.path.join(solution_dir, rel)) + if not any(os.path.isfile(os.path.join(d, target)) for d in (here, solution_dir)): + out.append(f"{rel}:{line}: #include <{target}> is not in the solution ({name}: {rule['why']})") + if "text" in rule: + for m in re.finditer(rule["text"], code): + line = code.count("\n", 0, m.start()) + 1 + out.append(f"{rel}:{line}: {m.group(0)!r} ({name}: {rule['why']})") + break + return out + + +# Run-time audit. + +# The library a kernel was launched from, by the path of the innermost library on the launch's native stack. +LAUNCHERS = [ + (r".*/libcublas(Lt)?[^/]*", "cuBLAS"), + (r".*/libcudnn[^/]*", "cuDNN"), + (r".*/lib(torch|c10)[^/]*", "ATen"), + (r".*triton.*", "Triton"), + (r".*/libnccl[^/]*", "NCCL"), + (r".*/lib(cufft|curand|cusparse|cusolver|cutensor|nvshmem|npp)[^/]*", "CUDA library"), +] +# Frames skipped on the way out: CUPTI and this module's callback, the Python interpreter, the CUDA runtime/driver. +SKIP = r"(.*/)?(libffi[^/]*|_ctypes[^/]*|python[^/]*|libpython[^/]*|callback_wrappers[^/]*|cupti[^/]*|libcupti[^/]*|" \ + r"libcuda\.so[^/]*|libcudart[^/]*|libc\.so[^/]*|ld-linux[^/]*|\?)" +# Fallback by kernel name, for launches that could not be traced. +NAMES = [ + (r"(void )?at::.*|.*at::native::.*|_ZN2at.*|.*at_cuda_detail.*|_ZN3c10.*", "ATen"), + (r".*(nvjet|cublas|gemv2?[NT]?_|gemmSN|splitKreduce|xmma|sgemm|hgemm|dgemm).*|cutlass3x_.*|cutlass_\d.*", "cuBLAS"), + (r".*cudnn.*", "cuDNN"), + (r"triton_.*", "Triton"), + (r"(void )?nccl.*|_Z\d+nccl.*", "NCCL"), +] +# Libraries whose kernels are communication: allowed outright, and counted apart from the placeholders. +COMMUNICATION = ("NCCL",) +# The CUPTI activity buffers: few and large, since a decode launches tens of thousands of kernels. +BUFFER_BYTES = 32 << 20 + +_libc = None +_demangler = None + + +class _DlInfo(ctypes.Structure): + _fields_ = [("fname", ctypes.c_char_p), ("fbase", ctypes.c_void_p), ("sname", ctypes.c_char_p), + ("saddr", ctypes.c_void_p)] + + +def _launcher(): + """Path of the innermost library on the current native stack that is not CUPTI, Python or the CUDA runtime.""" + global _libc + if _libc is None: + _libc = ctypes.CDLL("libc.so.6") + _libc.backtrace.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_int] + _libc.dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(_DlInfo)] + buf = (ctypes.c_void_p * 96)() + n = _libc.backtrace(buf, 96) + for i in range(n): + info = _DlInfo() + path = info.fname.decode() if buf[i] and _libc.dladdr(buf[i], ctypes.byref(info)) and info.fname else "?" + if not re.fullmatch(SKIP, path): + return path + return "" + + +def demangle(name): + global _demangler + if _demangler is None: + try: + lib = ctypes.CDLL("libstdc++.so.6") + lib.__cxa_demangle.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_int)] + lib.__cxa_demangle.restype = ctypes.c_char_p + _demangler = lib + except OSError: + _demangler = False + if not _demangler or not name.startswith("_Z"): + return name + status = ctypes.c_int() + out = _demangler.__cxa_demangle(name.encode(), None, None, ctypes.byref(status)) + return out.decode() if status.value == 0 and out else name + + +def short_name(name, limit=110): + """A kernel's demangled name without its return type, parameter list and the at::native/anonymous namespaces.""" + text = demangle(name).replace("(anonymous namespace)::", "") + depth, start, end = 0, 0, len(text) + for i, c in enumerate(text): + if c in "<(": + if c == "(" and depth == 0 and i > 0: + end = i + break + depth += 1 + elif c in ">)": + depth -= 1 + elif c == " " and depth == 0: + start = i + 1 + text = text[start:end] or text + for noise in ("at::native::", "at::", "std::"): + text = text.replace(noise, "") + return text if len(text) <= limit else text[:limit - 3] + "..." + + +def classify(name, launcher, traced): + """The library a kernel belongs to (None: the solution's own).""" + if traced: + for pattern, label in LAUNCHERS: + if re.fullmatch(pattern, launcher): + return label + return None + for pattern, label in NAMES: + if re.fullmatch(pattern, name) or re.fullmatch(pattern, demangle(name)): + return label + return None + + +Op = collections.namedtuple("Op", "kind name bytes launcher library ns") + + +def launch_callbacks(cupti): + """The driver API callbacks of kernel launches (cuLaunchKernel, cuLaunchKernelEx, cuGraphLaunch, ...).""" + return [int(getattr(cbid, "value", cbid)) for name, cbid in cupti.driver_api_trace_cbid.__members__.items() + if "Launch" in name and "HostFunc" not in name] + + +def trace(fn, *args): + """Calls fn(*args) once under CUPTI; returns its GPU operations in start order, as Op(kind, name, bytes, + launcher, library, ns): kind is kernel, memcpy or memset; library is None for the solution's own kernels; ns is + the operation's GPU time.""" + import torch + from cupti import cupti + + launches, records, errors = {}, [], [] + + def on_api(userdata, domain, cbid, data): + try: + if data.callback_site == cupti.ApiCallbackSite.API_ENTER: + graph = "Graph" in data.function_name + launches[data.correlation_id] = "" if graph else _launcher() + except Exception as e: # never let an exception cross into the driver + errors.append(repr(e)) + + def requested(): + return BUFFER_BYTES, 0 + + def completed(activities): + for a in activities: + if a.kind == cupti.ActivityKind.CONCURRENT_KERNEL: + records.append((a.start, a.end, "kernel", a.name, 0, a.correlation_id, + bool(a.graph_id or a.is_device_launched))) + elif a.kind == cupti.ActivityKind.MEMCPY: + records.append((a.start, a.end, "memcpy", "", a.bytes, a.correlation_id, False)) + elif a.kind == cupti.ActivityKind.MEMSET: + records.append((a.start, a.end, "memset", "", a.bytes, a.correlation_id, False)) + + kinds = (cupti.ActivityKind.CONCURRENT_KERNEL, cupti.ActivityKind.MEMCPY, cupti.ActivityKind.MEMSET) + torch.cuda.synchronize() + sub = cupti.subscribe(on_api, 0) + try: + for cbid in launch_callbacks(cupti): + cupti.enable_callback(1, sub, cupti.CallbackDomain.DRIVER_API, cbid) + for k in kinds: + cupti.activity_enable(k) + cupti.activity_register_callbacks(requested, completed) + fn(*args) + torch.cuda.synchronize() + cupti.activity_flush_all(1) + finally: + for k in kinds: + cupti.activity_disable(k) + cupti.unsubscribe(sub) + cupti.finalize() + if errors: + raise RuntimeError(f"the launch trace failed: {errors[0]}") + ops = [] + for start, end, kind, name, nbytes, corr, untraceable in sorted(records): + launcher = launches.get(corr) + traced = kind == "kernel" and not untraceable and bool(launcher) + library = classify(name, launcher or "", traced) if kind == "kernel" else None + ops.append(Op(kind, name, nbytes, os.path.basename(launcher or ""), library, end - start)) + return ops + + +def placeholder(library): + """Whether a library's kernels are placeholders (cuBLAS, cuDNN, ATen, ...), as opposed to the solution's own or + communication.""" + return bool(library) and library not in COMMUNICATION + + +def kernels(ops): + """What a call launched, by kernel: [[library or "", name, launches, ns]] by GPU time, the largest first; copies + and sets as the kernels "(memcpy)" and "(memset)".""" + table = collections.defaultdict(lambda: [0, 0]) + for op in ops: + key = (op.library or "", op.name if op.kind == "kernel" else f"({op.kind})") + table[key][0] += 1 + table[key][1] += op.ns + return sorted(([lib, name, n, ns] for (lib, name), (n, ns) in table.items()), key=lambda r: -r[3]) + + +def library_share(ops): + """The share of a call's GPU time (the summed time of its operations) spent in placeholder library kernels.""" + total = sum(op.ns for op in ops) + return sum(op.ns for op in ops if placeholder(op.library)) / total if total else 0.0 + + +def comm_share(ops): + """The share of a call's GPU time spent in communication kernels (NCCL).""" + total = sum(op.ns for op in ops) + return sum(op.ns for op in ops if op.library in COMMUNICATION) / total if total else 0.0 + + +def main(): + here = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, here) + import spec # noqa: E402 + + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("solution", nargs="?", default="solution") + a = p.parse_args() + table = rule_table(getattr(spec, "RULES", None)) + found = check(a.solution, table) + for v in found: + print(v) + print(f"stack rules: {'ok' if not found else f'{len(found)} violations'} ({len(table)} rules)") + sys.exit(1 if found else 0) + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/wan22_vae_decode/tools/bench.py b/fastvideo-kernel/wan22_vae_decode/tools/bench.py new file mode 100644 index 000000000..93826a539 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/tools/bench.py @@ -0,0 +1,771 @@ +"""Correctness, timing and library-kernel audit of a solution: the Wan2.2 VAE decode of latents into video, on one, two +or four GB200 GPUs of one node. The task's CI runs this same file (see ci.py). + +The task (spec.py has the details) +-------- +A solution decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T-3, 704, 1248] (bf16, in [-1, 1]) with +the Wan2.2 VAE's decoder, as diffusers' AutoencoderKLWan decodes them (reference.py in the kernel repo). It is a +directory (default solution/, see loader.py) compiled into one extension module that exports + + run(z, weights, rank, world, pg, out) positional, destination passing + prepare(weights, rank, world, pg) -> list[Tensor] optional + +called in one process per GPU: z the full latents on every rank; weights the decoder's 110 bf16 tensors on the rank's +GPU, in the order of spec.WEIGHTS (or what prepare returned); pg torch.distributed's NCCL ProcessGroup over the ranks; +out the video on rank 0, from torch.empty and never zeroed (every element must be written), and an empty tensor on the +other ranks. When run has returned on every rank, rank 0's out holds the video. prepare, if exported, runs once per +process before any timed call and is not timed; its result replaces the weights in every run. + +The workloads (spec.WORKLOADS) are clips decoded at a world size, named x. A clip is the first T latent +frames of one of three real 20 s clips, named by its video frames: f1, f5, f17 (T = 1, 2, 5) and f121, f241, f481 +(T = 31, 61, 121). The long clips at 1, 2 and 4 GPUs are the nine timed workloads (f121x1 ... f481x4), whose geomean +is the objective; the short clips at 1 and 4 GPUs (f1x1 ... f17x4) are checked and never timed. + +Running it +---------- +In a job with as many GPUs as the largest world size it runs, with the task's data at $DATA (default /data), from the +repo root: + + python tools/bench.py every workload this job has the GPUs for: check, audit, time + python tools/bench.py --workload f5x1 f17x1 a subset; the short clips take seconds + python tools/bench.py --workload f121x1 --no-time correctness only + python tools/bench.py --kernels also list the audited calls' kernels by GPU time + python tools/bench.py --solution other another solution directory + +A run +----- +The solution is built once, by this process, exactly as it is deployed (loader.py; a development copy, tools/, takes a +build its node already made of the same sources, flags and toolchain from the node's cache, unless --no-cache). Then, +per world size in turn (1, 2, 4), torch.multiprocessing spawns one process per rank, rank r on GPU r of the job's +CUDA_VISIBLE_DEVICES. Each initializes NCCL over 127.0.0.1 with its device (the default group, the solution's pg) and +a gloo group for the harness's own barriers, imports the build (in its working directory, as deployed), moves the +weights ($DATA/decoder_weights_bf16.pt, checked against spec.WEIGHTS) to its GPU and calls prepare, if the module +exports it. The ranks then run that world size's workloads in spec.WORKLOADS' order. A rank that raises, dies or +hangs (WORKLOAD_SEC per workload, SESSION_START_SEC before the first) fails the workload it was at: all the ranks' +processes are stopped, and the workloads after it run in new ones. This process runs no solution code. + +Every call gets z at a fresh address (a copy of the clip's latents) and out at a fresh address, 256-byte aligned and +filled with NaN on rank 0, so that an element run leaves unwritten fails the check. run must not write z, nor the +weights (compared after each workload with what prepare left), nor start threads: each fails the workload. + +Correctness +----------- +The first call of a workload, its warm-up, is checked on rank 0 against the clip's fp32 decode ($DATA/ref_.pt): +every element finite and within [-1, 1]; and with d = out - ref in fp32, max |d| <= min(2, 3 x bf16_max_abs) and +rms(d) <= 3 x bf16_rms, the error of the release's bf16 decode of that clip ($DATA/calibration.json; spec.py, +Correctness). The line shows max |d| and rms(d), each with its allowance and the share of it used, and the PSNR, +10 log10(4 / mse). Every later call's output (the audited call's, the trials') must agree with the checked one: +lose none of the elements it wrote (none left NaN or overflowed), and be within the same allowance of it everywhere +else. A race, or a state carried from call to call, fails the workload. + +Timing (the timed workloads) +------ +After the warm-up come the trials: 3 when the warm-up took under 1.5 s, else 2. A trial is one call timed by wall +clock between barriers on all ranks: synchronize the device, barrier, start, run, synchronize, barrier, stop. The +time of record is rank 0's; the latency of record the mean of the trials, and the fastest trial is shown too. While +the trials run, this process reads the SM clock of the world's GPUs through NVML every CLOCK_PERIOD_S (a rank inside +run may hold its GIL for the whole call, so the ranks do not read it themselves); the mean of rank 0's GPU is shown +with the latency, and the range over the world's GPUs where they differ by 2% or more: the clocks are not locked +alike on every GPU, and this says at what clock the times were taken. Which processes hold the world's GPUs is +asked (nvidia-smi) before every trial and after the last: a workload timed while another process held one of them is +marked SHARED, since its times do not compare with those of another job. + +Audit +----- +At world 1, one more call per timed clip runs under CUPTI, and each kernel is attributed to the library that launched +it (see audit.py). Library kernels (cuBLAS/cuBLASLt, cuDNN, ATen) are allowed as placeholders: their count per call +and their share of the call's GPU time show where the solution still depends on them. NCCL's kernels are +communication, reported apart and not counted as placeholders. --kernels lists the audited calls' kernels. + +Bound +----- +Each timed workload's latency is shown with its bound (bounds.py with spec.bounds: the decode's FLOPs spread evenly +over the world's GPUs at the 16-bit tensor-core peak, or the video and the weights at the memory bandwidth, whichever +is slower) and the gap, latency / bound. The peak assumes the SM clock GUILD_SM_CLOCK_MHZ where the environment sets +it (the guild's: 2062 MHz, the GB200's highest), else the GPU's highest. + +In CI, and from tools/ +---------------------- +ci.py runs this file in child processes: --save DIR has rank 0 save each workload's checked output there (node-local +scratch) instead of checking it, so that CI judges it in a process that runs no solution code; --no-check skips the +check (the timing runs), --no-audit the audit, and --json writes the report. The kernel repo holds copies of this +file, spec.py, loader.py, audit.py and bounds.py in tools/, which run as they do here: they read the data at $DATA +(default /data) and the solution from solution/ under the current directory. +""" + +import argparse +import collections +import ctypes +import functools +import json +import math +import os +import random +import socket +import statistics +import subprocess +import sys +import threading +import time +import traceback +from datetime import timedelta +from multiprocessing import connection + +# The caching allocator stays as torch configures it by default: without expandable segments, whose memory +# cudaIpcGetMemHandle cannot export, since a solution may share its buffers between ranks by CUDA IPC. +import torch +import torch.distributed as dist + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import audit # noqa: E402 +import bounds # noqa: E402 +import loader # noqa: E402 +import spec # noqa: E402 + +# Trials per timed workload, fewer when one call takes long. +TRIALS, LONG_TRIALS, LONG_CALL_S = 3, 2, 1.5 +# How often this process reads rank 0's SM clock while ranks run. +CLOCK_PERIOD_S = 0.02 +NVML_CLOCK_SM = 1 # the nvmlClockType_t of the SM clock +# Limits that stop a runaway: a session's start (NCCL, the module, the weights, prepare), and one workload (its calls, +# its check or its save); a call of the release's bf16 decode takes 15 s at the longest clip. +SESSION_START_SEC = 900 +WORKLOAD_SEC = 1800 +STOP_GRACE_SEC = 15 +# Elements of a video compared at once, in fp32 (1 GiB per operand). +CHUNK_ELEMENTS = 1 << 28 + + +def gpu_holders(): + """{GPU UUID (lower case, without "GPU-"): ids of the processes that hold it} now, as nvidia-smi gives them, or + None when it cannot say.""" + try: + got = subprocess.run(["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader"], + capture_output=True, text=True, timeout=30) + except (OSError, subprocess.TimeoutExpired): + return None + if got.returncode != 0: + return None + out = collections.defaultdict(set) + for line in got.stdout.splitlines(): + gpu, _, pid = line.partition(",") + if pid.strip().isdigit(): + out[gpu.strip().lower().removeprefix("gpu-")].add(int(pid)) + return out + + +def others(questions, uuids, ours): + """How many processes besides this run's held one of these GPUs at any of the questions (gpu_holders' answers), or + None when one could not be asked. This run's processes are known by their ids where nvidia-smi gives the ids this + process sees (one PID namespace, the usual case: then a rank that also holds another rank's GPU is still one of + ours); otherwise each GPU is taken to hold one of ours, the rank on it.""" + if not questions or any(q is None for q in questions): + return None + total = 0 + for uuid in uuids: + seen = set().union(*(q.get(uuid, set()) for q in questions)) + total += len(seen - ours) if seen & ours else max(0, len(seen) - 1) + return total + + +@functools.cache +def sm_clock(uuid): + """A function that reads the SM clock of the GPU with this UUID now, in MHz (None when a read fails), or None where + NVML cannot be loaded or does not know the GPU. NVML in this process takes microseconds per read.""" + try: + nvml = ctypes.CDLL("libnvidia-ml.so.1") + if nvml.nvmlInit_v2() != 0: + return None + handle, mhz = ctypes.c_void_p(), ctypes.c_uint() + if not uuid or nvml.nvmlDeviceGetHandleByUUID(f"GPU-{uuid}".encode(), ctypes.byref(handle)) != 0: + return None + except (OSError, AttributeError): + return None + + def read(): + ok = nvml.nvmlDeviceGetClockInfo(handle, NVML_CLOCK_SM, ctypes.byref(mhz)) == 0 + return mhz.value if ok and mhz.value else None + + return read if read() else None + + +class Clock: + """The SM clocks of the world's GPUs (rank order), read every CLOCK_PERIOD_S by a thread of this process, as + (monotonic time, [MHz or None per GPU]).""" + + def __init__(self, uuids): + self.reads, self.samples, self.done = [sm_clock(u) for u in uuids], [], threading.Event() + self.thread = threading.Thread(target=self._run, daemon=True) if any(self.reads) else None + if self.thread: + self.thread.start() + + def _run(self): + while not self.done.wait(CLOCK_PERIOD_S): + self.samples.append((time.monotonic(), [read() if read else None for read in self.reads])) + + def stop(self): + self.done.set() + if self.thread: + self.thread.join() + + def means(self, windows): + """Per GPU, the mean of its reads within these (start, stop) windows of time.monotonic(), or None without + any.""" + inside = [mhz for t, mhz in list(self.samples) if any(a <= t <= b for a, b in windows)] + out = [] + for i in range(len(self.reads)): + got = [m[i] for m in inside if m[i]] + out.append(statistics.mean(got) if got else None) + return out + + +class Pool: + """A workload's buffers for one tensor: each call's at a fresh address, shifted by multiples of 256 bytes.""" + + ALIGN, MAX_SHIFT = 256, 2048 + + def __init__(self, shape, dtype, device, calls, seed): + rng = random.Random(seed) + self.offsets = [0] + for _ in range(calls - 1): + self.offsets.append(self.offsets[-1] + rng.randint(1, self.MAX_SHIFT // self.ALIGN)) + self.shape, self.numel = tuple(shape), math.prod(shape) + self.block = max(1, self.ALIGN // torch.empty(0, dtype=dtype).element_size()) + self.buffer = torch.empty(self.numel + self.offsets[-1] * self.block, dtype=dtype, device=device) + self.calls = 0 + + def next(self): + start = self.offsets[self.calls] * self.block + self.calls += 1 + return self.buffer[start:start + self.numel].view(self.shape) + + +def fingerprint(tensors): + """Each tensor's bytes summed as integers: whether run() wrote the weights.""" + out = [] + for t in tensors: + raw = t.detach().contiguous().view(-1).view(torch.uint8) + try: + raw = raw.view(torch.int32) + except RuntimeError: + pass + out.append(int(raw.sum(dtype=torch.int64))) + return out + + +def judge(out, ref, allowed): + """A video against the reference, or a later call's video against the checked one's: {"passed", "nonfinite", + "outside" (elements beyond [-1, 1]), "max_abs", "rms", "psnr_db", "max_allowed", "rms_allowed"}, d = out - ref where + both are finite (rms and PSNR over those elements, None where there is none), in fp32 on out's GPU (the current one + for a video on the CPU), in chunks of frames, the squares summed in fp64.""" + device = out.device if out.is_cuda else torch.device("cuda", torch.cuda.current_device()) + step = max(1, CHUNK_ELEMENTS // max(1, out[:, :, :1].numel())) + nonfinite = outside = compared = unmatched = 0 + max_abs = squares = 0.0 + for f in range(0, out.shape[2], step): + x = out[:, :, f:f + step].to(device, torch.float32) + y = ref[:, :, f:f + step].to(device, torch.float32) + finite = torch.isfinite(x) + nonfinite += int((~finite).sum()) + outside += int((x.abs() > 1).sum()) + # Where the checked call itself left an element unfinite, a later call is not compared there (the check said). + both = finite & torch.isfinite(y) + compared += int(both.sum()) + unmatched += int((~finite & torch.isfinite(y)).sum()) + d = torch.where(both, x - y, 0.0) + max_abs = max(max_abs, float(d.abs().max())) + squares += float(d.square().sum(dtype=torch.float64)) + del x, y, finite, both, d + mse = squares / compared if compared else None + rms = math.sqrt(mse) if compared else None + passed = (nonfinite == 0 and outside == 0 and compared == out.numel() and max_abs <= allowed["max_abs"] + and rms <= allowed["rms"]) + return {"passed": passed, "nonfinite": nonfinite, "unmatched": unmatched, "outside": outside, + "max_abs": max_abs if compared else None, "rms": rms, "psnr_db": 10 * math.log10(4 / mse) if mse else None, + "max_allowed": allowed["max_abs"], "rms_allowed": allowed["rms"]} + + +def verdict(c): + """A check (judge's result) in a few words.""" + parts = [] + if c["nonfinite"]: + parts.append(f"{c['nonfinite']} elements not finite (NaN: never written, or overflowed)") + if c["outside"]: + parts.append(f"{c['outside']} elements outside [-1, 1]") + if c["max_abs"] is None: + parts.append("no element to compare") + else: + parts.append(f"max |d| {c['max_abs']:.4g} of {c['max_allowed']:.4g} (x{c['max_abs'] / c['max_allowed']:.2f})") + parts.append(f"rms {c['rms']:.4g} of {c['rms_allowed']:.4g} (x{c['rms'] / c['rms_allowed']:.2f})") + parts.append(f"PSNR {c['psnr_db']:.1f} dB" if c["psnr_db"] is not None else "identical") + return f"{'ok ' if c['passed'] else 'FAIL'} ({', '.join(parts)})" + + +def run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg): + """One workload on one rank, in step with the other ranks (see the module docstring); rank 0's result.""" + entry = spec.WORKLOADS[w] + clip, device = entry["clip"], torch.device("cuda", rank) + timed = entry["timed"] and cfg["time"] + audited = entry["timed"] and cfg["audit"] and world == 1 + calls = 1 + audited + (TRIALS if timed else 0) + seed = list(spec.WORKLOADS).index(w) + z0 = spec.make_inputs(w, device)[0] + zs = Pool(z0.shape, z0.dtype, device, calls, seed) + outs = Pool(spec.output_shape(w), torch.bfloat16, device, calls, seed + 1) if rank == 0 else None + allowed = spec.allowance(clip) + res = {"clip": clip, "world": world, "timed": entry["timed"]} + problems = [] + threads = threading.active_count() + + def arguments(): + z = zs.next() + z.copy_(z0) + out = outs.next().fill_(math.nan) if outs else torch.empty(0, dtype=torch.bfloat16, device=device) + return z, out + + def unchanged(z): + if not torch.equal(z, z0) and "run() wrote its input z" not in problems: + problems.append("run() wrote its input z") + + def call(): + z, out = arguments() + torch.cuda.synchronize(device) + dist.barrier(group=harness) + start = time.monotonic() + sol.fn(z, weights, rank, world, pg, out) + torch.cuda.synchronize(device) + dist.barrier(group=harness) + stop = time.monotonic() + unchanged(z) + return out, start, stop + + # The warm-up, which is the call checked. + out, start, stop = call() + warm = stop - start + first, worst = None, 0.0 + if rank == 0: + if cfg["save"]: + path = os.path.join(cfg["save"], f"{w}.pt") + torch.save(out.cpu(), path) + res["saved"] = path + elif cfg["check"]: + res["check"] = judge(out, spec.load_reference(clip), allowed) + if calls > 1: + first = out.clone() + + def agree(out, what): + nonlocal worst + # It agrees when it loses no element the checked call wrote, and is within the allowance of it elsewhere; + # the checked call's own faults are its check's to report. + c = judge(out, first, allowed) + worst = max(worst, c["max_abs"] or 0.0) + close = c["max_abs"] is None or (c["max_abs"] <= allowed["max_abs"] and c["rms"] <= allowed["rms"]) + if c["unmatched"] or not close: + lost = f"{c['unmatched']} elements not finite where it was; " if c["unmatched"] else "" + problems.append(f"{what}'s output does not agree with the checked call's: {lost}" + f"{verdict(c).removeprefix('FAIL').removeprefix('ok').strip()}") + + # World 1 alone: one call under CUPTI, which attributes every kernel. + if audited: + z, out = arguments() + ops = audit.trace(sol.fn, z, weights, rank, world, pg, out) + unchanged(z) + agree(out, "the audited call") + res |= {"kernels": audit.kernels(ops), "launches": sum(1 for op in ops if op.kind == "kernel"), + "lib": sum(1 for op in ops if audit.placeholder(op.library)), "lib_share": audit.library_share(ops), + "comm": sum(1 for op in ops if op.library in audit.COMMUNICATION), "comm_share": audit.comm_share(ops)} + del ops + + if timed: + n = [TRIALS if warm < LONG_CALL_S else LONG_TRIALS] + dist.broadcast_object_list(n, src=0, group=harness) + questions, trials, windows = [], [], [] + for k in range(n[0]): + if rank == 0: + questions.append(gpu_holders()) + out, start, stop = call() + trials.append(1e3 * (stop - start)) + windows.append((start, stop)) + if rank == 0: + agree(out, f"trial {k + 1}") + if rank == 0: + questions.append(gpu_holders()) + res |= {"trials": trials, "ms": statistics.mean(trials), "min_ms": min(trials), "windows": windows} + held = others(questions, [uuid for _, uuid in ranks], {pid for pid, _ in ranks} | {os.getppid()}) + if held: + res["shared"] = held + if first is not None: + res["agree_max_abs"] = worst + if fingerprint(weights) != marks: + problems.append("run() wrote its weights (what prepare returned is read-only to run: keep scratch elsewhere)") + if threading.active_count() > threads: + problems.append(f"the solution started threads ({threading.active_count() - threads} more)") + + # What every rank found, so that rank 0 reports it. + found = [None] * world + dist.all_gather_object(found, problems, group=harness) + text = [p if world == 1 else f"rank {r}: {p}" for r, ps in enumerate(found) for p in ps] + if text and rank == 0: + res["error"] = "; ".join(text) + return res + + +def worker(rank, world, port, cfg, conn): + """The process of one rank (see A run in the module docstring): it reports to the process that spawned it through + conn, ("ready", rank, uuid, prepare seconds), ("start", workload) and ("result", workload, result) on rank 0, + ("done", rank), or ("error", rank, workload or None, phase, traceback), after which it exits at once.""" + phase, current = "start", None + try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + timeout = timedelta(seconds=WORKLOAD_SEC) + dist.init_process_group("nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world, + device_id=device, timeout=timeout) + harness = dist.new_group(backend="gloo", timeout=timeout) + pg = dist.group.WORLD + phase = "load" + sol = loader.import_built(cfg["workdir"], cfg["spec"]) + phase = "weights" + weights = spec.load_weights(device) + phase = "prepare" + prepare_s = None + if sol.prepare is not None: + torch.cuda.synchronize(device) + start = time.monotonic() + got = sol.prepare(weights, rank, world, pg) + torch.cuda.synchronize(device) + prepare_s = time.monotonic() - start + if not isinstance(got, (list, tuple)) or not all(isinstance(t, torch.Tensor) for t in got): + raise TypeError(f"prepare returned {type(got).__name__}, not a list of tensors") + weights = list(got) + uuid = str(torch.cuda.get_device_properties(device).uuid).lower().removeprefix("gpu-") + ranks = [None] * world + dist.all_gather_object(ranks, (os.getpid(), uuid), group=harness) + conn.send(("ready", rank, uuid, prepare_s)) + phase = "workload" + marks = fingerprint(weights) + for w in cfg["workloads"]: + current = w + if rank == 0: + conn.send(("start", w)) + res = run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg) + if rank == 0: + conn.send(("result", w, res)) + current = None + torch.cuda.empty_cache() + phase = "end" + dist.barrier(group=harness) + conn.send(("done", rank)) + dist.destroy_process_group() + except BaseException: + try: + conn.send(("error", rank, current, phase, traceback.format_exc(limit=8))) + finally: + sys.stdout.flush() + sys.stderr.flush() + # At once: NCCL's teardown would wait for the other ranks, which may wait for this one. + os._exit(1) + + +def free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def stop(procs, grace=0.0): + """Ends rank processes: those that do not exit within `grace` seconds get SIGTERM, then SIGKILL after + STOP_GRACE_SEC.""" + deadline = time.monotonic() + grace + for p in procs: + p.join(max(0.0, deadline - time.monotonic())) + for p in procs: + if p.is_alive(): + p.terminate() + deadline = time.monotonic() + STOP_GRACE_SEC + for p in procs: + p.join(max(0.1, deadline - time.monotonic())) + if p.is_alive(): + p.kill() + p.join(10) + + +def exit_text(code): + if code is None: + return "did not exit" + if code < 0: + import signal + + try: + return f"was killed by {signal.Signals(-code).name}" + except ValueError: + return f"was killed by signal {-code}" + return f"exited with code {code}" + + +def session(world, workloads, cfg, report, show): + """One set of rank processes running these workloads at one world size, in order; results go into report as they + arrive, and show(w) prints each. Returns the workloads it did not reach: after a failure at one workload (which gets + the error), those after it; [] when every one ran, or when the failure came before any (every one gets it).""" + ctx = torch.multiprocessing.get_context("spawn") + port = free_port() + pipes = [ctx.Pipe(duplex=False) for _ in range(world)] + procs = [ctx.Process(target=worker, args=(r, world, port, dict(cfg, workloads=workloads), pipes[r][1]), + daemon=True) for r in range(world)] + for p in procs: + p.start() + for _, send in pipes: + send.close() + receivers = {pipes[r][0]: r for r in range(world)} + done, errors, finished, uuids = set(), [], [], {} + current, clock, failure = None, None, None + deadline = time.monotonic() + SESSION_START_SEC + started = time.monotonic() + + def handle(msg): + nonlocal current, clock, deadline + kind = msg[0] + if kind == "ready": + _, rank, uuid, prepare_s = msg + uuids[rank] = uuid + if rank == 0: + report.setdefault("prepare_s", {})[str(world)] = prepare_s + if len(uuids) == world: + clock = Clock([uuids[r] for r in range(world)]) + ready = f", prepare {report['prepare_s'][str(world)]:.1f} s" if prepare_s is not None else "" + print(f"world {world}: {world} rank{'s' if world > 1 else ''} up in {time.monotonic() - started:.0f} s" + f"{ready}", flush=True) + elif kind == "start": + current = msg[1] + deadline = time.monotonic() + WORKLOAD_SEC + elif kind == "result": + _, w, res = msg + windows = res.pop("windows", None) + means = clock.means(windows) if clock and windows else [] + if means and means[0]: + res["clock_mhz"] = means[0] + if world > 1 and means and all(means): + res["world_clock_mhz"] = [min(means), max(means)] + report["workloads"][w] = res + finished.append(w) + current = None + deadline = time.monotonic() + WORKLOAD_SEC + show(w) + elif kind == "error": + errors.append(msg[1:]) + elif kind == "done": + done.add(msg[1]) + + try: + while len(done) < world: + # A rank that exited cleanly may still have its last messages in its pipe: it fails once that is read. + bad = [(r, p.exitcode) for r, p in enumerate(procs) if r not in done and p.exitcode is not None + and (p.exitcode != 0 or pipes[r][0] not in receivers)] + if bad: + failure = "; ".join(f"rank {r} {exit_text(code)}" for r, code in bad) + break + left = deadline - time.monotonic() + if left <= 0: + failure = (f"timed out after {WORKLOAD_SEC} s" if current else + f"the ranks did not start within {SESSION_START_SEC} s") + break + waits = list(receivers) + [p.sentinel for p in procs if p.exitcode is None] + for obj in connection.wait(waits, timeout=min(left, 5.0)): + if obj in receivers: + try: + handle(obj.recv()) + except (EOFError, OSError): + del receivers[obj] + # What the other ranks still say, such as the error that made one exit. + drain = time.monotonic() + 2.0 + while failure and receivers and time.monotonic() < drain: + for obj in connection.wait(list(receivers), timeout=max(0.0, drain - time.monotonic())): + try: + handle(obj.recv()) + except (EOFError, OSError): + del receivers[obj] + finally: + stop(procs, 0.0 if failure else STOP_GRACE_SEC) + if clock: + clock.stop() + if not failure: + return [] + + # The first error a rank reported says what failed; else how a rank ended. + if errors: + rank, at, phase, text = errors[0] + if phase == "load": + report["load_error"] = text.strip().splitlines()[-1] + text = (f"rank {rank}: " if world > 1 else "") + text.strip() + else: + text = failure + rest = [w for w in workloads if w not in finished] + at = current or (errors and errors[0][1]) + if at not in rest: + at = rest[0] if rest and finished else None + if not at or report.get("load_error"): + for w in rest: + report["workloads"][w] = {"error": text, "world": world, "clip": spec.WORKLOADS[w]["clip"], + "timed": spec.WORKLOADS[w]["timed"]} + show(w) + return [] + report["workloads"][at] = {"error": text, "world": world, "clip": spec.WORKLOADS[at]["clip"], + "timed": spec.WORKLOADS[at]["timed"]} + show(at) + return rest[rest.index(at) + 1:] + + +def kernel_lines(res, limit=25): + """The audited call's kernels by GPU time, one line each.""" + rows = res.get("kernels", []) + total = sum(r[3] for r in rows) or 1 + lines = [f" {'[' + lib + ']' if lib else '':<9} {ns / 1e6:10.2f} ms {100 * ns / total:5.1f}% x{n:<6} " + f"{name if name.startswith('(') else audit.short_name(name, 110)}" for lib, name, n, ns in rows[:limit]] + if len(rows) > limit: + rest = rows[limit:] + lines.append(f" {'':<9} {sum(r[3] for r in rest) / 1e6:10.2f} ms " + f"{100 * sum(r[3] for r in rest) / total:5.1f}% {len(rest)} more kinds") + return lines + + +def main(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--solution", default="solution", help="the solution directory (default: solution)") + p.add_argument("--workload", nargs="*", help="workloads to run (default: every one this job has the GPUs for)") + p.add_argument("--no-time", action="store_true", help="correctness only: no audit, no trials") + p.add_argument("--no-check", action="store_true", help="do not check the outputs (CI's timing runs)") + p.add_argument("--no-audit", action="store_true", help="no call under CUPTI") + p.add_argument("--kernels", action="store_true", help="list the audited calls' kernels by GPU time") + p.add_argument("--build-dir", help="write and build the solution here; reused when it holds the same build") + p.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds") + p.add_argument("--save", help="CI: save each workload's checked output here instead of checking it") + p.add_argument("--json", help="write the results to this file") + a = p.parse_args() + unknown = [w for w in a.workload or [] if w not in spec.WORKLOADS] + if unknown: + p.error(f"unknown workloads {unknown}; known: {list(spec.WORKLOADS)}") + # Absolute paths, the data's too: the ranks run in the solution's working directory. + solution = os.path.abspath(a.solution) + if os.path.isfile(solution): + solution = os.path.dirname(solution) + a.build_dir, a.save, a.json = (x and os.path.abspath(x) for x in (a.build_dir, a.save, a.json)) + os.environ["DATA"] = os.path.abspath(spec.data_dir()) + report = {"workloads": {}} + + def finish(): + if a.json: + with open(a.json, "w") as f: + json.dump(report, f) + + rules = audit.rule_table(getattr(spec, "RULES", None)) + report["rules"] = audit.check(solution, rules) + for v in report["rules"]: + print(f"rule violated: {v}") + print(f"stack rules: {'ok' if not report['rules'] else 'FAIL'} ({len(rules)} rules)") + + # The workloads: those named, or every one this job has the GPUs for. + gpus = torch.cuda.device_count() + chosen = a.workload or [w for w, e in spec.WORKLOADS.items() if e["world"] <= gpus] + skipped = [w for w in spec.WORKLOADS if w not in chosen and not a.workload] + if skipped: + print(f"skipped {', '.join(skipped)}: this job has {gpus} GPU{'s' if gpus != 1 else ''}") + order = [w for w in spec.WORKLOADS if w in chosen] + timed = [w for w in order if spec.WORKLOADS[w]["timed"] and not a.no_time] + limits, dev, peak = bounds.workload_bounds(spec, timed) + if limits: + print(bounds.describe(dev, peak)) + + try: + workdir, build_spec, seconds = loader.stage(solution, a.build_dir, + cache=None if a.no_cache else loader.cache_dir()) + except Exception as e: + report["load_error"] = str(e) if isinstance(e, loader.SolutionError) else traceback.format_exc(limit=6) + print(f"FAILED to build {solution}:\n{report['load_error']}", flush=True) + finish() + sys.exit(1) + report |= {"language": build_spec["languages"][0], "build_s": seconds} + built = f"built in {seconds:.1f} s" if seconds is not None else "a build reused" + print(f"{build_spec['entry_point']} ({', '.join(build_spec['languages'])}): {built}", flush=True) + + width = max(len(w) for w in spec.WORKLOADS) + + def show(w): + res = report["workloads"][w] + line = f"{w:<{width}}" + if "ms" in res: + at = f", SM clock {res['clock_mhz']:.0f} MHz" if "clock_mhz" in res else "" + low, high = res.get("world_clock_mhz", (None, None)) + if low and "clock_mhz" in res and high - low >= 0.02 * res["clock_mhz"]: + at += f", the world's GPUs {low:.0f}-{high:.0f} MHz" + line += f" {res['ms']:9.1f} ms (fastest trial {res['min_ms']:.1f}{at})" + if w in limits: + res |= {"bound_ms": limits[w], "gap": res["ms"] / limits[w]} + line += f" bound {limits[w]:.1f} ms, gap x{res['gap']:.2f}" + if "check" in res: + line += " " + verdict(res["check"]) + elif "saved" in res: + line += " saved" + if not res.get("timed", True): + line += " (correctness only)" + if "lib" in res: + line += f" library kernels {res['lib']} of {res['launches']} ({100 * res['lib_share']:.0f}% of GPU time)" + if res.get("comm"): + line += f", NCCL {res['comm']} ({100 * res['comm_share']:.0f}%)" + if "shared" in res: + line += f" SHARED: {res['shared']} other process(es) held a GPU while it was timed" + if "error" in res: + line += " ERROR\n" + res["error"] + print(line, flush=True) + if a.kernels and "kernels" in res: + print("\n".join(kernel_lines(res)), flush=True) + finish() + + cfg = {"workdir": workdir, "spec": build_spec, "time": not a.no_time, "check": not a.no_check, + "audit": not a.no_audit and (not a.no_time or a.kernels), "save": a.save} + for world in sorted({spec.WORKLOADS[w]["world"] for w in order}): + todo = [w for w in order if spec.WORKLOADS[w]["world"] == world] + if world > gpus or report.get("load_error"): + why = report.get("load_error") or f"needs {world} GPUs; this job has {gpus}" + for w in todo: + report["workloads"][w] = {"error": why, "world": world, "clip": spec.WORKLOADS[w]["clip"], + "timed": spec.WORKLOADS[w]["timed"]} + show(w) + continue + while todo: + todo = session(world, todo, cfg, report, show) + + results = [report["workloads"][w] for w in order if w in report["workloads"]] + times = [report["workloads"].get(w, {}).get("ms") for w in timed] + if timed and all(times): + gaps = [report["workloads"][w]["gap"] for w in timed if "gap" in report["workloads"][w]] + gap = f", gap to the bound x{bounds.geomean(gaps):.2f}" if len(gaps) == len(timed) else "" + clocks = [report["workloads"][w].get("clock_mhz") for w in timed] + at = f", mean SM clock {statistics.mean(clocks):.0f} MHz" if all(clocks) else "" + print(f"geomean {bounds.geomean(times):.1f} ms over {len(timed)} timed workloads (mean of trials){gap}{at}") + shared = [w for w in order if "shared" in report["workloads"].get(w, {})] + if shared: + print(f"SHARED GPU: another process held one while {', '.join(shared)} " + f"{'was' if len(shared) == 1 else 'were'} timed, so those times do not compare with another job's") + checked = [r for r in results if "check" in r] + if checked: + ok = sum(1 for r in checked if r["check"]["passed"] and "error" not in r) + print(f"correctness: ok at {ok} of {len(checked)} workloads checked") + failed = [w for w in order if "error" in report["workloads"].get(w, {})] + if failed: + print(f"errors at {', '.join(failed)}") + audited = [r for r in results if "kernels" in r] + if audited: + libs = collections.Counter() + for r in audited: + for lib, _, n, _ in r["kernels"]: + if audit.placeholder(lib): + libs[lib] += n + detail = ", ".join(f"{k} {v}" for k, v in libs.most_common()) + share = statistics.mean(r["lib_share"] for r in audited) + print(f"library kernels: {sum(libs.values())} launches in {len(audited)} audited calls" + f"{' (' + detail + ')' if detail else ''}, {100 * share:.0f}% of GPU time on average") + finish() + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/wan22_vae_decode/tools/bounds.py b/fastvideo-kernel/wan22_vae_decode/tools/bounds.py new file mode 100644 index 000000000..caad07537 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/tools/bounds.py @@ -0,0 +1,156 @@ +"""The physical lower bound of a call's time, from the work it must do and what the GPUs can do per second. + +A task's spec.py may define + + bounds(workload) -> {"flops": float, "bytes": float, "gpus": int} optionally also "math": one of PEAKS ("f16") + +lower bounds for one call of run() at that workload: the floating-point operations any correct algorithm performs +(its convolutions and matrix products) and the bytes the busiest GPU must move to or from its memory at least once, +and the GPUs the call runs on (default 1). The bound time is then + + bound = max(flops / (gpus x peak FLOP/s of the chosen math), bytes / memory bandwidth) + +with the peaks of the GPU it runs on at its SM clock: a call cannot be faster than the slower of the two, even with +its FLOPs spread evenly over its GPUs and nothing else to do, so the gap, measured / bound, says how much of the way +remains. The peaks are estimates, derived as SM count x SM clock x per-SM rate per clock (DEVICES), not measured; they +are the same for every solution, so gaps compare. The SM clock is the one GUILD_SM_CLOCK_MHZ names, where a task's +environment sets it, else the one DEVICES gives the GPU, else the GPU's highest SM clock: GB200 clocks are not locked, +and read low while they idle. + + python tools/bounds.py [--spec spec.py] the bound of every workload on this GPU (or the reference device) +""" + +import argparse +import math +import os +import subprocess + +# Per compute capability: dense per-SM throughput per clock, and the memory bandwidth. Estimates from the +# architecture's published figures; the memory clock does not follow the SM clock. `sm_clock_mhz`, where given, is the +# clock to assume instead of the GPU's. +DEVICES = { + (10, 0): { + "name": "Blackwell, compute capability 10.0", + # FLOP per SM per clock, dense (no structured sparsity), fp32 accumulation. + "per_sm_clk": { + "f16": 8192, # FP16/BF16 tensor cores (tcgen05.mma kind::f16) + "tf32": 4096, # TF32 tensor cores (kind::tf32) + "f32": 256, # FP32 FMA on the CUDA cores: 128 lanes x 2 + }, + "bytes_per_s": 8.0e12, # HBM3e, about 8 TB/s per GPU + }, +} +PEAKS = {"f16": "16-bit tensor cores", "tf32": "TF32 tensor cores", "f32": "FP32 CUDA cores"} +# Where no GPU is visible (on a laptop, say): a GB200 GPU at its highest SM clock. +REFERENCE_DEVICE = {"capability": (10, 0), "sms": 152, "clock_mhz": 2062} + + +def sm_clock_mhz(uuid=None): + """The highest SM clock of the GPU with this UUID, or the reference device's.""" + try: + cmd = ["nvidia-smi", "--query-gpu=clocks.max.sm", "--format=csv,noheader,nounits"] + if uuid: + cmd.insert(1, f"--id={uuid}") + out = subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout.split() + return float(out[0]) if out else REFERENCE_DEVICE["clock_mhz"] + except (OSError, ValueError, subprocess.SubprocessError): + return REFERENCE_DEVICE["clock_mhz"] + + +def assumed_clock_mhz(): + """The SM clock a task's environment says the bound assumes (GUILD_SM_CLOCK_MHZ), or None where it says none.""" + try: + value = float(os.environ.get("GUILD_SM_CLOCK_MHZ", "")) + except ValueError: + return None + return value if value > 0 else None + + +def device(): + """{"capability", "sms", "clock_mhz"} of CUDA device 0, or REFERENCE_DEVICE where there is none. It asks the GPU's + properties only, which makes no CUDA context: a process that times others must hold none.""" + try: + import torch + + if torch.cuda.is_available(): + p = torch.cuda.get_device_properties(0) + uuid = str(getattr(p, "uuid", "") or "") + uuid = uuid if not uuid or uuid.startswith(("GPU-", "MIG-")) else f"GPU-{uuid}" + capability = (p.major, p.minor) + clock = assumed_clock_mhz() or DEVICES.get(capability, {}).get("sm_clock_mhz") or sm_clock_mhz(uuid) + return {"capability": capability, "sms": p.multi_processor_count, "clock_mhz": clock} + except ImportError: + pass + return dict(REFERENCE_DEVICE) + + +def peaks(dev): + """{"f16": FLOP/s, "tf32": ..., "f32": ..., "bytes_per_s": B/s} of one GPU, or None when DEVICES lacks it.""" + table = DEVICES.get(tuple(dev["capability"])) + if table is None: + return None + hz = dev["clock_mhz"] * 1e6 + out = {k: v * dev["sms"] * hz for k, v in table["per_sm_clk"].items()} + out["bytes_per_s"] = table["bytes_per_s"] + return out + + +def bound_ms(work, peak): + """The bound time in ms of one call doing `work` ({"flops", "bytes"[, "gpus", "math"]}) on GPUs with these peaks.""" + math_kind = work.get("math", "f16") + if math_kind not in PEAKS: + raise ValueError(f"bounds(): math must be one of {sorted(PEAKS)}, not {math_kind!r}") + flops, nbytes = float(work.get("flops", 0.0)), float(work.get("bytes", 0.0)) + gpus = int(work.get("gpus", 1)) + return 1e3 * max(flops / (gpus * peak[math_kind]), nbytes / peak["bytes_per_s"]) + + +def describe(dev, peak): + """One line on how the bound is computed here.""" + return (f"bound = max(FLOPs / (GPUs x {peak['f16'] / 1e15:.2f} PFLOP/s) ({PEAKS['f16']}, dense), bytes / " + f"{peak['bytes_per_s'] / 1e12:.1f} TB/s), estimated for {dev['sms']} SMs per GPU at " + f"{dev['clock_mhz']:.0f} MHz") + + +def workload_bounds(spec, workloads, dev=None): + """{workload: bound ms} for a spec that defines bounds(), with the device and its peaks; ({}, dev, None) when the + spec has no bounds() or the device no peak figures.""" + dev = dev or device() + peak = peaks(dev) + fn = getattr(spec, "bounds", None) + if fn is None or peak is None: + return {}, dev, peak + return {w: bound_ms(fn(w), peak) for w in workloads}, dev, peak + + +def geomean(xs): + xs = list(xs) + return math.exp(sum(math.log(x) for x in xs) / len(xs)) if xs else math.nan + + +def main(): + import importlib.util + + here = os.path.dirname(os.path.abspath(__file__)) + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--spec", default=os.path.join(os.environ.get("CHECKS", here), "spec.py")) + a = p.parse_args() + module_spec = importlib.util.spec_from_file_location("spec", a.spec) + spec = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(spec) + workloads = [w for w, e in spec.WORKLOADS.items() if e.get("timed", True)] + bounds, dev, peak = workload_bounds(spec, workloads) + if peak is None: + raise SystemExit(f"no peak figures for compute capability {dev['capability']}") + if not bounds: + raise SystemExit(f"{a.spec} defines no bounds(workload)") + print(describe(dev, peak)) + for w, ms in bounds.items(): + work = spec.bounds(w) + print(f"{w:<10} {ms:10.1f} ms {work['flops'] / 1e15:8.4f} PFLOP over {work.get('gpus', 1)} GPU(s) " + f"{work['bytes'] / 1e9:6.2f} GB") + print(f"geomean {geomean(bounds.values()):.1f} ms over {len(bounds)} workloads") + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/wan22_vae_decode/tools/loader.py b/fastvideo-kernel/wan22_vae_decode/tools/loader.py new file mode 100644 index 000000000..3c33506d1 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/tools/loader.py @@ -0,0 +1,542 @@ +"""Builds and loads a solution exactly as it is deployed, so that what runs here runs there. + +A solution is a directory holding exactly the files that are deployed, plus `solution.json`, the deployment's build +spec, with its field names and meaning: + + {"languages": ["cuda_cpp"], "entry_point": "binding.cpp::run", "destination_passing_style": true, + "compile_options": {"cflags": [...], "cuda_cflags": [...], "ld_flags": [...]}, "dependencies": []} + +- `languages`: always ["cuda_cpp"]. +- `entry_point`: "::", the function the deployment calls, in a C++/CUDA file. +- `destination_passing_style`: required, and true for this task: run() writes its output into `out`, its last + argument, and returns nothing. The calling convention must never be left to inference. +- `compile_options` (optional): extra flags; see the build below. `dependencies` (optional): a list of names, which + does not change the build. `binding` (optional): null or "torch". +Any other field is refused: the deployment would ignore it, so a misspelled field would silently fall back to its +default. Every other file in the directory is a source: non-empty UTF-8 text at a relative path, a plain file (no +links; `__pycache__` and `.pyc` files are ignored). + +The module exports the entry, called positionally as the deployment calls it (argument names play no part): + + run(z, weights, rank, world, pg, out) every call; writes out on rank 0 + prepare(weights, rank, world, pg) -> list[Tensor] optional, under that name: once per process, before run + +(spec.py says what each argument is.) Load one solution per process. + +stage(solution_dir) writes the sources into a fresh working directory and compiles every .cpp/.cc/.cxx/.c/.cu file at +its top into one extension module, in a separate process, with torch.utils.cpp_extension.load, build_directory = the +working directory, include paths [working directory, $CUTLASS_DIR/include, $CUTLASS_DIR/tools/util/include] ($CUTLASS_DIR +defaults to fastvideo-kernel/include/cutlass, the CUTLASS submodule), and these +flags on top of torch's own: + - no compile_options: cuda_cflags [], cflags [], ld_flags ["-lcuda"]; + - compile_options given: its lists, where a missing one defaults to cflags [], cuda_cflags ["-O3", + "--use_fast_math"], ld_flags ["-lcuda"]; + - unless a cuda_cflag names an architecture (-gencode, -arch), a -gencode flag for the GPU's own architecture with + its arch-specific features goes first (sm_100a on compute capability 10.0), as the deployment does for its GPU. +Files in subdirectories are not compiled, but can be included. import_built(working directory) then makes it the +current directory, as the deployment does, imports the module and takes run (and prepare) from it; load() does both. +The module's name is the build's, not the solution's: PYBIND11_MODULE must be given TORCH_EXTENSION_NAME, as the +deployment builds the module under a name of its own. bench.py stages a build once and its ranks import that build. + +Development builds are kept on each node, so that a job does not compile again what an earlier job on its node built: +a content-addressed cache in a directory of the node that setup/node.sh makes for the agent user (CACHE_DIR, or +$GUILD_BUILD_CACHE), keyed by everything that determines the build: the sources, solution.json, the flags, the target +architecture, this loader, and the toolchain (the machine, Python, torch with its CUDA, nvcc, the host compiler, and +the variables they read). A build found there is copied into the working directory in place of compiling it; one made +is put there with its sources and module, without what the compiler left beside them. Only the commands of a +development copy of the harness use it (tools/ in a repo: `loader.py check` and `build`, `bench.py`), unless they are +given --no-cache. The copy in the task's checks/, which CI, the hooks and the task's commands run, never reads or +writes it and builds fresh: the jobs of agents write that directory, so what it holds is theirs to forge. Past a +quarter of its filesystem, a limit that only stops a runaway, the least recently used builds are removed. + + python tools/loader.py check [solution] [--build-dir DIR] build and load it, print what was loaded + python tools/loader.py build [solution] [DIR] only write and build it in DIR (default build/): with + --arch (e.g. sm_100a), on a machine without a GPU, + such as a CPU builder, where the module is not loaded +""" + +import argparse +import atexit +import functools +import hashlib +import importlib.util +import json +import os +import platform +import re +import shutil +import stat +import subprocess +import sys +import tempfile +import time + +SPEC_FILE = "solution.json" +# The name the build gives the module; solutions never name it themselves (see the module docstring). It is this +# kernel's own: another kernel's loader may build an extension of a generic name (solution_ext), and in one process the +# second extension of a name loads as the first one. +MODULE = "wan22_vae_decode_ext" +# The optional function a solution exports beside its entry, called once per process before any call of the entry. +PREPARE = "prepare" +BINDINGS = ("torch",) +CPP_ENTRY_SUFFIXES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh") +# The fields solution.json may hold. +SPEC_FIELDS = ("languages", "entry_point", "dependencies", "destination_passing_style", "binding", "compile_options") +OPTION_FIELDS = ("cflags", "cuda_cflags", "ld_flags") +# The languages a solution may declare, and the files compiled. +LANGUAGES = ("cuda_cpp",) +CPP_SOURCES = (".cu", ".cpp", ".cc", ".cxx", ".c") +DEFAULT_OPTIONS = {"cflags": [], "cuda_cflags": ["-O3", "--use_fast_math"], "ld_flags": ["-lcuda"]} +BUILD_TIMEOUT_SEC = 1800 +# The first argument of a PYBIND11_MODULE in a source. +MODULE_NAME = re.compile(r"\bPYBIND11_MODULE\s*\(\s*([A-Za-z_]\w*)") +HERE = os.path.dirname(os.path.abspath(__file__)) +# The node's cache of development builds (see the module docstring), and what each build there holds beside its sources. +CACHE_DIR = "/scratch/kdev-builds" +STAMP, SIZE = ".stamp", ".size" +# The share of its filesystem past which the cache removes its least recently used builds: a limit that stops a runaway. +CACHE_SHARE = 0.25 +# What a job that stopped while it stored a build left in the cache is removed after this long. +STALE_SEC = 3600 +# The variables the compilers read, which change what they make. +BUILD_VARIABLES = ("CC", "CXX", "CUDA_HOME", "CUDA_PATH", "CUTLASS_DIR", "NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", + "CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH", "LIBRARY_PATH") + + +class SolutionError(Exception): + """The solution is malformed, or fails to build or load.""" + + +def _nonempty_str(x): + return isinstance(x, str) and len(x) >= 1 + + +def _str_list(x): + return isinstance(x, list) and all(isinstance(s, str) for s in x) + + +def spec_problems(spec): + """Problems of a build spec, as a list of strings (empty: valid); see the module docstring. + + Stricter than the deployment only where it would coerce a value (a "true" string as a boolean) or fall back to a + default (an unknown field, a missing destination_passing_style).""" + if not isinstance(spec, dict): + return [f"{SPEC_FILE} must hold an object"] + out = [] + unknown = sorted(set(spec) - set(SPEC_FIELDS)) + if unknown: + out.append(f"unknown fields {unknown} (known: {list(SPEC_FIELDS)})") + langs = spec.get("languages") + if not isinstance(langs, list) or not langs or any(lang not in LANGUAGES for lang in langs): + out.append(f"languages: required, {list(LANGUAGES)}") + entry = spec.get("entry_point") + if not _nonempty_str(entry): + out.append("entry_point: required, a non-empty string") + elif entry.count("::") != 1: + out.append(f'entry_point: {entry!r} is not "::"') + elif os.path.splitext(entry.split("::")[0])[1] not in CPP_ENTRY_SUFFIXES: + out.append(f"entry_point: needs a C++/CUDA entry file, not {entry.split('::')[0]!r}") + deps = spec.get("dependencies", []) + if not isinstance(deps, list) or not all(_nonempty_str(d) for d in deps): + out.append("dependencies: a list of non-empty strings") + if "destination_passing_style" not in spec: + out.append("destination_passing_style: must be given, and true here (run writes out, its last argument)") + elif spec["destination_passing_style"] is not True: + out.append("destination_passing_style: true: this task's run() writes the video into out, its last argument, " + "and returns nothing") + if spec.get("binding") not in (None,) + BINDINGS: + out.append(f"binding: null or one of {list(BINDINGS)}") + options = spec.get("compile_options") + if options is not None: + if not isinstance(options, dict): + out.append("compile_options: null or an object") + else: + for k in OPTION_FIELDS: + if k in options and not _str_list(options[k]): + out.append(f"compile_options.{k}: a list of strings") + return out + + +def read_spec(solution_dir): + """solution.json, checked against the build spec's rules (see the module docstring).""" + path = os.path.join(solution_dir, SPEC_FILE) + try: + with open(path) as f: + spec = json.load(f) + except (OSError, ValueError) as e: + raise SolutionError(f"cannot read {path}: {e}") from e + problems = spec_problems(spec) + if problems: + raise SolutionError(f"{SPEC_FILE}: " + "; ".join(problems)) + return spec + + +def sources(solution_dir): + """{relative path: text} of every file of the solution except solution.json.""" + out = {} + for root, dirs, files in os.walk(solution_dir): + dirs[:] = sorted(d for d in dirs if d != "__pycache__") + for name in sorted(files): + path = os.path.join(root, name) + rel = os.path.relpath(path, solution_dir).replace(os.sep, "/") + if rel == SPEC_FILE or name.endswith(".pyc"): + continue + if os.path.islink(path): + raise SolutionError(f"{rel}: sources must be plain files, not links") + try: + with open(path, encoding="utf-8") as f: + text = f.read() + except (OSError, UnicodeDecodeError) as e: + raise SolutionError(f"{rel}: sources must be UTF-8 text ({e})") from e + if not text: + raise SolutionError(f"{rel}: sources must not be empty") + out[rel] = text + if not out: + raise SolutionError(f"{solution_dir} holds no sources") + named = sorted({m for text in out.values() for m in MODULE_NAME.findall(text) if m != "TORCH_EXTENSION_NAME"}) + if named: + raise SolutionError(f"PYBIND11_MODULE names its module {', '.join(named)}: give it TORCH_EXTENSION_NAME, the " + f"name the build chooses (the deployment builds the module under a name of its own)") + return out + + +def read(solution_dir): + """(spec, {relative path: text}) of a solution, checked as stage() checks it before building: solution.json against + the build spec's rules, the sources as the deployment takes them, the entry file among them.""" + spec = read_spec(solution_dir) + files = sources(solution_dir) + entry = spec["entry_point"].split("::")[0] + if entry not in files: + raise SolutionError(f"entry file {entry} is not among the sources") + return spec, files + + +def gpu_arch(): + """The architecture solutions are built for, e.g. sm_100a: the GPU's own, with its arch-specific features. It asks + the GPU's properties only, which makes no CUDA context.""" + import torch + + if not torch.cuda.is_available(): + raise SolutionError("there is no GPU here to take the architecture from: name it (--arch sm_100a, or the " + "architecture of the GPUs the build is for)") + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" + ("a" if major >= 9 else "") + + +def compile_flags(spec, arch): + """(cflags, cuda_cflags, ld_flags) added to torch's own flags (see the module docstring).""" + given = spec.get("compile_options") + options = {"cflags": [], "cuda_cflags": [], "ld_flags": ["-lcuda"]} if given is None else { + k: list(given.get(k, v)) for k, v in DEFAULT_OPTIONS.items()} + if not any("-gencode" in f or "-arch" in f for f in options["cuda_cflags"]): + options["cuda_cflags"].insert(0, f"-gencode=arch=compute_{arch[3:]},code={arch}") + return options["cflags"], options["cuda_cflags"], options["ld_flags"] + + +def cache_dir(): + """The node's cache of development builds, or None where a build is made fresh: in the copy of the harness in the + task's checks/ (CI, hooks and the task's commands run it, which must not load what a job could have put there), and + on a node without a cache that this user may write.""" + checks = os.environ.get("CHECKS") + if checks and os.path.realpath(checks) == os.path.realpath(HERE): + return None + cache = os.environ.get("GUILD_BUILD_CACHE", CACHE_DIR) + return cache if os.path.isdir(cache) and os.access(cache, os.W_OK | os.X_OK) else None + + +@functools.lru_cache(maxsize=None) +def toolchain(): + """What determines a build beside the solution, its flags and the architecture: this loader, the machine, Python, + torch with its CUDA, nvcc, the host compiler, and the variables the compilers read.""" + import torch + from torch.utils import cpp_extension + + def version(cmd): + try: + p = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + return (p.stdout + p.stderr).strip() + except (OSError, subprocess.TimeoutExpired) as e: + return f"unavailable: {e}" + + with open(os.path.abspath(__file__), "rb") as f: + loader = hashlib.sha256(f.read()).hexdigest() + return {"loader": loader, "machine": platform.machine(), "python": sys.version, + "torch": [torch.__version__, torch.version.cuda, os.path.dirname(torch.__file__)], + "nvcc": version([os.path.join(cpp_extension.CUDA_HOME or "/usr/local/cuda", "bin", "nvcc"), "--version"]), + "cxx": version([os.environ.get("CXX", "c++"), "--version"]), + "variables": {k: os.environ[k] for k in BUILD_VARIABLES if k in os.environ}} + + +def cache_key(spec, files, flags, arch, tools): + """The key of a build in the cache: a hash of everything that determines it (see the module docstring).""" + what = {"spec": spec, "files": sorted(files.items()), "flags": flags, "arch": arch, "toolchain": tools} + return hashlib.sha256(json.dumps(what, sort_keys=True).encode()).hexdigest() + + +def disk_usage(path): + """The bytes a directory takes on disk.""" + total = 0 + for root, _, names in os.walk(path): + for name in names: + try: + st = os.lstat(os.path.join(root, name)) + except OSError: + continue + total += getattr(st, "st_blocks", 0) * 512 or st.st_size + return total + + +def fetch(cache, key, build_dir): + """Copies the build of a key from the cache into build_dir, in its place, and marks the build used; False when the + cache has none, or it went while it was copied.""" + entry = os.path.join(cache, key) + if not os.path.exists(os.path.join(entry, MODULE + ".so")): + return False + shutil.rmtree(build_dir, ignore_errors=True) + try: + shutil.copytree(entry, build_dir, symlinks=True, ignore=lambda d, _: [SIZE] if d == entry else []) + os.utime(entry) + except OSError: + shutil.rmtree(build_dir, ignore_errors=True) + return False + return True + + +def store(cache, key, build_dir, files): + """Puts the build in build_dir into the cache under its key: its sources, module and stamp, without what the + compiler left beside them. A build of the key that another job put there first stays. Then the cache is kept within + its limit. A cache that cannot take it (a full disk, say) is left as it is: the build stands without it.""" + entry = os.path.join(cache, key) + if os.path.exists(entry): + return + try: + tmp = tempfile.mkdtemp(prefix=f".{key[:16]}-", dir=cache) + except OSError: + return + try: + for rel in [*files, MODULE + ".so", STAMP]: + dest = os.path.join(tmp, rel) + os.makedirs(os.path.dirname(dest), exist_ok=True) + shutil.copy2(os.path.join(build_dir, rel), dest) + with open(os.path.join(tmp, SIZE), "w") as f: + f.write(str(disk_usage(tmp))) + os.rename(tmp, entry) + except OSError: + shutil.rmtree(tmp, ignore_errors=True) + return + try: + st = os.statvfs(cache) + evict(cache, CACHE_SHARE * st.f_blocks * st.f_frsize) + except OSError: + pass + + +def evict(cache, limit): + """Removes the least recently used builds of the cache until it holds at most `limit` bytes, and what jobs that + stopped while they stored a build left there (after STALE_SEC). Anything but a directory is left alone.""" + now = time.time() + builds = [] + for name in os.listdir(cache): + path = os.path.join(cache, name) + try: + st = os.lstat(path) + except OSError: + continue + if not stat.S_ISDIR(st.st_mode): + continue + if name.startswith("."): + if now - st.st_mtime > STALE_SEC: + shutil.rmtree(path, ignore_errors=True) + continue + try: + with open(os.path.join(path, SIZE)) as f: + size = int(f.read()) + except (OSError, ValueError): + size = disk_usage(path) + builds.append((st.st_mtime, path, size)) + total = sum(size for _, _, size in builds) + for _, path, size in sorted(builds): + if total <= limit: + break + shutil.rmtree(path, ignore_errors=True) + total -= size + + +def stage(solution_dir, build_dir=None, arch=None, extra_cuda_cflags=(), cache=None): + """Writes the sources into build_dir (a fresh temporary directory if None) and builds them there. + + Returns (workdir, spec, build seconds or None when an earlier build was reused). A build_dir that already holds + a build of the same sources and flags is reused, so CI builds each solution once for its rounds. With a cache + (cache_dir(): a development build), a build of the same key is copied from it, and one made is put there.""" + spec, files = read(solution_dir) + arch = arch or gpu_arch() + flags = compile_flags(spec, arch) + # Flags a tool adds that do not change the code, such as -lineinfo for source-level profiles. + flags[1].extend(extra_cuda_cflags) + digest = hashlib.sha256(json.dumps([spec, flags, sorted(files.items())]).encode()).hexdigest() + if build_dir is None: + build_dir = tempfile.mkdtemp(prefix="solution-") + atexit.register(shutil.rmtree, build_dir, True) + build_dir = os.path.abspath(build_dir) + stamp = os.path.join(build_dir, STAMP) + try: + with open(stamp) as f: + if f.read() == digest and os.path.exists(os.path.join(build_dir, MODULE + ".so")): + return build_dir, spec, None + except OSError: + pass + key = cache and cache_key(spec, files, flags, arch, toolchain()) + if key and fetch(cache, key, build_dir): + return build_dir, spec, None + shutil.rmtree(build_dir, ignore_errors=True) + os.makedirs(build_dir) + for rel, text in files.items(): + dest = os.path.join(build_dir, rel) + os.makedirs(os.path.dirname(dest), exist_ok=True) + with open(dest, "w", encoding="utf-8") as f: + f.write(text) + start = time.monotonic() + env = dict(os.environ, PYTORCH_ALLOC_CONF="expandable_segments:True") + cmd = [sys.executable, "-s", os.path.abspath(__file__), "_build", json.dumps(flags)] + try: + p = subprocess.run(cmd, cwd=build_dir, env=env, capture_output=True, text=True, timeout=BUILD_TIMEOUT_SEC) + except subprocess.TimeoutExpired as e: + raise SolutionError(f"the build timed out after {BUILD_TIMEOUT_SEC} s") from e + seconds = time.monotonic() - start + if p.returncode != 0 or not os.path.exists(os.path.join(build_dir, MODULE + ".so")): + log = (p.stdout + p.stderr).strip() + raise SolutionError(f"the build failed after {seconds:.0f} s:\n{_errors(log)}") + with open(stamp, "w") as f: + f.write(digest) + if key: + store(cache, key, build_dir, files) + return build_dir, spec, seconds + + +def _errors(log, limit=6000): + """The compiler's error lines with some context, else the log's tail.""" + lines = log.splitlines() + keep = [i for i, line in enumerate(lines) if "error" in line.lower()] + if not keep: + return log[-limit:] + picked = sorted({j for i in keep for j in range(max(0, i - 2), min(len(lines), i + 3))}) + return "\n".join(lines[j] for j in picked)[-limit:] + + +# fastvideo-kernel's CUTLASS submodule (git submodule update --init fastvideo-kernel/include/cutlass). +_CUTLASS_SUBMODULE = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "include", + "cutlass") + + +def _build(flags): + """The compile step, run in the working directory by stage() in its own process.""" + import torch.utils.cpp_extension as ext + + cflags, cuda_cflags, ld_flags = json.loads(flags) + here = os.getcwd() + srcs = sorted(os.path.join(here, p) for p in os.listdir(here) + if p.endswith(CPP_SOURCES) and os.path.isfile(os.path.join(here, p))) + if not srcs: + raise RuntimeError("No CUDA/C++ source files found in working directory") + cutlass = os.environ.get("CUTLASS_DIR") or _CUTLASS_SUBMODULE + try: + ext.load(name=MODULE, sources=srcs, extra_cuda_cflags=cuda_cflags, extra_cflags=cflags, extra_ldflags=ld_flags, + extra_include_paths=[here, f"{cutlass}/include", f"{cutlass}/tools/util/include"], + build_directory=here, verbose=True) + except ImportError as e: + # Built, and not loadable here: a machine without a GPU has no driver to load it against, and whoever loads it + # says whether it loads where it runs. Where a GPU is, a module that does not load (an undefined symbol, a + # missing library) is a build that failed, and is neither stamped nor kept in the cache. + import torch + + if torch.cuda.is_available(): + raise + print(f"built; not loaded on this machine: {e}", flush=True) + built = [p for p in os.listdir(here) if p.startswith(MODULE) and p.endswith(".so") and p != MODULE + ".so"] + if built: + os.rename(built[0], MODULE + ".so") + + +class Solution: + """A loaded solution: fn, the entry (run), called as the deployment calls it, fn(*inputs, *outputs); prepare, the + module's prepare function or None; spec (solution.json), entry, languages, language, module, workdir, and + build_seconds (None when an earlier build was reused, or when the build was staged elsewhere).""" + + def __init__(self, spec, fn, prepare, module, workdir, build_seconds): + self.spec = spec + self.languages = list(spec["languages"]) + self.language = self.languages[0] + self.entry = spec["entry_point"] + self.fn, self.prepare, self.module = fn, prepare, module + self.workdir, self.build_seconds = workdir, build_seconds + + +def import_built(workdir, spec, chdir=True, build_seconds=None): + """Imports the module that stage() built in workdir and takes its entry and prepare from it. chdir: make the + working directory the current one, as in deployment.""" + import torch # noqa: F401 (the module needs libtorch loaded first) + + if chdir: + os.chdir(workdir) + file, func = spec["entry_point"].split("::") + module_spec = importlib.util.spec_from_file_location(MODULE, os.path.join(workdir, MODULE + ".so")) + module = importlib.util.module_from_spec(module_spec) + try: + module_spec.loader.exec_module(module) + except Exception as e: + raise SolutionError(f"the built module does not load: {e}") from e + fn = getattr(module, func, None) + if not callable(fn): + raise SolutionError(f"{file} defines no function {func} in its PYBIND11_MODULE") + prepare = getattr(module, PREPARE, None) + if prepare is not None and not callable(prepare): + raise SolutionError(f"the module's {PREPARE} is not a function") + return Solution(spec, fn, prepare, module, workdir, build_seconds) + + +def load(solution_dir, build_dir=None, chdir=True, extra_cuda_cflags=(), cache=None): + """Builds and loads a solution as it is deployed; see the module docstring. + + build_dir: where to write and build it (a fresh temporary directory if None; reused if it holds the same build). + chdir: make the working directory the current one, as in deployment. extra_cuda_cflags: flags a tool adds to the + deployed ones, that must not change the code (e.g. -lineinfo). cache: the node's cache of development builds + (cache_dir()), or None to build fresh.""" + workdir, spec, seconds = stage(solution_dir, build_dir, extra_cuda_cflags=extra_cuda_cflags, cache=cache) + return import_built(workdir, spec, chdir=chdir, build_seconds=seconds) + + +def main(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + c = sub.add_parser("check", help="build and load a solution") + c.add_argument("solution", nargs="?", default="solution") + c.add_argument("--build-dir") + d = sub.add_parser("build", help="write and build a solution in a directory, without loading it") + d.add_argument("solution", nargs="?", default="solution") + d.add_argument("build_dir", nargs="?", default="build") + d.add_argument("--arch", help="the architecture to build for, e.g. sm_100a (default: this machine's GPU's)") + for q in (c, d): + q.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds") + b = sub.add_parser("_build") + b.add_argument("flags") + a = p.parse_args() + if a.cmd == "_build": + _build(a.flags) + return + cache = None if a.no_cache else cache_dir() + try: + if a.cmd == "build": + where, _, seconds = stage(os.path.abspath(a.solution), os.path.abspath(a.build_dir), a.arch, cache=cache) + module = os.path.join(where, MODULE + ".so") + print(f"built in {seconds:.1f} s: {module}" if seconds is not None else f"reused a build: {module}") + return + sol = load(os.path.abspath(a.solution), a.build_dir and os.path.abspath(a.build_dir), cache=cache) + except SolutionError as e: + sys.exit(f"FAILED: {e}") + built = f", built in {sol.build_seconds:.1f} s" if sol.build_seconds is not None else ", a build reused" + prepare = f", with {PREPARE}" if sol.prepare is not None else f", without {PREPARE}" + print(f"loaded {sol.entry} ({', '.join(sol.languages)}){prepare}{built}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/wan22_vae_decode/tools/spec.py b/fastvideo-kernel/wan22_vae_decode/tools/spec.py new file mode 100644 index 000000000..f995bbe5b --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/tools/spec.py @@ -0,0 +1,428 @@ +"""The task: decoding Wan2.2 VAE latents into video on one, two or four GB200 GPUs of one node, as bench.py, ci.py and +bounds.py read it. It defines the clips and workloads, the weights' order, the data, the error allowance and the work +of one decode. + +What is computed +---------------- +reference.py (in the kernel repo) defines the result: diffusers' AutoencoderKLWan._decode of the Wan2.2-TI2V-5B VAE, +latents z [1, 48, T, 44, 78] to a video [1, 3, 4T-3, 704, 1248] in [-1, 1]: post_quant_conv, then the decoder one +latent frame at a time with its causal caches (the first latent frame yields one frame, every later one four), +unpatchify 2 x 2, clamp. The judge compares a solution with the fp32 decode of the same latents, cached as videos in +$DATA (see Correctness). + +The interface (as loader.py builds it and bench.py calls it, positional, destination passing) +------------- + run(z, weights, rank, world, pg, out) + prepare(weights, rank, world, pg) -> list[Tensor] optional + +- z: [1, 48, T, 44, 78] bf16, contiguous, the full latents on every rank: what a pipeline passes to + vae.decode(latents), rounded to bf16. +- weights: the 110 tensors of WEIGHTS below, bf16, on this rank's GPU; or what prepare returned. +- rank, world: this process's rank and the number of ranks (1, 2 or 4), rank r on GPU r. +- pg: torch.distributed's NCCL ProcessGroup over the ranks; a C++ entry takes it as + `const c10::intrusive_ptr&` (from ). +- out: on rank 0, the video [1, 3, 4T-3, 704, 1248] bf16, contiguous, from torch.empty and not zeroed (every element + must be written); on the other ranks an empty tensor. When run has returned on every rank, rank 0's out holds the + video. +- prepare, when the module exports it, is called once per process before any timed call, and is not timed. It sees + the weights, not the latents; its result replaces `weights` in every run. + +Numerics: bf16 weights and activations with fp32 accumulation, and fp32 statistics in the RMS norms; nothing narrower +(no FP8, FP6 or FP4). This is a rule for the reviewers: no check reads it from the code. + +Workloads +--------- +A workload is a clip decoded at a world size, named x. A clip is the first T latent frames of one of +three real 20 s clips encoded by the VAE (latents_.pt in $DATA): + + clip source latent frames T video frames 4T-3 worlds timed + f1 row_000168 1 1 1, 4 no + f5 row_000168 2 5 1, 4 no + f17 row_000168 5 17 1, 4 no + f121 row_000168 31 121 1, 2, 4 yes + f241 row_000129 61 241 1, 2, 4 yes + f481 row_000090 121 481 1, 2, 4 yes + +The nine timed workloads (f121x1 ... f481x4) make the objective, their geomean; the six short ones (f1x1 ... f17x4) +are checked for correctness and never timed: a decoder sees every clip length, and the first latent frame is a +special case. + +Correctness +----------- +On rank 0's out, against ref_.pt, the fp32 decode of the clip's latents (fp32 latents, TF32 off), stored as fp16: +every element finite and within [-1, 1]; and with d = out - ref in fp32, + + max |d| <= min(MAX_ABS_CAP, ERROR_FACTOR x bf16_max_abs) + rms(d) <= ERROR_FACTOR x bf16_rms + +bf16_max_abs and bf16_rms being the error of the release's bf16 decode (bf16 weights under bf16 autocast, the +latents rounded to bf16) against that fp32 decode on the same clip, from calibration.json. PSNR, 10 log10(4 / mse), is +shown beside them. + +The weights +----------- +run() gets the decoder's parameters and post_quant_conv's, 110 tensors, as a list in diffusers' registration order +(that of AutoencoderKLWan.named_parameters(), post_quant_conv first, then the decoder from conv_in to conv_out; +$DATA/weight_names.json lists the same), bf16, in PyTorch's layouts: a conv3d's weight is [out, in, kt, kh, kw], a +conv2d's [out, in, kh, kw], a bias [out], an RMS norm's gain [C, 1, 1, 1] (the attention's [C, 1, 1]). +$DATA/decoder_weights_bf16.pt is that list: the release's fp32 weights rounded to bf16 to nearest even, as diffusers' +from_pretrained(torch_dtype=torch.bfloat16) casts them. The order is not the order of execution: the mid block +registers its attention before its two residual blocks, though it runs resnets.0, the attention, then resnets.1; an +upsampler registers its 3x3 conv (resample.1) before its time conv, though the time conv runs first. The attention's +to_qkv gives q, k and v as its output channels 0-1023, 1024-2047 and 2048-3071. A time conv C -> 2C gives frame 2t of +its output as channels 0..C-1 and frame 2t+1 as channels C..2C-1. conv_out's 12 channels are (colour k, column offset +j, row offset i) of 2 x 2 pixel patches: video[k, f, 2y + i, 2x + j] = channel 4k + 2j + i at (f, y, x). + + # name shape role + 0 post_quant_conv.weight [48, 48, 1, 1, 1] 1x1x1 conv 48 -> 48, on the latents + 1 post_quant_conv.bias [48] + 2 decoder.conv_in.weight [1024, 48, 3, 3, 3] causal 3x3x3 conv 48 -> 1024 + 3 decoder.conv_in.bias [1024] + 4 decoder.mid_block.attentions.0.norm.gamma [1024, 1, 1] RMS norm gain, before to_qkv + 5 decoder.mid_block.attentions.0.to_qkv.weight [3072, 1024, 1, 1] 1x1 conv2d 1024 -> 3072 (q, k, v) + 6 decoder.mid_block.attentions.0.to_qkv.bias [3072] + 7 decoder.mid_block.attentions.0.proj.weight [1024, 1024, 1, 1] 1x1 conv2d 1024 -> 1024 + 8 decoder.mid_block.attentions.0.proj.bias [1024] + 9 decoder.mid_block.resnets.0.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 10 decoder.mid_block.resnets.0.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 11 decoder.mid_block.resnets.0.conv1.bias [1024] + 12 decoder.mid_block.resnets.0.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 13 decoder.mid_block.resnets.0.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 14 decoder.mid_block.resnets.0.conv2.bias [1024] + 15 decoder.mid_block.resnets.1.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 16 decoder.mid_block.resnets.1.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 17 decoder.mid_block.resnets.1.conv1.bias [1024] + 18 decoder.mid_block.resnets.1.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 19 decoder.mid_block.resnets.1.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 20 decoder.mid_block.resnets.1.conv2.bias [1024] + 21 decoder.up_blocks.0.resnets.0.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 22 decoder.up_blocks.0.resnets.0.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 23 decoder.up_blocks.0.resnets.0.conv1.bias [1024] + 24 decoder.up_blocks.0.resnets.0.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 25 decoder.up_blocks.0.resnets.0.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 26 decoder.up_blocks.0.resnets.0.conv2.bias [1024] + 27 decoder.up_blocks.0.resnets.1.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 28 decoder.up_blocks.0.resnets.1.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 29 decoder.up_blocks.0.resnets.1.conv1.bias [1024] + 30 decoder.up_blocks.0.resnets.1.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 31 decoder.up_blocks.0.resnets.1.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 32 decoder.up_blocks.0.resnets.1.conv2.bias [1024] + 33 decoder.up_blocks.0.resnets.2.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 34 decoder.up_blocks.0.resnets.2.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 35 decoder.up_blocks.0.resnets.2.conv1.bias [1024] + 36 decoder.up_blocks.0.resnets.2.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 37 decoder.up_blocks.0.resnets.2.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 38 decoder.up_blocks.0.resnets.2.conv2.bias [1024] + 39 decoder.up_blocks.0.upsampler.resample.1.weight [1024, 1024, 3, 3] 3x3 conv2d 1024 -> 1024, after x2 + 40 decoder.up_blocks.0.upsampler.resample.1.bias [1024] + 41 decoder.up_blocks.0.upsampler.time_conv.weight [2048, 1024, 3, 1, 1] causal 3x1x1 conv 1024 -> 2048 + 42 decoder.up_blocks.0.upsampler.time_conv.bias [2048] + 43 decoder.up_blocks.1.resnets.0.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 44 decoder.up_blocks.1.resnets.0.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 45 decoder.up_blocks.1.resnets.0.conv1.bias [1024] + 46 decoder.up_blocks.1.resnets.0.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 47 decoder.up_blocks.1.resnets.0.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 48 decoder.up_blocks.1.resnets.0.conv2.bias [1024] + 49 decoder.up_blocks.1.resnets.1.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 50 decoder.up_blocks.1.resnets.1.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 51 decoder.up_blocks.1.resnets.1.conv1.bias [1024] + 52 decoder.up_blocks.1.resnets.1.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 53 decoder.up_blocks.1.resnets.1.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 54 decoder.up_blocks.1.resnets.1.conv2.bias [1024] + 55 decoder.up_blocks.1.resnets.2.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 56 decoder.up_blocks.1.resnets.2.conv1.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 57 decoder.up_blocks.1.resnets.2.conv1.bias [1024] + 58 decoder.up_blocks.1.resnets.2.norm2.gamma [1024, 1, 1, 1] RMS norm gain, before conv2 + 59 decoder.up_blocks.1.resnets.2.conv2.weight [1024, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 1024 + 60 decoder.up_blocks.1.resnets.2.conv2.bias [1024] + 61 decoder.up_blocks.1.upsampler.resample.1.weight [1024, 1024, 3, 3] 3x3 conv2d 1024 -> 1024, after x2 + 62 decoder.up_blocks.1.upsampler.resample.1.bias [1024] + 63 decoder.up_blocks.1.upsampler.time_conv.weight [2048, 1024, 3, 1, 1] causal 3x1x1 conv 1024 -> 2048 + 64 decoder.up_blocks.1.upsampler.time_conv.bias [2048] + 65 decoder.up_blocks.2.resnets.0.norm1.gamma [1024, 1, 1, 1] RMS norm gain, before conv1 + 66 decoder.up_blocks.2.resnets.0.conv1.weight [512, 1024, 3, 3, 3] causal 3x3x3 conv 1024 -> 512 + 67 decoder.up_blocks.2.resnets.0.conv1.bias [512] + 68 decoder.up_blocks.2.resnets.0.norm2.gamma [512, 1, 1, 1] RMS norm gain, before conv2 + 69 decoder.up_blocks.2.resnets.0.conv2.weight [512, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 512 + 70 decoder.up_blocks.2.resnets.0.conv2.bias [512] + 71 decoder.up_blocks.2.resnets.0.conv_shortcut.weight [512, 1024, 1, 1, 1] 1x1x1 conv 1024 -> 512, shortcut + 72 decoder.up_blocks.2.resnets.0.conv_shortcut.bias [512] + 73 decoder.up_blocks.2.resnets.1.norm1.gamma [512, 1, 1, 1] RMS norm gain, before conv1 + 74 decoder.up_blocks.2.resnets.1.conv1.weight [512, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 512 + 75 decoder.up_blocks.2.resnets.1.conv1.bias [512] + 76 decoder.up_blocks.2.resnets.1.norm2.gamma [512, 1, 1, 1] RMS norm gain, before conv2 + 77 decoder.up_blocks.2.resnets.1.conv2.weight [512, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 512 + 78 decoder.up_blocks.2.resnets.1.conv2.bias [512] + 79 decoder.up_blocks.2.resnets.2.norm1.gamma [512, 1, 1, 1] RMS norm gain, before conv1 + 80 decoder.up_blocks.2.resnets.2.conv1.weight [512, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 512 + 81 decoder.up_blocks.2.resnets.2.conv1.bias [512] + 82 decoder.up_blocks.2.resnets.2.norm2.gamma [512, 1, 1, 1] RMS norm gain, before conv2 + 83 decoder.up_blocks.2.resnets.2.conv2.weight [512, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 512 + 84 decoder.up_blocks.2.resnets.2.conv2.bias [512] + 85 decoder.up_blocks.2.upsampler.resample.1.weight [512, 512, 3, 3] 3x3 conv2d 512 -> 512, after x2 + 86 decoder.up_blocks.2.upsampler.resample.1.bias [512] + 87 decoder.up_blocks.3.resnets.0.norm1.gamma [512, 1, 1, 1] RMS norm gain, before conv1 + 88 decoder.up_blocks.3.resnets.0.conv1.weight [256, 512, 3, 3, 3] causal 3x3x3 conv 512 -> 256 + 89 decoder.up_blocks.3.resnets.0.conv1.bias [256] + 90 decoder.up_blocks.3.resnets.0.norm2.gamma [256, 1, 1, 1] RMS norm gain, before conv2 + 91 decoder.up_blocks.3.resnets.0.conv2.weight [256, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 256 + 92 decoder.up_blocks.3.resnets.0.conv2.bias [256] + 93 decoder.up_blocks.3.resnets.0.conv_shortcut.weight [256, 512, 1, 1, 1] 1x1x1 conv 512 -> 256, shortcut + 94 decoder.up_blocks.3.resnets.0.conv_shortcut.bias [256] + 95 decoder.up_blocks.3.resnets.1.norm1.gamma [256, 1, 1, 1] RMS norm gain, before conv1 + 96 decoder.up_blocks.3.resnets.1.conv1.weight [256, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 256 + 97 decoder.up_blocks.3.resnets.1.conv1.bias [256] + 98 decoder.up_blocks.3.resnets.1.norm2.gamma [256, 1, 1, 1] RMS norm gain, before conv2 + 99 decoder.up_blocks.3.resnets.1.conv2.weight [256, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 256 + 100 decoder.up_blocks.3.resnets.1.conv2.bias [256] + 101 decoder.up_blocks.3.resnets.2.norm1.gamma [256, 1, 1, 1] RMS norm gain, before conv1 + 102 decoder.up_blocks.3.resnets.2.conv1.weight [256, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 256 + 103 decoder.up_blocks.3.resnets.2.conv1.bias [256] + 104 decoder.up_blocks.3.resnets.2.norm2.gamma [256, 1, 1, 1] RMS norm gain, before conv2 + 105 decoder.up_blocks.3.resnets.2.conv2.weight [256, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 256 + 106 decoder.up_blocks.3.resnets.2.conv2.bias [256] + 107 decoder.norm_out.gamma [256, 1, 1, 1] RMS norm gain, before conv_out + 108 decoder.conv_out.weight [12, 256, 3, 3, 3] causal 3x3x3 conv 256 -> 12 + 109 decoder.conv_out.bias [12] + +Up blocks 0 and 1 run at 1024 channels, up block 2 from 1024 to 512, up block 3 from 512 to 256; a residual block's +first conv changes the channel count, and its 1x1x1 conv_shortcut does so on the shortcut. + +Data +---- +$DATA (default /data), read-only: latents_.pt ({"latents": fp32 [1, 48, 121, 44, 78]}), ref_.pt +({"video": fp16 [1, 3, 4T-3, 704, 1248], "source", "latent_frames"}), calibration.json (per clip: bf16_max_abs, +bf16_mean_abs, bf16_rms, bf16_psnr_db, fp32_decode_s, bf16_decode_s), weight_names.json and decoder_weights_bf16.pt. + +The work of one decode (bounds) +------------------------------- +FLOPs, a multiply-add counting 2: every convolution at its full kernel over its output positions (2 x the weight's +elements x frames x rows x columns, zero padding included, which adds a few percent at 44 x 78 and less than one +beyond), plus the attention's two products per latent frame (2 x 2 x 3432^2 x 1024). The frames at each level, for T +latent frames: T at 44 x 78 (post_quant_conv, conv_in, the mid block, up block 0's residual blocks); T - 1 for up +block 0's time conv, which skips the first latent frame; 2T - 1 at 88 x 156 (up block 0's 3x3 conv, up block 1's +residual blocks); 2T - 2 for up block 1's time conv; 4T - 3 at 176 x 312 (up block 1's 3x3 conv, up block 2's +residual blocks) and at 352 x 624 (up block 2's 3x3 conv, up block 3, conv_out). The FLOPs are divided over the +workload's GPUs. Bytes: the video written (on rank 0) and the weights read once. +""" + +import json +import math +import os + +LATENT_CHANNELS = 48 +LATENT_HW = (44, 78) +VIDEO_HW = (704, 1248) +# How far a solution may err, against what the release's bf16 decode errs by on the same clip (see Correctness). +ERROR_FACTOR = 3.0 +MAX_ABS_CAP = 2.0 +WEIGHTS_FILE = "decoder_weights_bf16.pt" +NAMES_FILE = "weight_names.json" + +CLIPS = { + "f1": {"source": "row_000168", "latent_frames": 1}, + "f5": {"source": "row_000168", "latent_frames": 2}, + "f17": {"source": "row_000168", "latent_frames": 5}, + "f121": {"source": "row_000168", "latent_frames": 31}, + "f241": {"source": "row_000129", "latent_frames": 61}, + "f481": {"source": "row_000090", "latent_frames": 121}, +} +# The short clips are checked at one GPU and at four, never timed; the long ones are checked and timed at 1, 2 and 4. +SHORT_CLIPS, LONG_CLIPS = ("f1", "f5", "f17"), ("f121", "f241", "f481") +WORKLOADS = { + **{f"{c}x{n}": {"clip": c, "world": n, "timed": False} for c in SHORT_CLIPS for n in (1, 4)}, + **{f"{c}x{n}": {"clip": c, "world": n, "timed": True} for c in LONG_CLIPS for n in (1, 2, 4)}, +} + +# decoder.up_blocks.{i}: (in channels, out channels, temporal upsampling, spatial upsampling) +UP_BLOCKS = ((1024, 1024, True, True), (1024, 1024, True, True), (1024, 512, False, True), (512, 256, False, False)) +RESNETS_PER_UP_BLOCK = 3 + + +def _weights(): + """(name, shape) of run()'s weights in diffusers' registration order (see The weights).""" + out = [] + + def conv(name, shape): + out.extend([(f"{name}.weight", tuple(shape)), (f"{name}.bias", (shape[0],))]) + + def resnet(name, cin, cout): + out.append((f"{name}.norm1.gamma", (cin, 1, 1, 1))) + conv(f"{name}.conv1", (cout, cin, 3, 3, 3)) + out.append((f"{name}.norm2.gamma", (cout, 1, 1, 1))) + conv(f"{name}.conv2", (cout, cout, 3, 3, 3)) + if cin != cout: + conv(f"{name}.conv_shortcut", (cout, cin, 1, 1, 1)) + + conv("post_quant_conv", (48, 48, 1, 1, 1)) + conv("decoder.conv_in", (1024, 48, 3, 3, 3)) + out.append(("decoder.mid_block.attentions.0.norm.gamma", (1024, 1, 1))) + conv("decoder.mid_block.attentions.0.to_qkv", (3072, 1024, 1, 1)) + conv("decoder.mid_block.attentions.0.proj", (1024, 1024, 1, 1)) + for j in range(2): + resnet(f"decoder.mid_block.resnets.{j}", 1024, 1024) + for i, (cin, cout, temporal, spatial) in enumerate(UP_BLOCKS): + for j in range(RESNETS_PER_UP_BLOCK): + resnet(f"decoder.up_blocks.{i}.resnets.{j}", cin if j == 0 else cout, cout) + if spatial: + conv(f"decoder.up_blocks.{i}.upsampler.resample.1", (cout, cout, 3, 3)) + if temporal: + conv(f"decoder.up_blocks.{i}.upsampler.time_conv", (2 * cout, cout, 3, 1, 1)) + out.append(("decoder.norm_out.gamma", (256, 1, 1, 1))) + conv("decoder.conv_out", (12, 256, 3, 3, 3)) + return tuple(out) + + +WEIGHTS = _weights() +WEIGHT_NAMES = [name for name, _ in WEIGHTS] + + +def data_dir(): + """The task's data: $DATA, else /data, where the jobs mount it.""" + return os.environ.get("DATA") or "/data" + + +def frames(latent_frames): + """The video frames of T latent frames: one for the first, four for every later one.""" + return 4 * latent_frames - 3 + + +def output_shape(workload): + """The shape of rank 0's out at a workload.""" + return (1, 3, frames(CLIPS[WORKLOADS[workload]["clip"]]["latent_frames"]), *VIDEO_HW) + + +def _load(path): + import torch + + try: + return torch.load(path, map_location="cpu", weights_only=True, mmap=True) + except (OSError, RuntimeError) as e: + raise ValueError(f"cannot read {path}: {e}") from e + + +def make_inputs(workload, device): + """[z]: the clip's latents, the first T latent frames of its source, rounded to bf16, contiguous on `device`.""" + import torch + + clip = CLIPS[WORKLOADS[workload]["clip"]] + path = os.path.join(data_dir(), f"latents_{clip['source']}.pt") + latents = _load(path)["latents"] + t = clip["latent_frames"] + if latents.dim() != 5 or tuple(latents.shape[:2]) != (1, LATENT_CHANNELS) or tuple(latents.shape[3:]) != LATENT_HW \ + or latents.shape[2] < t: + raise ValueError(f"{path}: latents {tuple(latents.shape)}, expected [1, 48, >= {t}, 44, 78]") + return [latents[:, :, :t].to(device=device, dtype=torch.bfloat16).contiguous()] + + +def load_weights(device): + """run()'s weights on `device`: decoder_weights_bf16.pt, checked against WEIGHTS (and weight_names.json where it + is), so that a solution indexes what this file documents.""" + import torch + + listed = os.path.join(data_dir(), NAMES_FILE) + if os.path.exists(listed): + with open(listed) as f: + names = [(e["name"], tuple(e["shape"])) for e in json.load(f)] + if names != list(WEIGHTS): + raise ValueError(f"{listed} does not list spec.WEIGHTS: the weights' order or shapes changed") + path = os.path.join(data_dir(), WEIGHTS_FILE) + weights = _load(path) + if not isinstance(weights, list) or len(weights) != len(WEIGHTS): + raise ValueError(f"{path}: expected a list of {len(WEIGHTS)} tensors") + for i, (t, (name, shape)) in enumerate(zip(weights, WEIGHTS)): + if not isinstance(t, torch.Tensor) or t.dtype != torch.bfloat16 or tuple(t.shape) != shape: + got = f"{t.dtype} {tuple(t.shape)}" if isinstance(t, torch.Tensor) else type(t).__name__ + raise ValueError(f"{path}: tensor {i} is {got}, expected {name}, bf16 {shape}") + return [t.to(device).contiguous() for t in weights] + + +def load_reference(clip): + """The clip's fp32 decode, as stored: a CPU fp16 tensor [1, 3, 4T-3, 704, 1248] (memory-mapped).""" + import torch + + path = os.path.join(data_dir(), f"ref_{clip}.pt") + ref = _load(path) + t, source = CLIPS[clip]["latent_frames"], CLIPS[clip]["source"] + video = ref.get("video") if isinstance(ref, dict) else None + shape = (1, 3, frames(t), *VIDEO_HW) + if not isinstance(video, torch.Tensor) or tuple(video.shape) != shape or video.dtype != torch.float16: + raise ValueError(f"{path}: expected video, fp16 {list(shape)}") + if ref.get("source") != source or ref.get("latent_frames") != t: + raise ValueError(f"{path} decodes {ref.get('source')}'s first {ref.get('latent_frames')} latent frames, not " + f"{source}'s first {t}") + return video + + +def calibration(clip): + """The release's bf16 decode of the clip against its fp32 decode (calibration.json's entry).""" + path = os.path.join(data_dir(), "calibration.json") + try: + with open(path) as f: + entry = json.load(f)[clip] + except (OSError, ValueError, KeyError) as e: + raise ValueError(f"{path} has no entry for {clip}: {e}") from e + if entry.get("source") != CLIPS[clip]["source"] or entry.get("latent_frames") != CLIPS[clip]["latent_frames"]: + raise ValueError(f"{path}: {clip} is calibrated on another clip ({entry.get('source')}, " + f"{entry.get('latent_frames')} latent frames)") + return entry + + +def allowance(clip): + """{"max_abs", "rms"}: how far a solution's video may be from the fp32 decode of the clip (see Correctness).""" + entry = calibration(clip) + return {"max_abs": min(MAX_ABS_CAP, ERROR_FACTOR * entry["bf16_max_abs"]), "rms": ERROR_FACTOR * entry["bf16_rms"]} + + +def _conv_level(name, t): + """(output frames, upsampling factor in space) of the conv whose weight is `name`, decoding t latent frames.""" + if name.startswith(("post_quant_conv.", "decoder.conv_in.", "decoder.mid_block.", "decoder.up_blocks.0.resnets.")): + return t, 1 + if name.startswith("decoder.up_blocks.0.upsampler.time_conv."): + return t - 1, 1 + if name.startswith(("decoder.up_blocks.0.upsampler.resample.", "decoder.up_blocks.1.resnets.")): + return 2 * t - 1, 2 + if name.startswith("decoder.up_blocks.1.upsampler.time_conv."): + return 2 * t - 2, 2 + if name.startswith(("decoder.up_blocks.1.upsampler.resample.", "decoder.up_blocks.2.resnets.")): + return 4 * t - 3, 4 + return 4 * t - 3, 8 + + +def decode_flops(t): + """FLOPs of one decode of t latent frames (see The work of one decode).""" + h, w = LATENT_HW + total = 0 + for name, shape in WEIGHTS: + if name.endswith(".weight") and len(shape) >= 4: + count, scale = _conv_level(name, t) + total += 2 * math.prod(shape) * count * (h * scale) * (w * scale) + return total + 2 * 2 * (h * w) ** 2 * 1024 * t + + +def bounds(workload): + """{"flops", "bytes", "gpus"} of one call at the workload, for bounds.py: the decode's FLOPs (spread over the + world's GPUs), and the bytes that must move: the video written and the weights read.""" + entry = WORKLOADS[workload] + t = CLIPS[entry["clip"]]["latent_frames"] + video = math.prod(output_shape(workload)) * 2 + weights = sum(math.prod(shape) for _, shape in WEIGHTS) * 2 + return {"flops": decode_flops(t), "bytes": video + weights, "gpus": entry["world"]} + + +# The stack rules beside audit.py's: run() reads nothing but its arguments. +RULES = { + "file-io": { + "include": r"(fstream|fcntl\.h|sys/mman\.h|filesystem|experimental/filesystem)", + "text": r"\b(fopen|fopen64|freopen|fdopen|fread|fwrite|ifstream|ofstream|fstream|filebuf|mmap|mmap64|shm_open|" + r"torch::load|torch::save|torch::pickle_load|torch::pickle_save|torch::jit::load)\b|" + r"::\s*(open|openat|read|pread|write|pwrite)\s*\(", + "why": "run() and prepare() read nothing but their arguments and write nothing but out: no files (the " + "reference videos lie on the job's disk) and nothing that outlives a call"}, +} + + +if __name__ == "__main__": + for w in WORKLOADS: + b = bounds(w) + print(f"{w:<8} {b['flops'] / 1e15:8.4f} PFLOP over {b['gpus']} GPU(s), {b['bytes'] / 1e9:6.2f} GB") diff --git a/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py b/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py new file mode 100644 index 000000000..b32446f67 --- /dev/null +++ b/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py @@ -0,0 +1,143 @@ +"""The decoder of the Wan2.2 VAE (Wan2.2-TI2V-5B, 48 latent channels) on sm_100a (B200 / GB200), on one or several GPUs +of a node, built from solution/ on first use. + +It decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T - 3, 704, 1248] (bf16, in [-1, 1]), as +diffusers' AutoencoderKLWan.decode(z).sample computes it (reference.py defines the result), for any T >= 1. Other +latent sizes and batches above one are refused: the kernels are specialised to 704 x 1248. + + decoder = WanVAEDecoder(vae.state_dict(), device="cuda") # diffusers' AutoencoderKLWan, or its state dict + video = decoder.decode(latents) # [1, 3, 4T - 3, 704, 1248] bf16 + +Over several GPUs, every rank builds a decoder with the same process group and calls decode with the full latents; +rank 0 gets the video and the other ranks None. Without a process group, decode uses torch.distributed's default +group, and initializes a single-rank NCCL group when there is none (one GPU only). With local=True each process decodes +on its own GPU alone, whatever torch.distributed holds: one rank of one, so the kernels never communicate (comm.cu +returns before any exchange below two ranks) and no group has to be created collectively. + +frames_per_pass sets how many latent frames each GPU carries through one pass of the decoder, which trades memory for +speed; the video is the same at any value. On one GB200 at 8 to 20 s of video: 2 frames peak at 14-16 GiB allocated +(x0.85 the speed of 16 frames), 4 (the default) at 21-22 GiB (x0.92), 8 at 34-35 GiB (x0.97), 16 at 60-61 GiB. The +allocator reserves more than it allocates; PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True keeps the difference small. +""" + +import atexit +import functools +import importlib.util +import os +import sys + +import torch +import torch.distributed as dist +from torch import Tensor + +HERE = os.path.dirname(os.path.abspath(__file__)) +LATENT_CHANNELS, LATENT_HW, VIDEO_HW = 48, (44, 78), (704, 1248) + + +def _own(path, name): + """The module at path, imported under a name of this kernel's own. Other kernels built the same way + keep tools and references of the same file names; a bare import would hand whichever loaded first to both.""" + name = f"wan22_vae_decode_{name}" + module = sys.modules.get(name) + if module is None: + module_spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(module_spec) + sys.modules[name] = module + try: + module_spec.loader.exec_module(module) + except BaseException: + del sys.modules[name] + raise + return module + + +@functools.lru_cache(maxsize=None) +def _tools(): + """tools/loader.py and tools/spec.py.""" + return (_own(os.path.join(HERE, "tools", "loader.py"), "loader"), + _own(os.path.join(HERE, "tools", "spec.py"), "spec")) + + +@functools.lru_cache(maxsize=None) +def _module(): + """The extension built from solution/, once per process. Builds are kept in WAN22_VAE_DECODE_BUILD_DIR (default + ~/.cache/fastvideo/wan22_vae_decode), one directory per local rank, and reused while the sources and flags match.""" + loader, _ = _tools() + base = os.environ.get("WAN22_VAE_DECODE_BUILD_DIR", os.path.expanduser("~/.cache/fastvideo/wan22_vae_decode")) + build_dir = os.path.join(base, f"rank{os.environ.get('LOCAL_RANK', '0')}") + os.makedirs(base, exist_ok=True) + return loader.load(os.path.join(HERE, "solution"), build_dir=build_dir, chdir=False).module + + +def weight_names() -> list[str]: + """The 110 parameters the decoder takes, under diffusers' AutoencoderKLWan names, in the order it takes them.""" + return list(_tools()[1].WEIGHT_NAMES) + + +def decoder_weights(state_dict, device=None) -> list[Tensor]: + """The decoder's weights from an AutoencoderKLWan state dict (or the module itself): its post_quant_conv and + decoder parameters, in weight_names() order, as contiguous bf16 tensors on `device`.""" + if isinstance(state_dict, torch.nn.Module): + state_dict = state_dict.state_dict() + _, spec = _tools() + weights = [] + for name, shape in spec.WEIGHTS: + if name not in state_dict: + raise KeyError(f"the state dict has no {name}: is it the Wan2.2-TI2V-5B VAE's?") + w = state_dict[name] + if tuple(w.shape) != shape: + raise ValueError(f"{name} is {tuple(w.shape)}, expected {shape}") + weights.append(w.detach().to(device=device, dtype=torch.bfloat16).contiguous()) + return weights + + +def _destroy_if_initialized(): + # The group may be gone already (a caller destroyed it), and destroying it twice raises. + if dist.is_initialized(): + dist.destroy_process_group() + + +def _default_group(): + if not dist.is_initialized(): + dist.init_process_group("nccl", store=dist.HashStore(), rank=0, world_size=1) + atexit.register(_destroy_if_initialized) + return dist.group.WORLD + + +class WanVAEDecoder: + """The decoder with its weights prepared once (the layout and the Winograd transforms the kernels read).""" + + def __init__(self, state_dict, device=None, pg=None, local=False, frames_per_pass=4): + if frames_per_pass < 1: + raise ValueError(f"frames_per_pass must be at least 1, not {frames_per_pass}") + self.frames_per_pass = int(frames_per_pass) + self.device = torch.device(device if device is not None else "cuda") + if self.device.index is None: + self.device = torch.device("cuda", torch.cuda.current_device()) + with torch.cuda.device(self.device): + if local: + if pg is not None: + raise ValueError("local=True decodes on this GPU alone; it takes no process group") + # At world 1 the group is only carried, never used: any initialized group serves. + self.pg = dist.group.WORLD if dist.is_initialized() else _default_group() + self.rank, self.world = 0, 1 + else: + self.pg = pg if pg is not None else _default_group() + self.rank, self.world = dist.get_rank(self.pg), dist.get_world_size(self.pg) + if self.world > 8: + raise ValueError(f"the decoder splits over at most 8 GPUs, not {self.world}") + self.weights = _module().prepare(decoder_weights(state_dict, self.device), self.rank, self.world, + self.pg) + + @torch.no_grad() + def decode(self, z: Tensor) -> Tensor | None: + """The video of latents z [1, 48, T, 44, 78] on rank 0, None on the other ranks.""" + if z.dim() != 5 or z.shape[0] != 1 or z.shape[1] != LATENT_CHANNELS or tuple(z.shape[3:]) != LATENT_HW: + raise ValueError(f"latents {tuple(z.shape)}: expected [1, 48, T, 44, 78]") + z = z.to(device=self.device, dtype=torch.bfloat16).contiguous() + frames = 4 * z.shape[2] - 3 + with torch.cuda.device(self.device): + out = torch.empty((1, 3, frames, *VIDEO_HW) if self.rank == 0 else (0,), dtype=torch.bfloat16, + device=self.device) + _module().run_with(z, self.weights, self.rank, self.world, self.pg, out, self.frames_per_pass) + return out if self.rank == 0 else None diff --git a/fastvideo/configs/pipelines/pan2.py b/fastvideo/configs/pipelines/pan2.py index 8e58c33b9..c56b58ec1 100644 --- a/fastvideo/configs/pipelines/pan2.py +++ b/fastvideo/configs/pipelines/pan2.py @@ -24,6 +24,17 @@ class PAN2PipelineConfig(PipelineConfig): dit_precision: str = "bf16" vae_precision: str = "fp32" + # Decode the video with fastvideo-kernel's Wan2.2 decoder (B200 / GB200 / B300; a 704 x 1248 video, batch 1): bf16 + # with fp32 accumulation, about 7x faster than the fp32 decode on a B300, at PSNR 61-63 dB and SSIM 0.998 against + # it. Other devices and sizes, and False, keep the fp32 decode. + vae_decode_kernel: bool = True + # With vae_decode_kernel, split the decode over the sequence-parallel GPUs (2, 4 or 8) rather than decoding the + # whole video on each of them. On B300s a 20 s video decodes in 3.0 s on one GPU, 1.55 s on 2 (the same video, bit + # for bit), 0.79 s on 4 and 0.47 s on 8 (rounded differently, at the same PSNR against the fp32 decode). + vae_decode_kernel_parallel: bool = True + # Run the DiT's fused Triton kernels (Q/K RMSNorm + RoPE, LayerNorm + modulation): 21% less denoising time on 8 + # B300s, as far from the eager path's samples as bf16 rounding differences carry them. + dit_fused_kernels: bool = True # Unset: the checkpoint's scheduler config decides, which is shift 9 for IFM/PAN2-TEST, the reference sampler's # fixed value. A fixed default here would override the released shift with whatever a test checkpoint declares. flow_shift: float | None = None diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 89b9b3953..db05a3801 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -162,6 +162,31 @@ def _quantize_video_frames_to_uint8(src: torch.Tensor) -> torch.Tensor: return rearrange(vid_u8, "b c t h w -> t b c h w") +def _video_frames(src: torch.Tensor) -> list[np.ndarray]: + """`[b, c, t, h, w]` decoded output -> one host uint8 frame `[h, w, 3]` per time step, the batch tiled six to a + row by `torchvision.utils.make_grid` (which also repeats a single channel three times). + + One RGB video, every request of most pipelines, is quantized and laid out frames-first and channels-last on + ``src``'s device, copied to the host once, and its frames are views of that one array. Building each frame on the + CPU instead (a channels-last copy per frame) took 0.48 s against 0.38 s for 241 frames of 704 x 1248 on a B300 + node, and 1-2 s in a job pinned to two CPUs. Same values either way; the out-of-memory fallback of + `_quantize_video_frames_to_uint8` applies to both. + """ + if src.ndim == 5 and src.shape[0] == 1 and src.shape[1] == 3: + try: + host = pixels_to_uint8(src)[0].permute(1, 2, 3, 0).contiguous().cpu() + except torch.cuda.OutOfMemoryError: + logger.warning("GPU out of memory laying out the decoded video's frames; doing it on the CPU instead.") + src_cpu = src.detach().to("cpu") + host = pixels_to_uint8(src_cpu if src_cpu.dtype == torch.uint8 else src_cpu.float())[0].permute( + 1, 2, 3, 0).contiguous() + return list(host.numpy()) + return [ + torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() + for x in _quantize_video_frames_to_uint8(src) + ] + + def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None: """Validate typed stage overrides against the model's registered preset.""" if not request.stage_overrides: @@ -1041,10 +1066,7 @@ class VideoGenerator: # `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or # `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same # data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback. - vid_u8 = _quantize_video_frames_to_uint8(output_batch.output) - frames = [ - torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8 - ] + frames = _video_frames(output_batch.output) postprocess_time = time.perf_counter() - postprocess_start logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time) if logging_info is not None: diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 56c9a8cb1..588bfe924 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -635,6 +635,12 @@ FASTVIDEO_H3_VSA_FINE_ATTENTION = EnvChoice( "emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized " "BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over " "FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.") +FASTVIDEO_WAN22_VAE_DECODE_DIR = EnvPath("", + category="performance", + doc="Directory of the Wan2.2 VAE decoder kernels " + "(fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds " + "on first use, when its config sets decode_kernel. Empty takes the copy in " + "this checkout.") FASTVIDEO_H3_VSA_HEADS_FIRST_TILE = EnvBool(False, category="performance", doc="VSA-H3 64/128-token tiles: scatter rows straight into the heads-first " @@ -924,6 +930,10 @@ FASTVIDEO_TEST_SD35_MODEL_DIR = EnvStr("stabilityai/stable-diffusion-3.5-medium" category="test", doc="Model for the Stable Diffusion 3.5 SSIM test.", deprecated_names=("SD35_MODEL_DIR", )) +FASTVIDEO_TEST_PAN2_MODEL_DIR = EnvPath(None, + category="test", + doc="PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU " + "tests; unset skips them.") FASTVIDEO_TEST_TAEH3_REFERENCE_DIR = EnvStr(None, category="test", doc="Upstream taehv checkout for the MLX TAEH3 parity test.", diff --git a/fastvideo/models/dits/pan2.py b/fastvideo/models/dits/pan2.py index 40c809d88..157ce7e62 100644 --- a/fastvideo/models/dits/pan2.py +++ b/fastvideo/models/dits/pan2.py @@ -16,16 +16,26 @@ from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, sequence_model_parallel_shard) -from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_1d_rotary_pos_embed from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.pan2_fusions import layer_norm_modulate, qk_rmsnorm_rope from fastvideo.platforms import AttentionBackendEnum +def _fusable_layer_norm(x: torch.Tensor) -> bool: + """Whether ``layer_norm_modulate`` computes what the eager ``norm(x) * (1 + scale) + shift`` does here. + + The fused op rounds once from FP32, which is the eager result under autocast (FP32 LayerNorm and modulation, cast + by the next linear layer) and in FP32; a low-precision tensor outside autocast rounds after every eager op instead. + """ + return x.dtype == torch.float32 or torch.is_autocast_enabled(x.device.type) + + class PAN2SwiGLU(nn.Module): """SwiGLU MLP of the text refiner; the fused input projection stores ``[value; gate]``.""" @@ -102,6 +112,9 @@ class PAN2TokenRefinerBlock(nn.Module): class PAN2AdaLayerNormZero(nn.Module): """Timestep modulation of the video stream: modulated attention input plus the gate/shift/scale for the rest.""" + # Set by PAN2Transformer3DModel.set_fused_kernels. + fused_kernels = False + def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""): super().__init__() self.linear = ReplicatedLinear(dim, 6 * dim, bias=True, prefix=f"{prefix}.linear") @@ -111,12 +124,18 @@ class PAN2AdaLayerNormZero(nn.Module): temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb, _ = self.linear(F.silu(temb)) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) - x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + if self.fused_kernels and _fusable_layer_norm(x): + x = layer_norm_modulate(x, 1 + scale_msa, shift_msa, self.norm.eps) + else: + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa, shift_mlp, scale_mlp, gate_mlp class PAN2AdaLayerNormOut(nn.Module): + # Set by PAN2Transformer3DModel.set_fused_kernels. + fused_kernels = False + def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""): super().__init__() self.linear = ReplicatedLinear(dim, 2 * dim, bias=True, prefix=f"{prefix}.linear") @@ -125,6 +144,8 @@ class PAN2AdaLayerNormOut(nn.Module): def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: emb, _ = self.linear(F.silu(temb).to(x.dtype)) scale, shift = emb.chunk(2, dim=-1) + if self.fused_kernels and _fusable_layer_norm(x): + return layer_norm_modulate(x, 1 + scale, shift, self.norm.eps) return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] @@ -132,9 +153,13 @@ class PAN2Attention(nn.Module): """Joint attention over ``[video; text]`` with per-stream projections and Q/K RMSNorm. RoPE applies to the video tokens only; ``DistributedAttention`` applies it after its all-to-all, where each rank - holds the full (unpadded) video sequence for its heads. + holds the full (unpadded) video sequence for its heads. With fused kernels the Q/K RMSNorm and RoPE run as one op + before the all-to-all, on this rank's video shard with ``local_freqs_cis``. """ + # Set by PAN2Transformer3DModel.set_fused_kernels. + fused_kernels = False + def __init__(self, dim: int, num_heads: int, @@ -172,16 +197,25 @@ class PAN2Attention(nn.Module): self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], + freqs_cis: tuple[torch.Tensor, torch.Tensor] | None, original_seq_len: int, + local_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: query, key, value = self._project(self.to_qkv, hidden_states) - query = self.norm_q(query).to(value.dtype) - key = self.norm_k(key).to(value.dtype) - encoder_query, encoder_key, encoder_value = self._project(self.add_qkv_proj, encoder_hidden_states) - encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype) - encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype) + if self.fused_kernels and local_freqs_cis is not None: + # RoPE is per token and per head, so it applies to this rank's shard before the all-to-all as well as to + # the gathered sequence after it. + query, key = qk_rmsnorm_rope(query, key, self.norm_q.weight, self.norm_k.weight, + self.norm_q.variance_epsilon, *local_freqs_cis) + encoder_query, encoder_key = qk_rmsnorm_rope(encoder_query, encoder_key, self.norm_added_q.weight, + self.norm_added_k.weight, self.norm_added_q.variance_epsilon) + freqs_cis = None + else: + query = self.norm_q(query).to(value.dtype) + key = self.norm_k(key).to(value.dtype) + encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype) + encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype) hidden_states, encoder_hidden_states = self.attn(query, key, @@ -198,6 +232,9 @@ class PAN2Attention(nn.Module): class PAN2TransformerBlock(nn.Module): + # Set by PAN2Transformer3DModel.set_fused_kernels. + fused_kernels = False + def __init__(self, dim: int, num_heads: int, @@ -234,6 +271,7 @@ class PAN2TransformerBlock(nn.Module): temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], original_seq_len: int, + local_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) ( @@ -244,11 +282,28 @@ class PAN2TransformerBlock(nn.Module): context_scale_mlp, context_gate_mlp, ) = self.scale_shift_table_context.to(temb.dtype).chunk(6, dim=-1) - norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) * - (1 + context_scale_msa) + context_shift_msa) + fused = self.fused_kernels and _fusable_layer_norm(encoder_hidden_states) + if fused: + norm_encoder_hidden_states = layer_norm_modulate(encoder_hidden_states, 1 + context_scale_msa, + context_shift_msa, self.norm1_context.eps) + else: + norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) * + (1 + context_scale_msa) + context_shift_msa) attn_output, context_attn_output = self.attn(norm_hidden_states, norm_encoder_hidden_states, freqs_cis, - original_seq_len) + original_seq_len, local_freqs_cis) + if fused: + # One rounding per gated residual, as PAN2's fused training kernels compute it. + hidden_states = torch.addcmul(hidden_states, attn_output, gate_msa.unsqueeze(1)) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, context_attn_output, context_gate_msa) + norm_hidden_states = layer_norm_modulate(hidden_states, 1 + scale_mlp, shift_mlp, self.norm2.eps) + norm_encoder_hidden_states = layer_norm_modulate(encoder_hidden_states, 1 + context_scale_mlp, + context_shift_mlp, self.norm2_context.eps) + hidden_states = torch.addcmul(hidden_states, self.ff(norm_hidden_states), gate_mlp.unsqueeze(1)) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, self.ff_context(norm_encoder_hidden_states), + context_gate_mlp) + return hidden_states, encoder_hidden_states + hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + context_attn_output * context_gate_msa @@ -329,8 +384,34 @@ class PAN2Transformer3DModel(BaseDiT): prefix=f"{config.prefix}.proj_out") self.gradient_checkpointing = False + # See set_fused_kernels; the pipeline turns it on with PAN2PipelineConfig.dit_fused_kernels. + self.fused_kernels = False self.__post_init__() + def set_fused_kernels(self, enabled: bool) -> None: + """Run the blocks' fused Q/K RMSNorm + RoPE and LayerNorm + modulation ops (``pan2_fusions``). + + The ops launch Triton kernels on CUDA and run their PyTorch reference elsewhere. The forward reproduces the + eager arithmetic up to FP32 summation order, except that each gated residual rounds once instead of twice. + """ + enabled = bool(enabled) + for module in self.modules(): + if isinstance(module, (PAN2TransformerBlock, PAN2Attention, PAN2AdaLayerNormZero, PAN2AdaLayerNormOut)): + module.fused_kernels = enabled + self.fused_kernels = enabled + + @staticmethod + def _local_rotary_pos_embed(freqs_cis: tuple[torch.Tensor, torch.Tensor], + local_len: int) -> tuple[torch.Tensor, torch.Tensor]: + """The RoPE rows of this rank's video shard (identity rotation on the sequence-parallel padding).""" + cos, sin = freqs_cis + start = get_sp_parallel_rank() * local_len + pad = start + local_len - cos.shape[0] + if pad > 0: + cos = F.pad(cos, (0, 0, 0, pad), value=1.0) + sin = F.pad(sin, (0, 0, 0, pad), value=0.0) + return cos[start:start + local_len], sin[start:start + local_len] + def _rotary_pos_embed(self, num_frames: int, height: int, width: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: """Half-dim cos/sin of the 3-axis ``(t, h, w)`` RoPE over the full (unsharded) video token grid.""" @@ -366,6 +447,9 @@ class PAN2Transformer3DModel(BaseDiT): hidden_states = self.x_embedder(hidden_states) hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) + local_freqs_cis = None + if self.fused_kernels: + local_freqs_cis = self._local_rotary_pos_embed(freqs_cis, hidden_states.shape[1]) if timestep.dim() == 0: timestep = timestep.expand(batch_size) @@ -379,10 +463,10 @@ class PAN2Transformer3DModel(BaseDiT): for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len) + block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len, local_freqs_cis) else: hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis, - original_seq_len) + original_seq_len, local_freqs_cis) hidden_states = self.norm_out(hidden_states, temb) hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) diff --git a/fastvideo/models/dits/pan2_fusions.py b/fastvideo/models/dits/pan2_fusions.py new file mode 100644 index 000000000..2785efe76 --- /dev/null +++ b/fastvideo/models/dits/pan2_fusions.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused PAN2 block kernels with backward passes, ported from the PAN2 reference trainer. + +The PAN2 reference fuses its memory-bound per-token work into Triton kernels with hand-written backwards +(``pan/models/transformers/modules/triton_kernels.py``). Two of them map onto FastVideo's unpacked layout: + +- ``qk_rmsnorm_rope``: the per-head query/key RMSNorm and, for the video stream, the interleaved 3-axis RoPE, in one + pass over the projection output (the reference's ``qkv_permute_rms`` + ``qkv_rstd_rope`` without its + concat/interleave permute). The eager path runs about twenty kernels with FP32 temporaries for this. +- ``layer_norm_modulate``: the affine-free LayerNorm and its adaptive ``* (1 + scale) + shift`` (the reference's + ``packed_modulate`` with ``ln_stats``), computed row-wise in one pass instead of from cached statistics. + +Each op is a ``torch.library.custom_op`` with a registered autograd formula, so activation checkpointing and +``torch.compile`` see one node. The forwards reproduce the eager path's roundings (the normalized value and the +weighted value round to the input dtype before the next step, as ``RMSNorm.forward_native`` and ``_apply_rotary_emb`` +do), so the fused and eager forwards differ only by FP32 summation order. The backwards keep FP32 throughout, where +eager autograd rounds some intermediate gradients to BF16. + +On CUDA the ops launch Triton; elsewhere they run the PyTorch reference of the same arithmetic. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - CPU-only installs + triton = None + tl = None + +HAVE_TRITON = triton is not None +# Tests set this to drive CPU tensors through the Triton kernels under ``TRITON_INTERPRET=1``. +_FORCE_TRITON = False + + +def fusions_available(x: torch.Tensor) -> bool: + """Whether the fused ops run their Triton kernels for ``x``.""" + return HAVE_TRITON and (x.is_cuda or _FORCE_TRITON) + + +def _use_triton(x: torch.Tensor) -> bool: + return fusions_available(x) + + +# -------------------------------------------------------------------------------------------------------------------- +# Query/key RMSNorm + RoPE +# -------------------------------------------------------------------------------------------------------------------- + +_QK_BLOCK_ROWS = 16 + +if HAVE_TRITON: + + @triton.jit + def _qk_rmsnorm_rope_fwd_kernel(X, W, COS, SIN, OUT, RSTD, n_rows, seq_len, num_heads, stride_xr, stride_xh, + stride_or, stride_oh, stride_cs, eps, HEAD_DIM: tl.constexpr, + BLOCK_ROWS: tl.constexpr, HAS_ROPE: tl.constexpr): + HALF: tl.constexpr = HEAD_DIM // 2 + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + head = tl.program_id(1) + pairs = tl.arange(0, HALF) + keep = (rows < n_rows)[:, None] + rows64 = rows.to(tl.int64) + x_base = X + rows64[:, None] * stride_xr + head * stride_xh + x_even = tl.load(x_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32) + x_odd = tl.load(x_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32) + mean_square = (tl.sum(x_even * x_even, axis=1) + tl.sum(x_odd * x_odd, axis=1)) / HEAD_DIM + rstd = tl.rsqrt(mean_square + eps) + # The eager RMSNorm rounds the normalized value, then the weighted value, to the input dtype. + n_even = (x_even * rstd[:, None]).to(OUT.dtype.element_ty).to(tl.float32) + n_odd = (x_odd * rstd[:, None]).to(OUT.dtype.element_ty).to(tl.float32) + w_even = tl.load(W + 2 * pairs).to(tl.float32) + w_odd = tl.load(W + 2 * pairs + 1).to(tl.float32) + y_even = (n_even * w_even[None, :]).to(OUT.dtype.element_ty).to(tl.float32) + y_odd = (n_odd * w_odd[None, :]).to(OUT.dtype.element_ty).to(tl.float32) + if HAS_ROPE: + positions = (rows % seq_len).to(tl.int64) + cs_offsets = positions[:, None] * stride_cs + pairs[None, :] + cos = tl.load(COS + cs_offsets, mask=keep, other=1.0) + sin = tl.load(SIN + cs_offsets, mask=keep, other=0.0) + out_even = y_even * cos - y_odd * sin + out_odd = y_odd * cos + y_even * sin + else: + out_even = y_even + out_odd = y_odd + out_base = OUT + rows64[:, None] * stride_or + head * stride_oh + tl.store(out_base + 2 * pairs[None, :], out_even.to(OUT.dtype.element_ty), mask=keep) + tl.store(out_base + 2 * pairs[None, :] + 1, out_odd.to(OUT.dtype.element_ty), mask=keep) + tl.store(RSTD + rows64 * num_heads + head, rstd, mask=rows < n_rows) + + @triton.jit + def _qk_rmsnorm_rope_bwd_kernel(X, W, COS, SIN, GRAD_OUT, RSTD, GRAD_X, GRAD_W_PARTIAL, n_rows, seq_len, num_heads, + stride_xr, stride_xh, stride_gr, stride_gh, stride_cs, HEAD_DIM: tl.constexpr, + BLOCK_ROWS: tl.constexpr, HAS_ROPE: tl.constexpr): + HALF: tl.constexpr = HEAD_DIM // 2 + block = tl.program_id(0) + head = tl.program_id(1) + rows = block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + pairs = tl.arange(0, HALF) + keep = (rows < n_rows)[:, None] + rows64 = rows.to(tl.int64) + x_base = X + rows64[:, None] * stride_xr + head * stride_xh + x_even = tl.load(x_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32) + x_odd = tl.load(x_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32) + rstd = tl.load(RSTD + rows64 * num_heads + head, mask=rows < n_rows, other=0.0) + n_even = x_even * rstd[:, None] + n_odd = x_odd * rstd[:, None] + g_base = GRAD_OUT + rows64[:, None] * stride_gr + head * stride_gh + g_even = tl.load(g_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32) + g_odd = tl.load(g_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32) + if HAS_ROPE: + positions = (rows % seq_len).to(tl.int64) + cs_offsets = positions[:, None] * stride_cs + pairs[None, :] + cos = tl.load(COS + cs_offsets, mask=keep, other=1.0) + sin = tl.load(SIN + cs_offsets, mask=keep, other=0.0) + # The transpose of each pair's rotation. + gy_even = g_even * cos + g_odd * sin + gy_odd = g_odd * cos - g_even * sin + else: + gy_even = g_even + gy_odd = g_odd + # The weight multiplies the rounded normalized value; its gradient sums over this block's rows. + rounded_even = n_even.to(GRAD_X.dtype.element_ty).to(tl.float32) + rounded_odd = n_odd.to(GRAD_X.dtype.element_ty).to(tl.float32) + partial = GRAD_W_PARTIAL + (block * num_heads + head).to(tl.int64) * HEAD_DIM + tl.store(partial + 2 * pairs, tl.sum(gy_even * rounded_even, axis=0)) + tl.store(partial + 2 * pairs + 1, tl.sum(gy_odd * rounded_odd, axis=0)) + w_even = tl.load(W + 2 * pairs).to(tl.float32) + w_odd = tl.load(W + 2 * pairs + 1).to(tl.float32) + gn_even = gy_even * w_even[None, :] + gn_odd = gy_odd * w_odd[None, :] + projection = (tl.sum(gn_even * n_even, axis=1) + tl.sum(gn_odd * n_odd, axis=1)) / HEAD_DIM + gx_even = rstd[:, None] * (gn_even - n_even * projection[:, None]) + gx_odd = rstd[:, None] * (gn_odd - n_odd * projection[:, None]) + gx_base = GRAD_X + rows64[:, None] * (num_heads * HEAD_DIM) + head * HEAD_DIM + tl.store(gx_base + 2 * pairs[None, :], gx_even.to(GRAD_X.dtype.element_ty), mask=keep) + tl.store(gx_base + 2 * pairs[None, :] + 1, gx_odd.to(GRAD_X.dtype.element_ty), mask=keep) + + +def _rows_view(x: torch.Tensor) -> tuple[torch.Tensor, int, int]: + """``[B, S, H, D]`` (last dim contiguous) as row/head strides over ``B * S`` rows, copying only when it must.""" + if x.stride(-1) != 1 or x.stride(0) != x.shape[1] * x.stride(1): + x = x.contiguous() + return x, x.stride(1), x.stride(2) + + +def _check_rope(x: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None) -> None: + if (cos is None) != (sin is None): + raise ValueError("cos and sin go together") + if cos is not None and (cos.shape != (x.shape[1], x.shape[-1] // 2) or sin.shape != cos.shape): + raise ValueError(f"RoPE tables must be [seq_len, head_dim / 2] = [{x.shape[1]}, {x.shape[-1] // 2}], got " + f"{tuple(cos.shape)} and {tuple(sin.shape)}") + + +def _qk_forward_triton(x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None, + eps: float) -> tuple[torch.Tensor, torch.Tensor]: + batch, seq_len, num_heads, head_dim = x.shape + x, stride_r, stride_h = _rows_view(x) + out = torch.empty((batch, seq_len, num_heads, head_dim), device=x.device, dtype=x.dtype) + rstd = torch.empty((batch * seq_len, num_heads), device=x.device, dtype=torch.float32) + n_rows = batch * seq_len + has_rope = cos is not None + tables = (cos.float().contiguous(), sin.float().contiguous()) if has_rope else (rstd, rstd) + grid = (triton.cdiv(n_rows, _QK_BLOCK_ROWS), num_heads) + _qk_rmsnorm_rope_fwd_kernel[grid](x, weight.contiguous(), tables[0], tables[1], out, rstd, n_rows, seq_len, + num_heads, stride_r, stride_h, out.stride(1), out.stride(2), + tables[0].stride(0), eps, HEAD_DIM=head_dim, BLOCK_ROWS=_QK_BLOCK_ROWS, + HAS_ROPE=has_rope) + return out, rstd + + +def _qk_backward_triton(grad_out: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, + sin: torch.Tensor | None, rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch, seq_len, num_heads, head_dim = x.shape + x, stride_r, stride_h = _rows_view(x) + grad_out, stride_gr, stride_gh = _rows_view(grad_out) + n_rows = batch * seq_len + num_blocks = triton.cdiv(n_rows, _QK_BLOCK_ROWS) + grad_x = torch.empty((batch, seq_len, num_heads, head_dim), device=x.device, dtype=x.dtype) + partial = torch.empty((num_blocks * num_heads, head_dim), device=x.device, dtype=torch.float32) + has_rope = cos is not None + tables = (cos.float().contiguous(), sin.float().contiguous()) if has_rope else (rstd, rstd) + _qk_rmsnorm_rope_bwd_kernel[(num_blocks, num_heads)](x, weight.contiguous(), tables[0], tables[1], grad_out, rstd, + grad_x, partial, n_rows, seq_len, num_heads, stride_r, + stride_h, stride_gr, stride_gh, tables[0].stride(0), + HEAD_DIM=head_dim, BLOCK_ROWS=_QK_BLOCK_ROWS, + HAS_ROPE=has_rope) + return grad_x, partial.sum(0).to(weight.dtype) + + +def _compute_dtype(x: torch.Tensor) -> torch.dtype: + """FP32, the eager kernels' arithmetic, except for FP64 inputs (gradient checks).""" + return torch.float64 if x.dtype == torch.float64 else torch.float32 + + +def _qk_forward_reference(x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None, + eps: float) -> tuple[torch.Tensor, torch.Tensor]: + """The eager arithmetic: ``RMSNorm.forward_native``, ``.to(dtype)``, then ``_apply_rotary_emb`` (GPT-J pairs).""" + compute = _compute_dtype(x) + x32 = x.to(compute) + rstd = torch.rsqrt(x32.pow(2).mean(-1) + eps) + normalized = (x32 * rstd[..., None]).to(x.dtype) + y = (normalized.to(compute) * weight.to(compute)).to(x.dtype) + if cos is not None: + y_even, y_odd = y[..., ::2].to(compute), y[..., 1::2].to(compute) + cos, sin = cos.to(compute)[None, :, None, :], sin.to(compute)[None, :, None, :] + out_even = (y_even * cos - y_odd * sin).to(x.dtype) + out_odd = (y_odd * cos + y_even * sin).to(x.dtype) + y = torch.stack((out_even, out_odd), dim=-1).flatten(-2) + return y.contiguous(), rstd.reshape(-1, x.shape[2]) + + +def _qk_backward_reference(grad_out: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, + sin: torch.Tensor | None, rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """FP32 formulas of the Triton backward (rounding steps count as identities, as eager autograd treats casts).""" + compute = _compute_dtype(x) + x32 = x.to(compute) + rstd = rstd.reshape(x.shape[:3])[..., None].to(compute) + normalized = x32 * rstd + grad = grad_out.to(compute) + if cos is not None: + g_even, g_odd = grad[..., ::2], grad[..., 1::2] + cos, sin = cos.to(compute)[None, :, None, :], sin.to(compute)[None, :, None, :] + grad = torch.stack((g_even * cos + g_odd * sin, g_odd * cos - g_even * sin), dim=-1).flatten(-2) + grad_weight = (grad * normalized.to(x.dtype).to(compute)).sum(dim=(0, 1, 2)).to(weight.dtype) + grad_normalized = grad * weight.to(compute) + projection = (grad_normalized * normalized).mean(-1, keepdim=True) + grad_x = rstd * (grad_normalized - normalized * projection) + return grad_x.to(x.dtype), grad_weight + + +@torch.library.custom_op("fastvideo::pan2_qk_rmsnorm_rope", mutates_args=()) +def _qk_rmsnorm_rope_op(q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, + cos: torch.Tensor | None, sin: torch.Tensor | None, + eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + forward = _qk_forward_triton if _use_triton(q) else _qk_forward_reference + q_out, q_rstd = forward(q, q_weight, cos, sin, eps) + k_out, k_rstd = forward(k, k_weight, cos, sin, eps) + return q_out, k_out, q_rstd, k_rstd + + +@_qk_rmsnorm_rope_op.register_fake +def _qk_rmsnorm_rope_fake(q, k, q_weight, k_weight, cos, sin, eps): + del q_weight, k_weight, cos, sin, eps + q_rows = q.shape[0] * q.shape[1] + k_rows = k.shape[0] * k.shape[1] + return (q.new_empty(q.shape), k.new_empty(k.shape), q.new_empty((q_rows, q.shape[2]), dtype=_compute_dtype(q)), + k.new_empty((k_rows, k.shape[2]), dtype=_compute_dtype(k))) + + +def _qk_setup_context(ctx, inputs, output) -> None: + q, k, q_weight, k_weight, cos, sin, _eps = inputs + _q_out, _k_out, q_rstd, k_rstd = output + ctx.save_for_backward(q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd) + ctx.mark_non_differentiable(q_rstd, k_rstd) + + +@torch.library.custom_op("fastvideo::pan2_qk_rmsnorm_rope_backward", mutates_args=()) +def _qk_rmsnorm_rope_backward_op(grad_q_out: torch.Tensor, grad_k_out: torch.Tensor, q: torch.Tensor, k: torch.Tensor, + q_weight: torch.Tensor, k_weight: torch.Tensor, cos: torch.Tensor | None, + sin: torch.Tensor | None, q_rstd: torch.Tensor, + k_rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The backward as its own op: compiled graphs (AOTAutograd traces the autograd formula) hold one opaque node + instead of tracing into the Triton launch.""" + backward = _qk_backward_triton if _use_triton(q) else _qk_backward_reference + grad_q, grad_q_weight = backward(grad_q_out, q, q_weight, cos, sin, q_rstd) + grad_k, grad_k_weight = backward(grad_k_out, k, k_weight, cos, sin, k_rstd) + return grad_q, grad_k, grad_q_weight, grad_k_weight + + +@_qk_rmsnorm_rope_backward_op.register_fake +def _qk_rmsnorm_rope_backward_fake(grad_q_out, grad_k_out, q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd): + del grad_q_out, grad_k_out, cos, sin, q_rstd, k_rstd + return q.new_empty(q.shape), k.new_empty(k.shape), q_weight.new_empty(q_weight.shape), k_weight.new_empty( + k_weight.shape) + + +def _qk_backward(ctx, grad_q_out, grad_k_out, _grad_q_rstd, _grad_k_rstd): + q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd = ctx.saved_tensors + grad_q, grad_k, grad_q_weight, grad_k_weight = _qk_rmsnorm_rope_backward_op(grad_q_out, grad_k_out, q, k, + q_weight, k_weight, cos, sin, q_rstd, + k_rstd) + return grad_q, grad_k, grad_q_weight, grad_k_weight, None, None, None + + +_qk_rmsnorm_rope_op.register_autograd(_qk_backward, setup_context=_qk_setup_context) + + +def qk_rmsnorm_rope(q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + cos: torch.Tensor | None = None, + sin: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """Per-head RMSNorm of ``q`` and ``k`` (``[B, S, H, D]``, e.g. views of a fused projection), then RoPE. + + ``cos``/``sin`` are the half-dimension tables ``[S, D / 2]`` of ``get_1d_rotary_pos_embed(use_real=False)``, + applied to the interleaved ``(x[2i], x[2i + 1])`` pairs as ``_apply_rotary_emb(is_neox_style=False)`` does; without + them only the normalization runs. Returns contiguous ``[B, S, H, D]`` tensors in the input dtype. + """ + if q.dim() != 4 or k.shape[:2] != q.shape[:2] or q.shape[-1] % 2: + raise ValueError(f"q and k must be [B, S, H, D] with an even D, got {tuple(q.shape)} and {tuple(k.shape)}") + _check_rope(q, cos, sin) + q_out, k_out, _, _ = _qk_rmsnorm_rope_op(q, k, q_weight, k_weight, cos, sin, float(eps)) + return q_out, k_out + + +# -------------------------------------------------------------------------------------------------------------------- +# LayerNorm + adaptive modulation +# -------------------------------------------------------------------------------------------------------------------- + +_LN_BWD_ROWS_PER_PROGRAM = 16 + +if HAVE_TRITON: + + @triton.jit + def _layer_norm_modulate_fwd_kernel(X, SCALE, SHIFT, OUT, MEAN, RSTD, seq_len, n_cols, stride_xr, stride_or, + stride_scale_b, stride_shift_b, eps, BLOCK_C: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + batch = row // seq_len + cols = tl.arange(0, BLOCK_C) + valid = cols < n_cols + x = tl.load(X + row * stride_xr + cols, mask=valid, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / n_cols + centered = tl.where(valid, x - mean, 0.0) + rstd = tl.rsqrt(tl.sum(centered * centered, axis=0) / n_cols + eps) + scale = tl.load(SCALE + batch * stride_scale_b + cols, mask=valid, other=0.0).to(tl.float32) + shift = tl.load(SHIFT + batch * stride_shift_b + cols, mask=valid, other=0.0).to(tl.float32) + out = centered * rstd * scale + shift + tl.store(OUT + row * stride_or + cols, out.to(OUT.dtype.element_ty), mask=valid) + tl.store(MEAN + row, mean) + tl.store(RSTD + row, rstd) + + @triton.jit + def _layer_norm_modulate_bwd_kernel(X, SCALE, GRAD_OUT, MEAN, RSTD, GRAD_X, GRAD_SCALE_PARTIAL, + GRAD_SHIFT_PARTIAL, seq_len, n_cols, rows_per_program, num_chunks, stride_xr, + stride_gr, stride_scale_b, BLOCK_C: tl.constexpr): + chunk = tl.program_id(0) + batch = tl.program_id(1).to(tl.int64) + cols = tl.arange(0, BLOCK_C) + valid = cols < n_cols + scale = tl.load(SCALE + batch * stride_scale_b + cols, mask=valid, other=0.0).to(tl.float32) + grad_scale = tl.zeros([BLOCK_C], dtype=tl.float32) + grad_shift = tl.zeros([BLOCK_C], dtype=tl.float32) + first = chunk * rows_per_program + last = tl.minimum(first + rows_per_program, seq_len) + for offset in range(first, last): + row = batch * seq_len + offset + x = tl.load(X + row * stride_xr + cols, mask=valid, other=0.0).to(tl.float32) + grad = tl.load(GRAD_OUT + row * stride_gr + cols, mask=valid, other=0.0).to(tl.float32) + normalized = tl.where(valid, (x - tl.load(MEAN + row)) * tl.load(RSTD + row), 0.0) + grad_scale += grad * normalized + grad_shift += grad + grad_normalized = grad * scale + mean_grad = tl.sum(grad_normalized, axis=0) / n_cols + mean_projection = tl.sum(grad_normalized * normalized, axis=0) / n_cols + grad_x = tl.load(RSTD + row) * (grad_normalized - mean_grad - normalized * mean_projection) + tl.store(GRAD_X + row * n_cols + cols, grad_x.to(GRAD_X.dtype.element_ty), mask=valid) + partial = (batch * num_chunks + chunk) * n_cols + tl.store(GRAD_SCALE_PARTIAL + partial + cols, grad_scale, mask=valid) + tl.store(GRAD_SHIFT_PARTIAL + partial + cols, grad_shift, mask=valid) + + +def _batch_rows(table: torch.Tensor, batch: int, n_cols: int) -> torch.Tensor: + """A ``[B, C]`` (or broadcastable ``[C]``) table with a contiguous last dimension.""" + table = table.expand(batch, n_cols) + return table if table.stride(-1) == 1 else table.contiguous() + + +def _ln_forward_triton(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, + eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch, seq_len, n_cols = x.shape + rows = x.reshape(-1, n_cols) + if rows.stride(-1) != 1: + rows = rows.contiguous() + out = torch.empty((batch, seq_len, n_cols), device=x.device, dtype=x.dtype) + mean = torch.empty(batch * seq_len, device=x.device, dtype=torch.float32) + rstd = torch.empty_like(mean) + scale, shift = _batch_rows(scale, batch, n_cols), _batch_rows(shift, batch, n_cols) + block = triton.next_power_of_2(n_cols) + _layer_norm_modulate_fwd_kernel[(batch * seq_len, )](rows, scale, shift, out, mean, rstd, seq_len, n_cols, + rows.stride(0), n_cols, scale.stride(0), shift.stride(0), eps, + BLOCK_C=block, num_warps=8 if block >= 2048 else 4) + return out, mean, rstd + + +def _ln_backward_triton(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor, + rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch, seq_len, n_cols = x.shape + rows = x.reshape(-1, n_cols) + if rows.stride(-1) != 1: + rows = rows.contiguous() + grad_rows = grad_out.reshape(-1, n_cols) + if grad_rows.stride(-1) != 1: + grad_rows = grad_rows.contiguous() + scale_rows = _batch_rows(scale, batch, n_cols) + num_chunks = triton.cdiv(seq_len, _LN_BWD_ROWS_PER_PROGRAM) + grad_x = torch.empty((batch, seq_len, n_cols), device=x.device, dtype=x.dtype) + grad_scale = torch.empty((batch, num_chunks, n_cols), device=x.device, dtype=torch.float32) + grad_shift = torch.empty_like(grad_scale) + block = triton.next_power_of_2(n_cols) + _layer_norm_modulate_bwd_kernel[(num_chunks, batch)](rows, scale_rows, grad_rows, mean, rstd, grad_x, grad_scale, + grad_shift, seq_len, n_cols, _LN_BWD_ROWS_PER_PROGRAM, + num_chunks, rows.stride(0), grad_rows.stride(0), + scale_rows.stride(0), BLOCK_C=block, + num_warps=8 if block >= 2048 else 4) + return grad_x, grad_scale.sum(1), grad_shift.sum(1) + + +def _ln_forward_reference(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, + eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The autocast arithmetic: FP32 LayerNorm, ``* (1 + scale) + shift`` in FP32, one rounding at the end.""" + compute = _compute_dtype(x) + x32 = x.to(compute) + mean = x32.mean(-1) + rstd = torch.rsqrt(x32.var(-1, unbiased=False) + eps) + normalized = (x32 - mean[..., None]) * rstd[..., None] + batch, _, n_cols = x.shape + out = (normalized * _batch_rows(scale, batch, n_cols).to(compute)[:, None] + + _batch_rows(shift, batch, n_cols).to(compute)[:, None]) + return out.to(x.dtype), mean.reshape(-1), rstd.reshape(-1) + + +def _ln_backward_reference(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor, + rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + compute = _compute_dtype(x) + batch, seq_len, n_cols = x.shape + mean, rstd = mean.reshape(batch, seq_len, 1).to(compute), rstd.reshape(batch, seq_len, 1).to(compute) + normalized = (x.to(compute) - mean) * rstd + grad = grad_out.to(compute) + grad_normalized = grad * _batch_rows(scale, batch, n_cols).to(compute)[:, None] + grad_x = rstd * (grad_normalized - grad_normalized.mean(-1, keepdim=True) - + normalized * (grad_normalized * normalized).mean(-1, keepdim=True)) + return grad_x.to(x.dtype), (grad * normalized).sum(1), grad.sum(1) + + +@torch.library.custom_op("fastvideo::pan2_layer_norm_modulate", mutates_args=()) +def _layer_norm_modulate_op(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, + eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + forward = _ln_forward_triton if _use_triton(x) else _ln_forward_reference + return forward(x, scale, shift, eps) + + +@_layer_norm_modulate_op.register_fake +def _layer_norm_modulate_fake(x, scale, shift, eps): + del scale, shift, eps + rows = x.shape[0] * x.shape[1] + stats = _compute_dtype(x) + return x.new_empty(x.shape), x.new_empty(rows, dtype=stats), x.new_empty(rows, dtype=stats) + + +def _ln_setup_context(ctx, inputs, output) -> None: + x, scale, shift, _eps = inputs + _out, mean, rstd = output + ctx.save_for_backward(x, scale, mean, rstd) + ctx.scale_shape, ctx.shift_shape = scale.shape, shift.shape + ctx.scale_dtype, ctx.shift_dtype = scale.dtype, shift.dtype + ctx.mark_non_differentiable(mean, rstd) + + +@torch.library.custom_op("fastvideo::pan2_layer_norm_modulate_backward", mutates_args=()) +def _layer_norm_modulate_backward_op(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor, + rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(grad x, grad scale, grad shift), the tables' gradients per sample ``[B, C]``; an op of its own for the same + reason as ``_qk_rmsnorm_rope_backward_op``.""" + backward = _ln_backward_triton if _use_triton(x) else _ln_backward_reference + return backward(grad_out, x, scale, mean, rstd) + + +@_layer_norm_modulate_backward_op.register_fake +def _layer_norm_modulate_backward_fake(grad_out, x, scale, mean, rstd): + del grad_out, scale, mean, rstd + tables = (x.shape[0], x.shape[2]) + stats = _compute_dtype(x) + return x.new_empty(x.shape), x.new_empty(tables, dtype=stats), x.new_empty(tables, dtype=stats) + + +def _ln_backward(ctx, grad_out, _grad_mean, _grad_rstd): + x, scale, mean, rstd = ctx.saved_tensors + grad_x, grad_scale, grad_shift = _layer_norm_modulate_backward_op(grad_out, x, scale, mean, rstd) + # Tables shared by the batch ([C]) receive the batch's summed gradient. + grad_scale = grad_scale.sum_to_size(ctx.scale_shape).to(ctx.scale_dtype) + grad_shift = grad_shift.sum_to_size(ctx.shift_shape).to(ctx.shift_dtype) + return grad_x, grad_scale, grad_shift, None + + +_layer_norm_modulate_op.register_autograd(_ln_backward, setup_context=_ln_setup_context) + + +def layer_norm_modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float) -> torch.Tensor: + """``LayerNorm(x) * scale + shift`` for ``x`` ``[B, S, C]``, affine-free LayerNorm, in one pass. + + ``scale`` (PAN2's ``1 + scale``) and ``shift`` are ``[B, C]`` per sample or ``[C]`` shared. The arithmetic is the + autocast path's: an FP32 LayerNorm and FP32 modulation, rounded once to ``x``'s dtype. Under autocast that is + exactly the value the next linear layer would have cast its FP32 input to. + """ + if x.dim() != 3 or scale.shape[-1] != x.shape[-1] or shift.shape[-1] != x.shape[-1]: + raise ValueError(f"x must be [B, S, C] with [B, C] or [C] tables, got {tuple(x.shape)}, " + f"{tuple(scale.shape)} and {tuple(shift.shape)}") + out, _, _ = _layer_norm_modulate_op(x, scale, shift, float(eps)) + return out + + + +__all__ = [ + "HAVE_TRITON", + "fusions_available", + "layer_norm_modulate", + "qk_rmsnorm_rope", +] diff --git a/fastvideo/models/wan/vae.py b/fastvideo/models/wan/vae.py index e26cd8955..46a82b4b5 100644 --- a/fastvideo/models/wan/vae.py +++ b/fastvideo/models/wan/vae.py @@ -15,18 +15,61 @@ # limitations under the License. import contextvars +import functools +import importlib.util +import os +import sys from contextlib import contextmanager +from types import ModuleType import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange +import fastvideo.envs as envs +from fastvideo.distributed.parallel_state import GroupCoordinator, get_sp_group, model_parallel_is_initialized from fastvideo.layers.activation import get_act_fn +from fastvideo.logger import init_logger from fastvideo.models.vaes.common import (DiagonalGaussianDistribution, ParallelTiledVAE) from fastvideo.models.wan.vae_config import WanVAEConfig +logger = init_logger(__name__) + CACHE_T = 2 +# fastvideo-kernel/wan22_vae_decode decodes the Wan2.2-TI2V-5B VAE's latents [1, 48, T, 44, 78] (704 x 1248) on +# sm_100a (B200 / GB200) and sm_103a (B300). +_DECODE_KERNEL_LATENT = (48, 44, 78) +_DECODE_KERNEL_VIDEO = (704, 1248) +_DECODE_KERNEL_CAPABILITIES = ((10, 0), (10, 3)) +# The group sizes it splits a decode over (decode_kernel_parallel), each checked to decode the same video as one GPU. +_DECODE_KERNEL_WORLDS = (2, 4, 8) + + +@functools.cache +def _wan22_vae_decode() -> ModuleType: + """The wrapper of fastvideo-kernel/wan22_vae_decode (``FASTVIDEO_WAN22_VAE_DECODE_DIR``, by default this + checkout's), loaded once per process as ``wan22_vae_decode``, as its own tests import it. Its extension builds on + first use (into ``WAN22_VAE_DECODE_BUILD_DIR``, see the wrapper).""" + module = sys.modules.get("wan22_vae_decode") + if module is None: + root = envs.FASTVIDEO_WAN22_VAE_DECODE_DIR.get() or os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), + "fastvideo-kernel", "wan22_vae_decode") + path = os.path.join(root, "wan22_vae_decode.py") + if not os.path.isfile(path): + raise RuntimeError(f"decode_kernel needs fastvideo-kernel/wan22_vae_decode; {path} does not exist " + "(FASTVIDEO_WAN22_VAE_DECODE_DIR).") + spec = importlib.util.spec_from_file_location("wan22_vae_decode", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules["wan22_vae_decode"] = module + try: + spec.loader.exec_module(module) + except BaseException: + del sys.modules["wan22_vae_decode"] + raise + return module is_first_frame = contextvars.ContextVar("is_first_frame", default=False) feat_cache = contextvars.ContextVar("feat_cache", default=None) @@ -1096,6 +1139,10 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): ) self.use_feature_cache = config.use_feature_cache + # The kernel decoders of decode_kernel, one per device and process group (their weights prepared once; None + # where the group's ranks could not all build one), and why the kernel cannot run. + self._kernel_decoders: dict[tuple[torch.device, tuple[int, ...] | None], object | None] = {} + self._kernel_unavailable: str | None = None def clear_cache(self) -> None: @@ -1187,6 +1234,10 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): return enc def decode(self, z: torch.Tensor) -> torch.Tensor: + if getattr(self.config, "decode_kernel", False): + video = self._kernel_decode(z) + if video is not None: + return video if self.use_feature_cache: self.clear_cache() iter_ = z.shape[2] @@ -1215,6 +1266,80 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): return out + def _kernel_decode(self, z: torch.Tensor) -> torch.Tensor | None: + """``decode`` by fastvideo-kernel's Wan2.2 decoder (``decode_kernel``): bf16 with fp32 accumulation, within the + error of a bf16 decode of the fp32 one, or None where it does not apply (other devices, geometries, batches or + VAEs), and the decode below runs as before. Every rank of a multi-GPU run returns the whole video, as the + decode below does: decoded on its GPU alone, or with ``decode_kernel_parallel`` split over the + sequence-parallel group and broadcast from the group's first rank.""" + # Only what is the same on every rank decides here, so that the ranks of a group take the split together; a + # rank whose own build failed still joins the group's agreement (see _kernel_decoder). + reason = None + if not z.is_cuda: + reason = "the latents are not on a GPU" + elif torch.cuda.get_device_capability(z.device) not in _DECODE_KERNEL_CAPABILITIES: + major, minor = torch.cuda.get_device_capability(z.device) + reason = f"the GPU is sm_{major}{minor}, not sm_100 or sm_103" + elif z.dim() != 5 or z.shape[0] != 1 or (z.shape[1], *z.shape[3:]) != _DECODE_KERNEL_LATENT: + reason = f"latents {tuple(z.shape)} are not [1, 48, T, 44, 78] (a 704 x 1248 video, batch 1)" + if reason is not None: + logger.warning_once(f"decode_kernel: {reason}; decoding with AutoencoderKLWan.") + return None + group = self._kernel_group(z) + if group is not None and (decoder := self._kernel_decoder(z.device, group)) is not None: + logger.info_once(f"decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder over the " + f"{group.world_size} GPUs of the sequence-parallel group.") + video = decoder.decode(z) # the video on the group's first rank, None on the others + if video is None: + video = torch.empty((1, 3, 4 * z.shape[2] - 3, *_DECODE_KERNEL_VIDEO), dtype=torch.bfloat16, + device=z.device) + return group.broadcast(video, src=0).float() + decoder = self._kernel_decoder(z.device, None) + if decoder is None: + logger.warning_once(f"decode_kernel: {self._kernel_unavailable}; decoding with AutoencoderKLWan.") + return None + logger.info_once("decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder on this GPU.") + return decoder.decode(z).float() + + def _kernel_group(self, z: torch.Tensor) -> GroupCoordinator | None: + """The sequence-parallel group ``decode_kernel_parallel`` splits the decode of ``z`` over, or None (decode + alone): the switch is off, there is no group of a size the kernel splits over, or the video has fewer latent + frames than the group has GPUs, where the split's fixed cost (about 50 ms over 4 B300s, 100 ms over 8) exceeds + the decode.""" + if not getattr(self.config, "decode_kernel_parallel", False) or not model_parallel_is_initialized(): + return None + group = get_sp_group() + return group if group.world_size in _DECODE_KERNEL_WORLDS and z.shape[2] >= group.world_size else None + + def _kernel_decoder(self, device: torch.device, group: GroupCoordinator | None) -> object | None: + """The kernel's decoder on ``device``, alone (``group`` None) or as this rank of ``group``, built on first use + (its extension built or loaded, its weights prepared), or None where it cannot be. Over a group every rank + builds its own without communicating, then the ranks agree: the split decode needs every one of them, so a + rank that could not build makes the whole group decode alone.""" + key = (device, None if group is None else tuple(group.ranks)) + if key in self._kernel_decoders: + return self._kernel_decoders[key] + decoder = None + if self._kernel_unavailable is None: + try: + decoder = _wan22_vae_decode().WanVAEDecoder(self.state_dict(), device=device, + pg=None if group is None else group.device_group, + local=group is None) + except (KeyError, ValueError) as error: # not the Wan2.2-TI2V-5B VAE's decoder + self._kernel_unavailable = f"this VAE's decoder does not fit the kernel ({error})" + except Exception as error: # the kernel failed to load or build; generation goes on without it + self._kernel_unavailable = f"the kernel did not load ({type(error).__name__}: {error})" + if group is not None: + ready = torch.tensor([int(decoder is not None)], dtype=torch.int32) + torch.distributed.all_reduce(ready, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group) + if not ready.item(): + if decoder is not None: + logger.warning("decode_kernel_parallel: a rank of the sequence-parallel group could not build " + "the kernel's decoder; every rank decodes alone.") + decoder = None + self._kernel_decoders[key] = decoder + return decoder + def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: x = self.post_quant_conv(z) with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae): diff --git a/fastvideo/models/wan/vae_config.py b/fastvideo/models/wan/vae_config.py index 8c6e67955..89935b69d 100644 --- a/fastvideo/models/wan/vae_config.py +++ b/fastvideo/models/wan/vae_config.py @@ -79,6 +79,13 @@ class WanVAEConfig(VAEConfig): use_parallel_tiling: bool = False use_light_vae: bool = False + # Decode with fastvideo-kernel's Wan2.2 decoder (wan22_vae_decode: B200 / GB200 / B300, the Wan2.2-TI2V-5B VAE, + # 704 x 1248 video, batch 1) in bf16 where it applies; other inputs and devices take the decode here. + decode_kernel: bool = False + # With decode_kernel, split each decode over the sequence-parallel group (2, 4 or 8 GPUs) instead of decoding on + # every GPU alone. Every rank of the group must decode the same latents at the same time, as a pipeline's decoding + # stage does; the group's first rank assembles the video and broadcasts it, so every rank returns it as before. + decode_kernel_parallel: bool = False def __post_init__(self): self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2 diff --git a/fastvideo/pipelines/basic/pan2/pan2_pipeline.py b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py index f3f7dc3f5..b0c392cd0 100644 --- a/fastvideo/pipelines/basic/pan2/pan2_pipeline.py +++ b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py @@ -39,7 +39,17 @@ class PAN2Pipeline(ComposedPipelineBase): stage=PAN2DenoisingStage(transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), pipeline=self)) - self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) + # The performance switches of PAN2PipelineConfig: the DiT's fused kernels and the VAE's decode kernel. + pipeline_config = getattr(fastvideo_args, "pipeline_config", None) + transformer = self.get_module("transformer") + if hasattr(transformer, "set_fused_kernels"): + transformer.set_fused_kernels(bool(getattr(pipeline_config, "dit_fused_kernels", False))) + vae = self.get_module("vae") + vae_config = getattr(vae, "config", None) + if hasattr(vae_config, "decode_kernel"): + vae_config.decode_kernel = bool(getattr(pipeline_config, "vae_decode_kernel", False)) + vae_config.decode_kernel_parallel = bool(getattr(pipeline_config, "vae_decode_kernel_parallel", False)) + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=vae, pipeline=self)) class PAN2ModularPipeline(PAN2Pipeline): diff --git a/fastvideo/pipelines/basic/pan2/presets.py b/fastvideo/pipelines/basic/pan2/presets.py index 10409a1d5..af93fe8f0 100644 --- a/fastvideo/pipelines/basic/pan2/presets.py +++ b/fastvideo/pipelines/basic/pan2/presets.py @@ -10,7 +10,8 @@ _DENOISE_STAGE = PresetStageSpec( allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}), ) -# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. The preset sends no negative prompt (an empty one is encoded +# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 2.5 (the release launcher's default since ifm-ai/PAN2 +# 86d5023, 2026-10-23; 3 before). The preset sends no negative prompt (an empty one is encoded # as a single space); the reference sampler sends the release's structured negative caption, so pass it as the # request's ``negative_prompt`` to follow that recipe. _SHARED_DEFAULTS = { @@ -19,7 +20,7 @@ _SHARED_DEFAULTS = { "num_frames": 313, "fps": 24, "num_inference_steps": 50, - "guidance_scale": 3.0, + "guidance_scale": 2.5, "negative_prompt": "", } @@ -33,6 +34,7 @@ PAN2_T2V = InferencePreset( defaults=dict(_SHARED_DEFAULTS), ) +# The same pipeline runs image-to-video when the request carries a first frame (image_path). PAN2_I2V = InferencePreset( name="pan2_i2v", version=1, diff --git a/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py index 6344fafb8..04dd61018 100644 --- a/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py +++ b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""``_quantize_video_frames_to_uint8`` and its CUDA-OOM fallback. +"""``_quantize_video_frames_to_uint8``, ``_video_frames`` and their CUDA-OOM fallback. Found on a real GPU run of the K6 T2IVA-distilled example (30B params, ``dit_cpu_offload=False``): the worker process keeps the transformer resident on the device after denoising and decode, so the client process's on-device quantize @@ -75,3 +75,54 @@ def test_a_non_oom_runtime_error_is_not_swallowed(sample_video, monkeypatch): monkeypatch.setattr(torch.Tensor, "clamp_", always_fails) with pytest.raises(RuntimeError, match="unrelated failure"): _quantize_video_frames_to_uint8(sample_video.clone()) + + +def _frames_by_grid(src: torch.Tensor) -> list: + """The frames as the generator built them before `_video_frames`: make_grid and a channels-last copy per frame.""" + import torchvision + + return [ + torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() + for x in _quantize_video_frames_to_uint8(src) + ] + + +@pytest.mark.parametrize("shape", [(1, 3, 4, 5, 6), (1, 1, 4, 5, 6), (2, 3, 4, 5, 6)]) +def test_video_frames_match_the_per_frame_grid(shape): + import numpy as np + + from fastvideo.entrypoints.video_generator import _video_frames + + torch.manual_seed(1) + src = torch.rand(*shape) * 1.3 - 0.15 + frames, expected = _video_frames(src.clone()), _frames_by_grid(src.clone()) + assert len(frames) == len(expected) == shape[2] + for frame, reference in zip(frames, expected): + assert frame.dtype == np.uint8 and frame.shape == reference.shape + np.testing.assert_array_equal(frame, reference) + + +def test_video_frames_fall_back_to_the_cpu_on_cuda_oom(monkeypatch, caplog): + import numpy as np + + from fastvideo.entrypoints.video_generator import _video_frames + + torch.manual_seed(2) + src = torch.rand(1, 3, 4, 5, 6) + expected = _frames_by_grid(src.clone()) + calls = {"n": 0} + real_clamp_ = torch.Tensor.clamp_ + + def flaky_clamp_(self, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise torch.cuda.OutOfMemoryError("simulated") + return real_clamp_(self, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "clamp_", flaky_clamp_) + with caplog.at_level(logging.WARNING): + frames = _video_frames(src.clone()) + assert calls["n"] == 2 + for frame, reference in zip(frames, expected): + np.testing.assert_array_equal(frame, reference) + assert any("out of memory" in message.lower() for message in caplog.messages) diff --git a/fastvideo/tests/stages/test_pan2_latent_preparation.py b/fastvideo/tests/stages/test_pan2_latent_preparation.py index 80c09dd4d..cc4515245 100644 --- a/fastvideo/tests/stages/test_pan2_latent_preparation.py +++ b/fastvideo/tests/stages/test_pan2_latent_preparation.py @@ -65,3 +65,26 @@ def test_pan2_pipeline_requests_fp32_latents(monkeypatch): pipeline.create_pipeline_stages(None) assert any(kwargs.get("latent_dtype") is torch.float32 for kwargs in seen["kwargs"]) + + +def test_pan2_pipeline_applies_its_performance_switches(monkeypatch): + for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage", + "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage", "LatentPreparationStage"): + monkeypatch.setattr(pan2_pipeline, name, lambda *args, **kwargs: None) + + class _Transformer: + fused = None + + def set_fused_kernels(self, enabled): + self.fused = enabled + + modules = {"transformer": _Transformer(), "vae": types.SimpleNamespace(config=types.SimpleNamespace(decode_kernel=False))} + pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline) + pipeline.add_stage = lambda stage_name, stage: None + pipeline.get_module = modules.get + config = types.SimpleNamespace(dit_fused_kernels=True, vae_decode_kernel=True, vae_decode_kernel_parallel=False) + + pipeline.create_pipeline_stages(types.SimpleNamespace(pipeline_config=config)) + + assert modules["transformer"].fused is True + assert modules["vae"].config.decode_kernel is True and modules["vae"].config.decode_kernel_parallel is False diff --git a/fastvideo/tests/transformers/test_pan2_fusions.py b/fastvideo/tests/transformers/test_pan2_fusions.py new file mode 100644 index 000000000..49f48e393 --- /dev/null +++ b/fastvideo/tests/transformers/test_pan2_fusions.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2's fused Q/K RMSNorm + RoPE and LayerNorm + modulation ops (``fastvideo.models.dits.pan2_fusions``). + +The PyTorch references must reproduce the eager arithmetic they replace, their backward formulas must be the true +gradients (fp64 gradcheck), the Triton kernels must reproduce the references, and a PAN2 model with fused kernels must +match the eager model forward and backward. On a GPU-less host the Triton kernels run under the Triton interpreter in +a subprocess, in FP32: the interpreter truncates FP32 to BF16 instead of rounding to nearest even, so the BF16 kernel +comparison needs a GPU. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig, PAN2VideoConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.rotary_embedding import _apply_rotary_emb, get_1d_rotary_pos_embed +from fastvideo.models.dits import pan2 as pan2_module +from fastvideo.models.dits import pan2_fusions +from fastvideo.models.dits.pan2_fusions import layer_norm_modulate, qk_rmsnorm_rope +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +REPO_ROOT = Path(__file__).resolve().parents[3] + +LATENT_CHANNELS = 4 +TINY_ARCH = dict( + patch_size=(1, 2, 2), + in_channels=2 * LATENT_CHANNELS + 1, + out_channels=LATENT_CHANNELS, + num_attention_heads=4, + attention_head_dim=16, + num_layers=2, + num_refiner_layers=1, + mlp_ratio=2.0, + context_mlp_ratio=2.0, + refiner_mlp_ratio=2.0, + text_embed_dim=16, + rope_axes_dim=(4, 6, 6), +) + + +@pytest.fixture() +def cpu_attention(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + + +@pytest.fixture() +def single_process_group(env_overrides): + """A 1x1 sequence/tensor-parallel group, which DistributedAttention needs even on one CPU process.""" + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + + env_overrides.enter_context(envs.override_external("MASTER_ADDR", os.environ.get("MASTER_ADDR") or "127.0.0.1")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", os.environ.get("MASTER_PORT") or "29593")) + env_overrides.enter_context(envs.override_external("RANK", "0")) + env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1")) + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + yield + finally: + cleanup_dist_env_and_memory() + + +def _model() -> pan2_module.PAN2Transformer3DModel: + return pan2_module.PAN2Transformer3DModel(PAN2VideoConfig(arch_config=PAN2ArchConfig(**TINY_ARCH)), hf_config={}) + + +def _inputs(seed: int = 0) -> dict[str, torch.Tensor]: + generator = torch.Generator().manual_seed(seed) + return dict( + hidden_states=torch.randn(1, TINY_ARCH["in_channels"], 3, 4, 6, generator=generator), + encoder_hidden_states=torch.randn(1, 5, TINY_ARCH["text_embed_dim"], generator=generator), + timestep=torch.tensor([600.0]), + ) + + +def _randomize(model: torch.nn.Module, seed: int) -> None: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1) + + + +def _qkv(batch: int, seq: int, heads: int, dim: int, dtype: torch.dtype, seed: int = 0): + generator = torch.Generator().manual_seed(seed) + qkv = torch.randn(batch, seq, 3 * heads * dim, generator=generator).to(dtype) + return qkv.unflatten(-1, (3, heads, -1)).unbind(2) + + +def _rope_tables(seq: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: + positions = torch.arange(seq, dtype=torch.float32) * (16 / 24) + return get_1d_rotary_pos_embed(dim, positions, theta=256.0, use_real=False) + + +def _eager_qk(x: torch.Tensor, weight: torch.Tensor, eps: float, rope: tuple[torch.Tensor, torch.Tensor] | None): + norm = RMSNorm(x.shape[-1], eps=eps) + with torch.no_grad(): + norm.weight.copy_(weight) + y = norm.forward_native(x).to(x.dtype) + if rope is None: + return y + rows = y.flatten(0, 1) + cos, sin = (table.repeat(x.shape[0], 1) for table in rope) + return _apply_rotary_emb(rows, cos, sin, is_neox_style=False).unflatten(0, x.shape[:2]) + + +# ---- references against the eager arithmetic ----------------------------------------------------------------------- + + +@pytest.mark.parametrize("with_rope", [False, True]) +def test_qk_reference_is_the_eager_chain_bit_for_bit(with_rope: bool) -> None: + q, k, _ = _qkv(2, 13, 3, 16, torch.bfloat16) + weights = torch.rand(16) + 0.5, torch.rand(16) + 0.5 + rope = _rope_tables(13, 16) if with_rope else None + fused_q, fused_k = qk_rmsnorm_rope(q, k, *weights, 1e-6, *(rope or (None, None))) + assert torch.equal(fused_q, _eager_qk(q, weights[0], 1e-6, rope)) + assert torch.equal(fused_k, _eager_qk(k, weights[1], 1e-6, rope)) + assert fused_q.is_contiguous() and fused_q.dtype == torch.bfloat16 + + +def test_layer_norm_modulate_reference_is_the_autocast_arithmetic() -> None: + generator = torch.Generator().manual_seed(1) + x = (torch.randn(2, 11, 48, generator=generator) * 4 + 2).to(torch.bfloat16) + scale = (1 + torch.randn(2, 48, generator=generator)).to(torch.bfloat16) + shift = torch.randn(48, generator=generator).to(torch.bfloat16) + norm = torch.nn.LayerNorm(48, elementwise_affine=False, eps=1e-6) + # Under autocast the eager LayerNorm and modulation run in FP32 and the next linear rounds the result once. + expected = (norm(x.float()) * scale.float()[:, None] + shift.float()).to(torch.bfloat16) + actual = layer_norm_modulate(x, scale, shift, 1e-6) + assert actual.dtype == torch.bfloat16 + torch.testing.assert_close(actual, expected, atol=0, rtol=2**-7) + + +def test_fused_ops_have_true_gradients() -> None: + q, k, _ = _qkv(2, 5, 2, 8, torch.float64, seed=2) + q, k = q.detach().requires_grad_(), k.detach().requires_grad_() + weights = [(torch.rand(8, dtype=torch.float64) + 0.5).requires_grad_() for _ in range(2)] + cos, sin = (table.double() for table in _rope_tables(5, 8)) + assert torch.autograd.gradcheck(lambda *args: qk_rmsnorm_rope(*args, 1e-6, cos, sin), (q, k, *weights)) + assert torch.autograd.gradcheck(lambda *args: qk_rmsnorm_rope(*args, 1e-6), (q, k, *weights)) + + x = torch.randn(2, 4, 12, dtype=torch.float64, requires_grad=True) + scale = torch.randn(2, 12, dtype=torch.float64, requires_grad=True) + shift = torch.randn(12, dtype=torch.float64, requires_grad=True) # shared by the batch, as PAN2's text tables + assert torch.autograd.gradcheck(lambda *args: layer_norm_modulate(*args, 1e-6), (x, scale, shift)) + + +@pytest.mark.gpu +def test_fused_ops_compile_with_their_backwards_on_gpu() -> None: + """Inside a fullgraph compiled region, as regional compile captures PAN2's blocks, both fused ops and their + backwards are opaque ops: tracing never launches a Triton kernel, and the values and gradients are eager's.""" + if not torch.cuda.is_available() or not pan2_fusions.HAVE_TRITON: + pytest.skip("needs CUDA and Triton") + q, k, _ = (t.cuda() for t in _qkv(2, 21, 3, 16, torch.bfloat16, seed=6)) + weights = [(torch.rand(16) + 0.5).cuda() for _ in range(2)] + cos, sin = (table.cuda() for table in _rope_tables(21, 16)) + generator = torch.Generator().manual_seed(7) + x = torch.randn(2, 9, 40, generator=generator).cuda().to(torch.bfloat16) + scale = (1 + torch.randn(2, 40, generator=generator)).cuda() + shift = torch.randn(40, generator=generator).cuda() + + def fn(q, k, q_weight, k_weight, x, scale, shift): + q_out, k_out = qk_rmsnorm_rope(q, k, q_weight, k_weight, 1e-6, cos, sin) + return (q_out.float().square().sum() + k_out.float().square().sum() + + layer_norm_modulate(x, scale, shift, 1e-6).float().square().sum()) + + def run(function): + leaves = [t.detach().clone().requires_grad_() for t in (q, k, *weights, x, scale, shift)] + loss = function(*leaves) + loss.backward() + return loss.detach(), [leaf.grad for leaf in leaves] + + eager_loss, eager_grads = run(fn) + loss, grads = run(torch.compile(fn, fullgraph=True, dynamic=True)) + torch.testing.assert_close(loss, eager_loss, atol=0, rtol=1e-5) + for actual, expected in zip(grads, eager_grads, strict=True): + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + +# ---- the Triton kernels against the references ---------------------------------------------------------------------- + + +def check_triton_kernels(device: str, dtypes: tuple[torch.dtype, ...]) -> None: + """Every Triton kernel against its reference, forward and backward; raises on a mismatch.""" + for dtype in dtypes: + tol = dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else dict(atol=0, rtol=2**-6) + # 21 rows: a partial row block. q and k are strided views of one projection, as in the model. + q, k, _ = (t.to(device) for t in _qkv(2, 21, 3, 16, dtype, seed=3)) + weight = (torch.rand(16) + 0.5).to(device) + cos, sin = (table.to(device) for table in _rope_tables(21, 16)) + grad = torch.randn(2, 21, 3, 16, generator=torch.Generator().manual_seed(4)).to(device=device, dtype=dtype) + for tables in ((cos, sin), (None, None)): + expected, expected_rstd = pan2_fusions._qk_forward_reference(q, weight, *tables, 1e-6) + actual, rstd = pan2_fusions._qk_forward_triton(q, weight, *tables, 1e-6) + torch.testing.assert_close(actual.float(), expected.float(), **tol) + torch.testing.assert_close(rstd, expected_rstd, atol=1e-6, rtol=1e-6) + expected_grads = pan2_fusions._qk_backward_reference(grad, q, weight, *tables, expected_rstd) + actual_grads = pan2_fusions._qk_backward_triton(grad, q, weight, *tables, rstd) + for a, e in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(a.float(), e.float(), atol=1e-4, rtol=1e-3 if dtype != torch.float32 else 1e-5) + + generator = torch.Generator().manual_seed(5) + x = (torch.randn(2, 9, 40, generator=generator) * 3 + 1).to(device=device, dtype=dtype) + scale = (1 + torch.randn(2, 40, generator=generator)).to(device=device, dtype=dtype) + shift = torch.randn(40, generator=generator).to(device=device, dtype=dtype) + grad = torch.randn(2, 9, 40, generator=generator).to(device=device, dtype=dtype) + expected = pan2_fusions._ln_forward_reference(x, scale, shift, 1e-6) + actual = pan2_fusions._ln_forward_triton(x, scale, shift, 1e-6) + torch.testing.assert_close(actual[0].float(), expected[0].float(), **tol) + torch.testing.assert_close(actual[1:], expected[1:], atol=1e-5, rtol=1e-5) + expected_grads = pan2_fusions._ln_backward_reference(grad, x, scale, expected[1], expected[2]) + actual_grads = pan2_fusions._ln_backward_triton(grad, x, scale, actual[1], actual[2]) + for a, e in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(a.float(), e.float(), atol=1e-4, rtol=1e-3 if dtype != torch.float32 else 1e-5) + + +def test_triton_kernels_match_their_references(env_overrides) -> None: + if not pan2_fusions.HAVE_TRITON: + pytest.skip("Triton is not installed") + if torch.cuda.is_available(): + check_triton_kernels("cuda", (torch.float32, torch.bfloat16)) + return + code = (f"import sys; sys.path.insert(0, {str(REPO_ROOT)!r}); " + "import fastvideo.models.dits.pan2_fusions as fusions; fusions._FORCE_TRITON = True; " + "import torch; from fastvideo.tests.transformers.test_pan2_fusions import check_triton_kernels; " + "check_triton_kernels('cpu', (torch.float32, ))") + # The child inherits this environment with the Triton interpreter switched on. + env_overrides.enter_context(envs.override_external("TRITON_INTERPRET", "1")) + process = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600) + assert process.returncode == 0, f"Triton interpreter check failed:\n{process.stdout}\n{process.stderr}" + + +# ---- the fused model against the eager model ----------------------------------------------------------------------- + + +def _forward_backward(model: torch.nn.Module, inputs: dict[str, torch.Tensor]): + model.zero_grad(set_to_none=True) + with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + output = model(**inputs) + weight = torch.randn(output.shape, generator=torch.Generator().manual_seed(6), dtype=output.dtype) + (output * weight).sum().backward() + return output.detach(), {name: p.grad.detach().clone() for name, p in model.named_parameters() if p.grad is not None} + + +def test_fused_model_matches_the_eager_model(cpu_attention, single_process_group) -> None: + model = _model() + _randomize(model, seed=7) + inputs = _inputs(seed=8) + eager_out, eager_grads = _forward_backward(model, inputs) + model.set_fused_kernels(True) + assert all(block.fused_kernels and block.attn.fused_kernels for block in model.transformer_blocks) + fused_out, fused_grads = _forward_backward(model, inputs) + torch.testing.assert_close(fused_out, eager_out, atol=1e-5, rtol=1e-5) + assert fused_grads.keys() == eager_grads.keys() + for name, grad in eager_grads.items(): + torch.testing.assert_close(fused_grads[name], grad, atol=1e-5, rtol=1e-4, msg=name) + + +@pytest.mark.gpu +def test_fused_model_matches_the_eager_model_in_bf16_on_gpu(cpu_attention, single_process_group) -> None: + """The sampling numerics: BF16 activations under autocast, Triton kernels, on one GPU.""" + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + model = _model().cuda() + _randomize(model, seed=7) + inputs = {name: value.cuda() for name, value in _inputs(seed=8).items()} + + def run(): + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + output = model(**inputs) + output.float().square().mean().backward() + return output.float(), {name: p.grad.float() for name, p in model.named_parameters() if p.grad is not None} + + eager_out, eager_grads = run() + model.set_fused_kernels(True) + fused_out, fused_grads = run() + scale = eager_out.abs().mean() + assert float((fused_out - eager_out).abs().mean() / scale) < 1e-2 + for name, grad in eager_grads.items(): + error = float((fused_grads[name] - grad).norm() / grad.norm().clamp_min(1e-12)) + assert error < 5e-2, f"{name}: relative gradient error {error:.3e}" diff --git a/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py b/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py new file mode 100644 index 000000000..a279f0163 --- /dev/null +++ b/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``AutoencoderKLWan.decode`` through fastvideo-kernel's Wan2.2 decoder (``WanVAEConfig.decode_kernel``). + +CPU tests: the switch is off unless asked for (PAN2 sampling asks), and a decode it cannot take runs the module's own +decode without loading the kernel. GPU tests (sm_100 / sm_103, the PAN2 VAE's weights in +``FASTVIDEO_TEST_PAN2_MODEL_DIR``): the kernel's decode against the module's fp32 decode, within three times the error +of a bf16 decode (the kernel's own rule); another geometry decoding exactly as before; two ranks decoding on their own +GPUs with no collective. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig +from fastvideo.models.wan import vae as wan_vae +from fastvideo.models.wan.vae_config import WanVAEArchConfig, WanVAEConfig + + +def _tiny_vae(decode_kernel: bool) -> wan_vae.AutoencoderKLWan: + torch.manual_seed(0) + arch = WanVAEArchConfig(base_dim=8, z_dim=4, dim_mult=(1, 1), num_res_blocks=1, temperal_downsample=(False, ), + latents_mean=(0.0, ) * 4, latents_std=(1.0, ) * 4) + config = WanVAEConfig(arch_config=arch) + config.load_encoder, config.load_decoder, config.decode_kernel = False, True, decode_kernel + return wan_vae.AutoencoderKLWan(config).eval() + + +def test_the_kernel_is_opt_in_and_pan2_sampling_opts_in() -> None: + assert WanVAEConfig().decode_kernel is False and WanVAEConfig().decode_kernel_parallel is False + pan2 = PAN2PipelineConfig() + assert pan2.vae_decode_kernel is True and pan2.vae_decode_kernel_parallel is True and pan2.dit_fused_kernels + + +def test_a_decode_the_kernel_cannot_take_runs_the_module(monkeypatch) -> None: + def no_kernel(): + raise AssertionError("a CPU decode must not load the kernel") + + monkeypatch.setattr(wan_vae, "_wan22_vae_decode", no_kernel) + z = torch.randn(1, 4, 2, 4, 4) + with torch.no_grad(): + plain = _tiny_vae(False).decode(z) + routed = _tiny_vae(True).decode(z) + assert torch.equal(plain, routed) + + +_ON_KERNEL_GPU = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="needs an sm_100 / sm_103 GPU") + + +@pytest.fixture +def pan2_vae_dir() -> str: + root = envs.FASTVIDEO_TEST_PAN2_MODEL_DIR.get() + if not root: + pytest.skip("set FASTVIDEO_TEST_PAN2_MODEL_DIR to a PAN2 root with vae/") + return os.path.join(root, "vae") + + +def _pan2_vae(vae_dir: str) -> wan_vae.AutoencoderKLWan: + from fastvideo.models.loader.component_loader import VAELoader + + config = WanVAEConfig() + config.load_encoder, config.load_decoder = False, True + loader_args = SimpleNamespace(model_paths={}, vae_cpu_offload=False, + pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32")) + return VAELoader().load(vae_dir, loader_args).to("cuda").eval() + + +@_ON_KERNEL_GPU +def test_the_kernel_decodes_pan2_latents_within_the_bf16_error(pan2_vae_dir) -> None: + vae = _pan2_vae(pan2_vae_dir) + z = torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(3)).cuda() + with torch.no_grad(): + vae.config.decode_kernel = False + reference = vae.decode(z) + with torch.autocast("cuda", dtype=torch.bfloat16): + bf16 = vae.decode(z) + vae.config.decode_kernel = True + kernel = vae.decode(z) + assert list(vae._kernel_decoders) == [(z.device, None)], "the kernel path was not taken" + assert kernel.shape == reference.shape == (1, 3, 9, 704, 1248) and kernel.dtype == torch.float32 + assert float(kernel.abs().max()) <= 1.0 + d, d_bf16 = kernel - reference, bf16.float() - reference + assert float(d.abs().max()) <= 3 * float(d_bf16.abs().max()) + assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_bf16.pow(2).mean().sqrt()) + + +@_ON_KERNEL_GPU +def test_another_geometry_decodes_as_before_on_the_gpu(pan2_vae_dir) -> None: + vae = _pan2_vae(pan2_vae_dir) + z = torch.randn((1, 48, 2, 32, 32), generator=torch.Generator().manual_seed(4)).cuda() + with torch.no_grad(): + vae.config.decode_kernel = False + plain = vae.decode(z) + vae.config.decode_kernel = True + routed = vae.decode(z) + assert torch.equal(plain, routed) and not vae._kernel_decoders + + +def _decode_alone(rank: int, store_path: str, vae_dir: str, z: torch.Tensor, out_path: str) -> None: + import torch.distributed as dist + from safetensors.torch import load_file + + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=f"file://{store_path}", rank=rank, world_size=2) + try: + weights = load_file(os.path.join(vae_dir, "diffusion_pytorch_model.safetensors")) + decoder = wan_vae._wan22_vae_decode().WanVAEDecoder(weights, device=f"cuda:{rank}", local=True) + video = decoder.decode(z) + torch.save(video.cpu(), f"{out_path}.{rank}") + dist.barrier(device_ids=[rank]) + finally: + dist.destroy_process_group() + + +@_ON_KERNEL_GPU +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs") +def test_two_ranks_decode_on_their_own_gpus(tmp_path, pan2_vae_dir) -> None: + import torch.multiprocessing as mp + + # Build the extension here first, so the two ranks load it rather than both compiling into one directory. + wan_vae._wan22_vae_decode()._module() + z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(5)) + out = str(tmp_path / "video") + mp.start_processes(_decode_alone, args=(str(tmp_path / "store"), pan2_vae_dir, z, out), nprocs=2, join=True, + start_method="spawn") + first, second = torch.load(f"{out}.0"), torch.load(f"{out}.1") + assert first.shape == (1, 3, 5, 704, 1248) and torch.equal(first, second) diff --git a/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py b/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py new file mode 100644 index 000000000..4e2252f59 --- /dev/null +++ b/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``AutoencoderKLWan.decode`` through fastvideo-kernel's Wan2.2 decoder split over the sequence-parallel group +(``WanVAEConfig.decode_kernel_parallel``), as PAN2 sampling decodes on several GPUs. + +Needs a multi-GPU torchrun launch (real NCCL collectives over the group), sm_100 / sm_103 GPUs and the PAN2 VAE's +weights in ``FASTVIDEO_TEST_PAN2_MODEL_DIR``; skipped otherwise: + + FASTVIDEO_TEST_PAN2_MODEL_DIR= torchrun --nproc-per-node=4 -m pytest \\ + fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py -q + +Every rank must return the same video, the group's first rank's, within the kernel's own error rule against the +module's fp32 decode (three times the error of a bf16 decode, as the one-GPU test); a video with fewer latent frames +than GPUs decodes on each GPU alone, and a rank that cannot build the kernel makes the group decode alone without +hanging the others. Over two GPUs the split decode is +bit for bit the video one GPU decodes; over four and eight the bands of rows round differently, at the same error. +""" + +import os +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.models.wan.vae_config import WanVAEConfig + +_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="needs sm_100 / sm_103 GPUs"), + pytest.mark.skipif(_WORLD_SIZE < 2, reason="needs a torchrun launch with WORLD_SIZE > 1"), +] + + +@pytest.fixture(scope="module") +def sp_group(): + from fastvideo.distributed import get_sp_group, maybe_init_distributed_environment_and_model_parallel + maybe_init_distributed_environment_and_model_parallel(1, _WORLD_SIZE) + return get_sp_group() + + +@pytest.fixture(scope="module") +def pan2_vae_dir() -> str: + root = envs.FASTVIDEO_TEST_PAN2_MODEL_DIR.get() + if not root: + pytest.skip("set FASTVIDEO_TEST_PAN2_MODEL_DIR to a PAN2 root with vae/") + return os.path.join(root, "vae") + + +def _load_pan2_vae(vae_dir: str): + from fastvideo.models.loader.component_loader import VAELoader + + config = WanVAEConfig() + config.load_encoder, config.load_decoder, config.decode_kernel = False, True, True + loader_args = SimpleNamespace(model_paths={}, vae_cpu_offload=False, + pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32")) + return VAELoader().load(vae_dir, loader_args).to("cuda").eval() + + +@pytest.fixture(scope="module") +def pan2_vae(pan2_vae_dir): + return _load_pan2_vae(pan2_vae_dir) + + +@pytest.mark.parametrize("latent_frames", (1, 2, 13)) +@torch.no_grad() +def test_every_rank_returns_the_split_decode(sp_group, pan2_vae, latent_frames: int) -> None: + z = torch.randn((1, 48, latent_frames, 44, 78), generator=torch.Generator().manual_seed(latent_frames)).cuda() + pan2_vae.config.decode_kernel = False + reference = pan2_vae.decode(z) + with torch.autocast("cuda", dtype=torch.bfloat16): + bf16 = pan2_vae.decode(z) + pan2_vae.config.decode_kernel, pan2_vae.config.decode_kernel_parallel = True, False + alone = pan2_vae.decode(z) + pan2_vae.config.decode_kernel_parallel = True + split = pan2_vae.decode(z) + if latent_frames < sp_group.world_size: # fewer latent frames than GPUs: each GPU decodes alone + assert torch.equal(split, alone) + return + assert pan2_vae._kernel_decoders.get((z.device, tuple(sp_group.ranks))) is not None, "the split was not taken" + assert split.shape == reference.shape == (1, 3, 4 * latent_frames - 3, 704, 1248) and split.dtype == torch.float32 + assert torch.equal(pan2_vae.decode(z), split), "a second split decode (the group's decoder reused) differs" + + checksum = split.double().sum().reshape(1) + low, high = checksum.clone(), checksum.clone() + torch.distributed.all_reduce(low, op=torch.distributed.ReduceOp.MIN, group=sp_group.device_group) + torch.distributed.all_reduce(high, op=torch.distributed.ReduceOp.MAX, group=sp_group.device_group) + assert torch.equal(low, high), "the ranks hold different videos" + + d, d_bf16 = split - reference, bf16.float() - reference + assert float(d.abs().max()) <= 3 * float(d_bf16.abs().max()) + assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_bf16.pow(2).mean().sqrt()) + if sp_group.world_size == 2: + assert torch.equal(split, alone) + + +@torch.no_grad() +def test_a_rank_without_the_kernel_makes_the_group_decode_alone(sp_group, pan2_vae_dir) -> None: + """Group rank 1 could not build the kernel (here: an earlier failure on that rank alone). Every rank still joins the + group's one agreement, so none waits on it forever, and they all decode alone: rank 1 with AutoencoderKLWan, the + others with the kernel on their own GPU.""" + vae = _load_pan2_vae(pan2_vae_dir) + vae.config.decode_kernel_parallel = True + if sp_group.rank_in_group == 1: + vae._kernel_unavailable = "simulated: this rank could not build the kernel" + z = torch.randn((1, 48, 13, 44, 78), generator=torch.Generator().manual_seed(7)).cuda() + for _ in range(2): # the second decode reads the group's cached verdict without communicating + video = vae.decode(z) + assert video.shape == (1, 3, 49, 704, 1248) and video.dtype == torch.float32 + assert vae._kernel_decoders[(z.device, tuple(sp_group.ranks))] is None + assert (vae._kernel_decoders[(z.device, None)] is None) == (sp_group.rank_in_group == 1)