[perf] PAN2 sampling: Wan2.2 VAE decode kernel, split decode, fused DiT kernels
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
This commit is contained in:
William Lin
@@ -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`. |
|
||||
|
||||
|
||||
+15
-1
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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_<source>.pt`), the weights in `weight_names()` order
|
||||
(`decoder_weights_bf16.pt`, `weight_names.json`), the fp32 decodes (`ref_<clip>.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.
|
||||
@@ -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).
|
||||
@@ -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)
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#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 <class ElementD, int M, int N, int CL>
|
||||
static void gemm1x1(const at::Tensor& x, const at::Tensor& w, const at::Tensor& y) {
|
||||
using Tile = Shape<Int<M>, Int<N>, Shape<_64>>;
|
||||
using Cluster = Shape<Int<CL>, _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<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
|
||||
cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
|
||||
using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
|
||||
using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
|
||||
using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
|
||||
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<cute::rank<0>(sc)>{}, [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
|
||||
cute::for_each(cute::make_seq<cute::rank<0>(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<<<dim3((HW + 31) / 32, (C + 31) / 32, T), dim3(32, 32), 0, stream>>>(
|
||||
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<float, 256, 256, 2>(x, w, y);
|
||||
}
|
||||
|
||||
float* s_p = (float*)s.data_ptr();
|
||||
float scale = 1.f / sqrtf((float)C);
|
||||
row_softmax_kernel<<<T * HW, 256, HW * sizeof(float), stream>>>(
|
||||
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<ElementBF16, 256, 256, 2>(x, w, y);
|
||||
}
|
||||
return o;
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
#pragma once
|
||||
#include <ATen/ATen.h>
|
||||
// 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);
|
||||
@@ -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 <torch/extension.h>
|
||||
#include <ATen/ops/cudnn_convolution.h>
|
||||
#include <ATen/ops/scaled_dot_product_attention.h>
|
||||
#include <ATen/record_function.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
|
||||
#include "kernels.h"
|
||||
#include "prefix.h"
|
||||
#include "comm.h"
|
||||
#include "resample.h"
|
||||
#include "attn.h"
|
||||
#include "head.h"
|
||||
#include <array>
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
|
||||
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<c10d::ProcessGroup>;
|
||||
using at::Tensor;
|
||||
using Weights = std::vector<Tensor>;
|
||||
|
||||
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<int,int>{72,70},std::pair<int,int>{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<Tensor,110> 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<c10::IValue>());
|
||||
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_trim<kt,"Invalid causal prefix trim");
|
||||
int pt=kt-1-temporal_trim,ph=(sharded&&world>1&&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<c10::IValue>());
|
||||
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<int64_t> 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<Tensor> 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<c10::IValue>());
|
||||
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<c10::IValue>());
|
||||
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<c10::IValue>());
|
||||
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<c10::IValue>());
|
||||
// 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<c10::IValue>());
|
||||
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<Tensor> 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<world,"Invalid rank/world");
|
||||
TORCH_CHECK(weights.size()==770,"Call prepare() once before run()");
|
||||
TORCH_CHECK(frames_per_pass>=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<int64_t>(frames_per_pass*world,kMaxChunk);
|
||||
int T=z.size(2);
|
||||
for(int start=0;start<T;start+=chunk_size) {
|
||||
int count=std::min((int)chunk_size,T-start);
|
||||
Tensor pixels;
|
||||
if(world>1) 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<world;++peer) {
|
||||
int lo=44*peer/world,hi=44*(peer+1)/world;
|
||||
auto part=nhwc_empty(12,pixels.size(2),(hi-lo)*8,624,pixels);
|
||||
std::vector<Tensor> buffers={part};
|
||||
pg->recv(buffers,peer,0)->wait();
|
||||
write_video(part,out,frame_offset,lo*16);
|
||||
}
|
||||
} else {
|
||||
std::vector<Tensor> 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<py::gil_scoped_release>());
|
||||
m.def("run_with",&decode_video,py::call_guard<py::gil_scoped_release>());
|
||||
}
|
||||
@@ -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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
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<<<grid_for(n8), 256, 0, stream()>>>((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<<<grid_for(n8), 256, 0, stream()>>>((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<at::Tensor> 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;
|
||||
}
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
|
||||
#include <torch/csrc/distributed/c10d/Work.hpp>
|
||||
#include <c10/util/intrusive_ptr.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector>
|
||||
|
||||
using PG = c10::intrusive_ptr<c10d::ProcessGroup>;
|
||||
|
||||
// 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<c10::intrusive_ptr<c10d::Work>> works;
|
||||
std::vector<at::Tensor> 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);
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.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;
|
||||
|
||||
template<int M,int N,int CL>
|
||||
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<M>,Int<N>,Shape<_64>>;
|
||||
using Cluster=Shape<Int<CL>,_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<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
|
||||
cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
|
||||
using Problem=cutlass::conv::ConvProblemShape<Fprop,3>;
|
||||
using Kernel=cutlass::conv::kernel::ConvUniversal<Problem,Mainloop,Epilogue>;
|
||||
using Conv=cutlass::conv::device::ConvUniversalAdapter<Kernel>;
|
||||
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<cute::rank<0>(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];});
|
||||
cute::for_each(cute::make_seq<cute::rank<0>(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);
|
||||
}
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_bf16.h>
|
||||
#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<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(
|
||||
(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();
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
#include <ATen/ATen.h>
|
||||
|
||||
// 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);
|
||||
@@ -0,0 +1,99 @@
|
||||
#include "kernels.h"
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <algorithm>
|
||||
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<at::BFloat16>() : nullptr; }
|
||||
static B* ptr(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
|
||||
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<n/8;i+=(int64_t)gridDim.x*blockDim.x) {
|
||||
Pack v, bias, skip;
|
||||
v.v=((uint4*)x)[i];bias.v=((const uint4*)b)[i%(C/8)];
|
||||
if(s) skip.v=((const uint4*)s)[i];
|
||||
#pragma unroll
|
||||
for(int j=0;j<4;j++) { v.h[j]=__hadd2(v.h[j],bias.h[j]); if(s) v.h[j]=__hadd2(v.h[j],skip.h[j]); }
|
||||
((uint4*)x)[i]=v.v;
|
||||
}
|
||||
}
|
||||
__global__ void finish_scalar(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<n;i+=(int64_t)gridDim.x*blockDim.x) {
|
||||
B v=__float2bfloat16(__bfloat162float(x[i])+__bfloat162float(b[i%C]));
|
||||
if(s) v=__float2bfloat16(__bfloat162float(v)+__bfloat162float(s[i])); x[i]=v;
|
||||
}
|
||||
}
|
||||
void finish(Tensor& x,const Tensor& b,const Tensor& s) {
|
||||
int C=x.size(1);int64_t n=x.numel();
|
||||
if(C%8==0) finish_vector<<<std::min((int64_t)2048,(n+2047)/2048),256,0,stream()>>>(ptr(x),ptr(b),ptr(s),n,C);
|
||||
else finish_scalar<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(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<n) {
|
||||
int c=i%12;int64_t p=i/12;int w=p%W,h=(p/W)%H,t=p/((int64_t)W*H);
|
||||
int color=c/4,dy=c%2,dx=(c%4)/2;
|
||||
int64_t o=(((int64_t)color*fullT+t+offsetT)*fullH+offsetH+2*h+dy)*(2*W)+2*w+dx;
|
||||
out[o]=__float2bfloat16(fminf(1.f,fmaxf(-1.f,__bfloat162float(x[i]))));
|
||||
}
|
||||
}
|
||||
void write_video(const Tensor& x,Tensor& out,int frame_offset,int row_offset) {
|
||||
write_kernel<<<(x.numel()+255)/256,256,0,stream()>>>(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<n;i+=(int64_t)gridDim.x*blockDim.x) {
|
||||
int c=i%CV,p=i/CV,tx=p%TW,ty=(p/TW)%TH,t=p/(TW*TH);
|
||||
float d[2][6][6];
|
||||
#pragma unroll
|
||||
for(int a=0;a<6;a++) {
|
||||
int y=4*ty-1+a+PH;
|
||||
#pragma unroll
|
||||
for(int b=0;b<6;b++) {int xx=4*tx-1+b;float2 f=(y>=0&&y<HP&&xx>=0&&xx<W)?__bfloat1622float2(((const __nv_bfloat162*)x)[(((int64_t)t*HP+y)*W+xx)*CV+c]):make_float2(0,0);d[0][a][b]=f.x;d[1][a][b]=f.y;}
|
||||
}
|
||||
#pragma unroll
|
||||
for(int lane=0;lane<2;lane++) {
|
||||
#pragma unroll
|
||||
for(int b=0;b<6;b++) {float q[6];bt6(d[lane][0][b],d[lane][1][b],d[lane][2][b],d[lane][3][b],d[lane][4][b],d[lane][5][b],q);
|
||||
#pragma unroll
|
||||
for(int a=0;a<6;a++)d[lane][a][b]=q[a];}
|
||||
}
|
||||
#pragma unroll
|
||||
for(int a=0;a<6;a++) {float q0[6],q1[6];bt6(d[0][a][0],d[0][a][1],d[0][a][2],d[0][a][3],d[0][a][4],d[0][a][5],q0);bt6(d[1][a][0],d[1][a][1],d[1][a][2],d[1][a][3],d[1][a][4],d[1][a][5],q1);
|
||||
#pragma unroll
|
||||
for(int b=0;b<6;b++)((__nv_bfloat162*)v)[(6*a+b)*n+i]=__floats2bfloat162_rn(q0[b],q1[b]);}
|
||||
}
|
||||
}
|
||||
Tensor winograd_input(const Tensor& x,int ph) {
|
||||
int C=x.size(1),D=x.size(2),H=x.size(3)-2*ph,W=x.size(4),TH=(H+3)/4,TW=(W+3)/4;
|
||||
auto v=at::empty({36,D,TH,TW,C},x.options());int64_t n=(int64_t)D*TH*TW*C;
|
||||
winograd_input_kernel<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(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.)
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
#pragma once
|
||||
#include <ATen/ATen.h>
|
||||
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<Tensor>& 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);
|
||||
@@ -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<C> / 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 <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <algorithm>
|
||||
#include <cstdlib>
|
||||
#include <type_traits>
|
||||
|
||||
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 <int C, int P, int MODE>
|
||||
__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<at::BFloat16>();
|
||||
const B* gp = gamma.defined() ? (const B*)gamma.data_ptr<at::BFloat16>() : nullptr;
|
||||
const B* hp = hist.defined() ? (const B*)hist.data_ptr<at::BFloat16>() : nullptr;
|
||||
B* yp = (B*)y.data_ptr<at::BFloat16>();
|
||||
const int64_t rows = (int64_t)(T + pt) * (H + ph * 2);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
const int blocks = (int)std::min<int64_t>(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<CC, P, 2><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
|
||||
else if (silu) norm_pad_kernel<CC, P, 0><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
|
||||
else norm_pad_kernel<CC, P, 1><<<blocks, threads, 0, stream>>>(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<int, 1024>{});
|
||||
else if (C == 512) go(std::integral_constant<int, 512>{});
|
||||
else go(std::integral_constant<int, 256>{});
|
||||
} else {
|
||||
TORCH_CHECK(C % 8 == 0, "norm_pad copy needs C divisible by 8, got ", C);
|
||||
copy_row_kernel<<<blocks, threads, 0, stream>>>((const uint4*)xp, (const uint4*)hp, (uint4*)yp, T, H, W,
|
||||
pt, ph, C / 8);
|
||||
}
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return y;
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
// Prefix ownership transpose: no arithmetic, only bf16 NDHWC packets.
|
||||
#include "prefix.h"
|
||||
#include "kernels.h"
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
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<n;i+=(int64_t)gridDim.x*blockDim.x) {
|
||||
const int64_t t=i/((int64_t)(hi-lo)*WC8);
|
||||
const int64_t rowcol=i%((int64_t)(hi-lo)*WC8);
|
||||
send[offset+i]=x[(t*H+lo)*WC8+rowcol];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames,
|
||||
int rank,int world,const c10::intrusive_ptr<c10d::ProcessGroup>& 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<<<grid,256,0,at::cuda::getCurrentCUDAStream()>>>(
|
||||
(const uint4*)local.data_ptr(),(uint4*)send.data_ptr(),T,H,W*(C/8),world);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
std::vector<int64_t> in_split(world),out_split(world);
|
||||
int64_t sent=0,received=0;
|
||||
for(int p=0;p<world;++p) {
|
||||
in_split[p]=(int64_t)T*(H*(p+1)/world-H*p/world)*W*C;
|
||||
out_split[p]=(total_frames*(p+1)/world-total_frames*p/world)*band_height*W*C;
|
||||
sent+=in_split[p]; received+=out_split[p];
|
||||
}
|
||||
TORCH_CHECK(sent==send.numel() && received==recv.numel(),"Prefix packet size mismatch");
|
||||
// WorkNCCL::wait orders the caller's stream after the collective; it need
|
||||
// not block the host. ProcessGroup's tensor/allocator stream tracking keeps
|
||||
// send/recv storage live through that dependency. Do not drop this wait
|
||||
// when tail kernels consume output on the caller's current stream.
|
||||
pg->alltoall_base(recv,send,out_split,in_split)->wait();
|
||||
return output;
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
#pragma once
|
||||
#include <ATen/ATen.h>
|
||||
#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
|
||||
// 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<c10d::ProcessGroup>& pg);
|
||||
@@ -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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
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<at::BFloat16>(); }
|
||||
static B* pm(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
|
||||
|
||||
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<Tensor>& 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<Tensor> 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<<<grid, block, 0, stream()>>>(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;
|
||||
}
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <vector>
|
||||
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<Tensor>& pw, const Tensor& bias,
|
||||
const Tensor& skip, bool halo, bool temporal, bool first);
|
||||
@@ -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":[]}
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.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"
|
||||
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 <int M, int N, int CL>
|
||||
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<M>, Int<N>, Shape<_64>>;
|
||||
using Cluster = Shape<Int<CL>, _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<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
|
||||
cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
|
||||
using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
|
||||
using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
|
||||
using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
|
||||
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<cute::rank<0>(sc)>{},
|
||||
[&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
|
||||
cute::for_each(cute::make_seq<cute::rank<0>(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;
|
||||
}
|
||||
@@ -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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <algorithm>
|
||||
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<at::BFloat16>():nullptr;}
|
||||
static B* e_mut(Tensor& t){return t.defined()?(B*)t.data_ptr<at::BFloat16>():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<int C>
|
||||
__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<int C,bool EXACT>
|
||||
__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<C>(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<H&&xx<W);
|
||||
__nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
|
||||
if(bias) v=__hadd2(v,bv);
|
||||
if(residual&&ok) v=__hadd2(v,((const __nv_bfloat162*)residual)[yb+off]);
|
||||
if(ok) ((__nv_bfloat162*)y)[yb+off]=v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if(n){
|
||||
// Remapped epilogue (w01's M655, measured 1.58x on the standalone inverse): the thread that
|
||||
// computes the transform owns a channel pair, so it would write 4 bytes per output, which
|
||||
// caps the pass at the 4-byte row of the stream table. Instead the 16 outputs go to shared
|
||||
// memory with the bias already applied, one barrier, and the epilogue is re-laid so a warp
|
||||
// owns one position and each lane an octet of channels: the residual read and the y write are
|
||||
// one 512-byte transaction per warp instruction, and the sum of squares becomes a single-warp
|
||||
// butterfly because a warp now holds every channel of its position.
|
||||
__syncthreads(); // the components are dead; their buffer is reused
|
||||
B* so=smem; // so[p*C + ch], bias applied
|
||||
#pragma unroll
|
||||
for(int i=0;i<4;i++)
|
||||
#pragma unroll
|
||||
for(int b=0;b<4;b++){
|
||||
__nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
|
||||
if(bias) v=__hadd2(v,bv);
|
||||
((__nv_bfloat162*)so)[(i*4+b)*CV+c]=v;
|
||||
}
|
||||
__syncthreads();
|
||||
constexpr int NPW=16/WP; // positions per warp (4, 2 or 1)
|
||||
constexpr int CW=C/32; // channels per lane (8, 16 or 32)
|
||||
constexpr int NU=CW/8; // 16-byte chunks per lane (1, 2 or 4)
|
||||
const int wp=c>>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<NU;k++)
|
||||
gv[k]=gain?*(const uint4*)((const char*)gain+((long)(k*32+ln)*8)*2)
|
||||
:make_uint4(0x3f803f80u,0x3f803f80u,0x3f803f80u,0x3f803f80u);
|
||||
#pragma unroll
|
||||
for(int j=0;j<NPW;j++){
|
||||
const int p=wp*NPW+j, yy=4*ty+(p>>2), xx=4*tx+(p&3);
|
||||
const bool ok=EXACT?true:(yy<H&&xx<W);
|
||||
uint4 yv4[NU]; float sq=0.f;
|
||||
#pragma unroll
|
||||
for(int k=0;k<NU;k++){
|
||||
const int off=(k*32+ln)*8;
|
||||
uint4 v=*(const uint4*)((const char*)so+((long)p*C+off)*2);
|
||||
if(residual&&ok){
|
||||
// (long)t first: the whole offset must be 64-bit, not the last factor. `H` here is this
|
||||
// RANK'S BAND height, not the full 352 -- chunk() narrows dim 3 of the latent grid (44
|
||||
// rows) to [44r/W,44(r+1)/W) before resnet15 and every up block, so block 3 carries 88
|
||||
// rows at world 4 and 44 at world 8. The band shrinks as the temporal chunk grows, which
|
||||
// is why a rank's working set is near-constant -- BUT THE CENSUS THAT STOOD HERE WAS
|
||||
// TRUE ONLY OF `chunk_size = 8*world`, AND THIS PR MOVES THE RULE TO `min(16*world, 61)`.
|
||||
// Post-change THREE cells pass INT_MAX and THE LARGEST IS WORLD 2, so it is listed
|
||||
// first: narrowing against world 4 because "chunk 61 looks biggest" is exactly the
|
||||
// mislabelled-maximum this paragraph exists to prevent.
|
||||
// world 2 / chunk 32 D=125, rows 176, 624x256 = 3.51e9 <-- THE MAXIMUM
|
||||
// world 1 / chunk 16 D=61, rows 352 = 3.43e9
|
||||
// world 4 / chunk 61 D=241, rows 88 = 3.39e9
|
||||
// World 2's chunk 32 is THIS PR's OWN value at world 2 (min(16*2, 61) = 32), not a
|
||||
// foreign or historical form. ALL THREE ARE PAST INT_MAX. This site is
|
||||
// correct because it chains `(long)` here and again on `nbase` below, so the 64-bit chain
|
||||
// is now LOAD-BEARING at those three cells rather than merely belt-and-braces. Do not
|
||||
// narrow it back on the strength of an `inside INT_MAX` reading from the old rule.
|
||||
const uint4 r=*(const uint4*)((const char*)residual
|
||||
+((long)((((long)t*H+yy)*W+xx)*C+off))*2);
|
||||
v.x=add2(v.x,r.x); v.y=add2(v.y,r.y); v.z=add2(v.z,r.z); v.w=add2(v.w,r.w);
|
||||
}
|
||||
yv4[k]=v; sq+=sq8(v);
|
||||
}
|
||||
#pragma unroll
|
||||
for(int d=16;d;d>>=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<NU;k++){
|
||||
const long cof=(long)((k*32+ln)*8)*2;
|
||||
if(wy) *(uint4*)((char*)y+ybase+cof)=yv4[k];
|
||||
uint4 u=yv4[k], g=gv[k], o;
|
||||
float2 f0=b2f(u.x), f1=b2f(u.y), f2=b2f(u.z), f3=b2f(u.w);
|
||||
float2 g0=b2f(g.x), g1=b2f(g.y), g2=b2f(g.z), g3=b2f(g.w);
|
||||
o.x=pack2(silu2(f0.x*s*g0.x),silu2(f0.y*s*g0.y));
|
||||
o.y=pack2(silu2(f1.x*s*g1.x),silu2(f1.y*s*g1.y));
|
||||
o.z=pack2(silu2(f2.x*s*g2.x),silu2(f2.y*s*g2.y));
|
||||
o.w=pack2(silu2(f3.x*s*g3.x),silu2(f3.y*s*g3.y));
|
||||
*(uint4*)((char*)n+nbase+cof)=o;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Zero the spatial padding rows of every frame of the prepared tensor.
|
||||
__global__ void epi_border(B* n,int C,int T,int Hp,int W){
|
||||
const int CV=C/2; const long per=(long)2*W*CV;
|
||||
const long total=(long)T*per;
|
||||
for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<total;i+=(long)gridDim.x*blockDim.x){
|
||||
const int t=(int)(i/per); const long r=i%per; const int edge=(int)(r/((long)W*CV));
|
||||
const long r2=r%((long)W*CV); const int w=(int)(r2/CV),cc=(int)(r2%CV);
|
||||
const long idx=(((long)t*Hp+(edge?Hp-1:0))*W+w)*CV+cc;
|
||||
((__nv_bfloat162*)n)[idx]=__floats2bfloat162_rn(0.f,0.f);
|
||||
}
|
||||
}
|
||||
__global__ void epi_zero(B* n,long pairs){
|
||||
const __nv_bfloat162 z=__floats2bfloat162_rn(0.f,0.f);
|
||||
for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<pairs;i+=(long)gridDim.x*blockDim.x)
|
||||
((__nv_bfloat162*)n)[i]=z;
|
||||
}
|
||||
|
||||
template<int C,bool EXACT>
|
||||
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<C,EXACT>;
|
||||
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<<<grid,block,smem,s>>>(p,y,n,b,r,g,H,W,TH,TW,PT,PH,wy);
|
||||
}
|
||||
|
||||
WinoOut wino_epilogue(const std::vector<Tensor>& 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<at::BFloat16>();
|
||||
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;
|
||||
}
|
||||
@@ -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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <algorithm>
|
||||
|
||||
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 <int CP, int SPLIT>
|
||||
__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<at::BFloat16>();
|
||||
B* vp = (B*)v.data_ptr<at::BFloat16>();
|
||||
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<CP, SPLIT><<<g, CP, 0, st>>>(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<CP, SPLIT><<<g, CP, 0, st>>>(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<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
|
||||
break;
|
||||
}
|
||||
}
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return v;
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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 <clip>x<world>. 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_<clip>.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()
|
||||
@@ -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()
|
||||
@@ -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`: "<file>::<function>", 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 "<file_path>::<function_name>"')
|
||||
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()
|
||||
@@ -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<c10d::ProcessGroup>&` (from <torch/csrc/distributed/c10d/ProcessGroup.hpp>).
|
||||
- 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 <clip>x<world>. A clip is the first T latent frames of one of
|
||||
three real 20 s clips encoded by the VAE (latents_<source>.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_<clip>.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_<source>.pt ({"latents": fp32 [1, 48, 121, 44, 78]}), ref_<clip>.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")
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
@@ -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)
|
||||
@@ -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=<PAN2 root> 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)
|
||||
Reference in New Issue
Block a user