diff --git a/.gitignore b/.gitignore index 8d5751ec..1921ef6a 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,13 @@ venv/ # IDE files .vscode/ .idea/ + +# FlashVSR model weights & outputs (examples) +examples/WanVSR/FlashVSR/ +examples/WanVSR/FlashVSR-v1.1/ +examples/WanVSR/results/ + +# Profiling artifacts (reports/caches/raw runs are large binaries) +examples/WanVSR/profiling/reports/ +examples/WanVSR/profiling/cache/ +examples/WanVSR/profiling/runs/ diff --git a/README.md b/README.md index 1cfb6ee8..a360051b 100644 --- a/README.md +++ b/README.md @@ -187,6 +187,100 @@ python infer_flashvsr_v1.1_tiny_long_video.py --- +### ⚡ Hopper Acceleration (optional, GH200 / sm_90) + +FlashVSR ships a set of **opt-in** fast paths for NVIDIA Hopper GPUs (e.g. GH200, `sm_90`). They are controlled by environment variables and are **all OFF by default**: with no variables set, the output is **bit-for-bit identical** to the standard path and Ampere / A100 are unaffected. Recoverable kernel/setup failures fall back to the original path and are exposed through telemetry. Direct-output failures after stateful decoding begins fail closed because replaying mutated decoder state is unsafe. + +On a GH200 at 768x1408, the current production stack reaches **57.26 FPS core +E2E** for v1.1 Tiny (F=81): the bit-identical stack reaches 54.61 FPS +(+10.9% over the 49.23 FPS Phase-5 baseline), and the quality-gated +`TCDECODER_CUDNN_FUSED` path (55.4 dB PSNR vs the lossless stack, gate 49 dB) +plus the Phase-7 stack (49.59 dB PSNR vs the prior production set, gate 49 dB) +add the remaining gain. + +| Env var | Values | Default | Effect | +|---|---|---|---| +| `FLASHVSR_CONV3D_BACKEND` | `auto`, `gemm` | `auto` | `gemm` = im2col + WGMMA conv3d for the LQ projector (largest single win) | +| `FLASHVSR_TCDECODER_CHANNELS_LAST` | `0`, `1` | `0` | NHWC TCDecoder (bit-identical) | +| `FLASHVSR_FUSE_NORM` | `0`, `1` | `0` | fuse norm / modulate / gate via `torch.compile` | +| `FLASHVSR_DIT_ROW_FUSION` | `0`, `1` | `0` | Triton affine-free LayerNorm + AdaLN and residual-gate fusion — quality-gated | +| `FLASHVSR_ATTN_BACKEND` | `sparse`, `triton`, `triton2`, `auto`, `dense` | `sparse` | `triton2` = warp-specialized Hopper block-sparse kernel (same mask) | +| `FLASHVSR_ATTN_TMA` | `0`, `1` | `1` | TMA bulk loads (only used by the `triton` backend) | +| `FLASHVSR_CONV3D_IM2COL_BUDGET_GB` | float | `2.0` | chunked im2col memory budget for the `gemm` backend | +| `FLASHVSR_CACHE_MOD` | `0`, `1` | `0` | cache step-invariant modulation (bit-identical) | +| `FLASHVSR_CACHE_MASK_BIAS` | `0`, `1` | `0` | cache the geometry-only attention bias (bit-identical) | +| `FLASHVSR_FUSE_ROPE` | `0`, `1` | `0` | fused single-kernel RoPE apply, same fp64 math (bit-identical) | +| `FLASHVSR_KV_RINGBUF` | `0`, `1` | `0` | preallocated KV-cache arena, removes the per-chunk KV concat (bit-identical; retains a little extra memory, see `FLASHVSR_KV_RINGBUF_SPARE`) | +| `FLASHVSR_KV_RINGBUF_SPARE` | int | `2` | arena spare slots: higher = rarer compaction copies, more retained memory | +| `FLASHVSR_ATTN_STRIDED_IO` | `0`, `1` | `0` | strided q/k/v/out for the `triton` backend — removes all attention-path transpose copies (bit-identical) | +| `FLASHVSR_MASKGEN_LEAN` | `0`, `1` | `0` | mask-generation cleanup (kthvalue select, no repeat copy, cached seqlens) — exact same mask (bit-identical) | +| `FLASHVSR_MASKGEN_THRESHOLD_CACHE` | `0`, `1` | `0` | reuse each DiT block's previous steady-chunk threshold — quality-gated | +| `FLASHVSR_LQPROJ_LEAN` | `0`, `1` | `0` | LQ projector: single-materialization causal pad + no cache clones (bit-identical) | +| `FLASHVSR_CACHE_ROPE_FREQS` | `0`, `1` | `0` | assemble RoPE freqs on-device in a cached buffer (bit-identical; useful when the CPU is loaded) | +| `FLASHVSR_CONV3D_PACKER` | `eager`, `triton` | `eager` | exact Triton im2col packing; retains the existing `torch.addmm` numerics | +| `FLASHVSR_TCDECODER_POINTER_STATE` | `0`, `1` | `0` | rotate recurrent-state tensor references instead of copying them | +| `FLASHVSR_TCDECODER_DIRECT_OUTPUT` | `0`, `1` | `0` | write overlapped decoder chunks directly into the final output | +| `FLASHVSR_TCDECODER_FUSE_POINTWISE` | `0`, `1` | `0` | exact BF16 MemBlock bias/ReLU/residual Triton fusion | +| `FLASHVSR_TCDECODER_UPSAMPLE` | `0`, `1` | `0` | exact channels-last nearest-neighbor Triton kernel | +| `FLASHVSR_TCDECODER_CONCAT` | `0`, `1` | `0` | exact channels-last recurrent concat Triton kernel | +| `FLASHVSR_TCDECODER_TGROW_UP` | `0`, `1` | `0` | run the 1x1 TGrow conv at low resolution and fuse temporal unpack + nearest upsample into one Triton kernel (measured bit-identical E2E) | +| `FLASHVSR_TCDECODER_CUDNN_FUSED` | `0`, `1` | `0` | cuDNN runtime-fused Conv+Bias(+Add)+ReLU decoder engines — **quality-gated** (~55 dB PSNR vs the lossless stack), not bit-exact | +| `FLASHVSR_TCDECODER_SPLITK_CONV` | `0`, `1` | `0` | avoid recurrent MemBlock concat via split-weight cuDNN convs; requires `TCDECODER_CUDNN_FUSED=1`, quality-gated | + +**Recommended full-speed config** (run from `examples/WanVSR`): + +```bash +FLASHVSR_CONV3D_BACKEND=gemm \ +FLASHVSR_TCDECODER_CHANNELS_LAST=1 \ +FLASHVSR_FUSE_NORM=1 \ +FLASHVSR_DIT_ROW_FUSION=1 \ +FLASHVSR_ATTN_BACKEND=triton2 \ +FLASHVSR_CACHE_MOD=1 \ +FLASHVSR_CACHE_MASK_BIAS=1 \ +FLASHVSR_CACHE_ROPE_FREQS=0 \ +FLASHVSR_FUSE_ROPE=1 \ +FLASHVSR_KV_RINGBUF=1 \ +FLASHVSR_ATTN_STRIDED_IO=1 \ +FLASHVSR_MASKGEN_LEAN=1 \ +FLASHVSR_MASKGEN_THRESHOLD_CACHE=1 \ +FLASHVSR_LQPROJ_LEAN=1 \ +FLASHVSR_FUSED_CSR=1 \ +FLASHVSR_ROPE_KERNEL=triton \ +FLASHVSR_POOLED_K_CACHE=1 \ +FLASHVSR_ATTN_ZEROCOPY=1 \ +FLASHVSR_DECODER_OVERLAP=1 \ +FLASHVSR_FP8_GEMM=0 \ +FLASHVSR_CONV3D_PACKER=triton \ +FLASHVSR_TCDECODER_POINTER_STATE=1 \ +FLASHVSR_TCDECODER_DIRECT_OUTPUT=1 \ +FLASHVSR_TCDECODER_FUSE_POINTWISE=1 \ +FLASHVSR_TCDECODER_UPSAMPLE=1 \ +FLASHVSR_TCDECODER_CONCAT=1 \ +FLASHVSR_TCDECODER_TGROW_UP=1 \ +FLASHVSR_TCDECODER_CUDNN_FUSED=1 \ +FLASHVSR_TCDECODER_SPLITK_CONV=1 \ +python infer_flashvsr_v1.1_tiny.py +``` + +The same preset is available as: + +```bash +./run_flashvsr_v1.1_tiny_gh200.sh +``` + +> **Notes** +> - The `triton`/`triton2` attention backends, Triton decoder/LQ kernels, and `gemm` Conv3D require a Hopper GPU (`sm_90`); unsupported paths fall back automatically. +> - `TCDECODER_DIRECT_OUTPUT` requires `DECODER_OVERLAP=1`; decoder Triton paths and `TCDECODER_SPLITK_CONV` require `TCDECODER_CHANNELS_LAST=1`; split-K also requires `TCDECODER_CUDNN_FUSED=1`; and `CONV3D_PACKER=triton` requires `CONV3D_BACKEND=gemm`. +> - `channels_last`, `CACHE_MOD`, `CACHE_MASK_BIAS`, and the Phase-2A knobs (`FUSE_ROPE`, `KV_RINGBUF`, `ATTN_STRIDED_IO`, `MASKGEN_LEAN`, `LQPROJ_LEAN`, `CACHE_ROPE_FREQS`) are bit-identical vs the same config without them (`max|diff| = 0`, see `examples/WanVSR/profiling/PHASE_BENCH_LOG.md`). `FUSE_NORM` and the `triton` backend are near-identical (~49-50 dB PSNR vs the default), not bit-exact, due to fp/accumulation order, so they are opt-in. +> - Parity + speed for each path can be checked with the `examples/WanVSR/test_*.py` scripts (`test_phase2a_lossless.py` covers the Phase-2A knobs; `test_phase5_lossless.py` modes `phase6`/`tgrowup`/`cudnnfuse` cover the Phase-6/6b decoder knobs; `test_tcdecoder_tgrow_up.py` is the isolated TGROW_UP kernel test). +> - Phase-2A stack measured on GH200: 38.6 → 41.6 FPS @768x1408 (steady chunk 156 → 138 ms) and 11.0 → 11.5 FPS @1536x2560; `KV_RINGBUF` retains ~+3 GiB @768 (tunable via `FLASHVSR_KV_RINGBUF_SPARE`). +> - Phase-6 production stack measured on GH200: 49.23 → 53.42 FPS @768x1408 and 13.44 → 14.68 FPS @1536x2560. All Phase-6 output changes are bit-identical against the Phase-5 production stack, including the accepted `8n+5` frame-count edge case. +> - Phase-6b: `TCDECODER_TGROW_UP` (bit-identical vs the Phase-6 stack, 53.40 → 54.61 FPS @768 3-run median; single-run spot-check 14.36 → 14.69 FPS @1536x2560 F=41, +2.30%) and `TCDECODER_CUDNN_FUSED` (quality-gated, not bit-exact: 54.58 → 55.91 FPS @768 3-run median; 14.69 → 15.07 FPS @1536 with TGROW_UP already on, +2.59%; E2E 55.4–55.8 dB PSNR vs the lossless stack, gate 49 dB). Disable `TCDECODER_CUDNN_FUSED` for a strictly `max|diff|=0` pipeline; `TCDECODER_TGROW_UP` stays on in that case. +> - Phase-7: row-wise DiT fusion, steady-threshold reuse, and MemBlock split-K move the quality-gated production stack **55.91 → 57.26 FPS** @768x1408 F=81 (3-run medians, +2.41%). Combined E2E PSNR is 49.81 dB at F=29 and 49.59 dB at F=89 vs the preceding production stack (gate 49 dB). The @1536x2560 spot reached 15.42 FPS at F=41 with 49.95 dB at F=29. +> - Use `FLASHVSR_REQUIRE_FASTPATHS=1` with `profiling/run_pipe_target.py` to fail the benchmark if any requested backend silently falls back. + +--- + ### 🛠️ Method The overview of **FlashVSR**. This framework features: diff --git a/diffsynth/models/fp8_gemm.py b/diffsynth/models/fp8_gemm.py new file mode 100644 index 00000000..70250830 --- /dev/null +++ b/diffsynth/models/fp8_gemm.py @@ -0,0 +1,305 @@ +"""Phase 2B-2: FP8 GEMM infrastructure (FLASHVSR_FP8_GEMM, default OFF). + +Routes the DiT hot-path GEMMs (self-attn q/k/v/o, ffn1/ffn2, LQ-projector +per-layer linears) through `torch._scaled_mm` with float8_e4m3fn operands: + + * Weights are pre-cast ONCE per module (lazy, cached on the module) with + per-output-channel scales -> (N, K) e4m3 + (1, N) fp32 scales. + * Activations are quantized dynamically per call with per-row scales via a + single Triton kernel (row amax + scale + cast in one launch; the eager + equivalent chain of abs/amax/div/clamp/cast kernels measured ~5x slower + than bandwidth and would erase the GEMM win). + * ffn2's input quantization is fused with the GELU(tanh) activation in the + same kernel (fp32 gelu -> e4m3), replacing the eager bf16 GELU pass + entirely; on the widest activation in the network this is the difference + between FP8 being a net win and a net loss. + * Bias is applied inside `_scaled_mm`'s epilogue (bf16 bias, bf16 out). + +Numerics: this is NOT a lossless path (e4m3 mantissa = 3 bits). It ships +permanently default-OFF; the enable decision belongs to the Phase-4 quality +protocol (PSNR/SSIM vs flag-OFF). Rollback: FLASHVSR_FP8_GEMM=0 (default) +leaves every call site on the original eager path, bit-for-bit. + +Scope control (comma list, default all): FLASHVSR_FP8_GEMM_SCOPE=qkv,o,ffn,lq +— used for per-site attribution if the Phase-4 quality audit needs to bisect. + +Any failure inside the FP8 path (unsupported build, shape, dtype) disables it +for the rest of the process (warn once) and falls back to eager — same silent +degradation pattern as the attention backends. +""" +import os + +import torch + +_FP8_GEMM = os.environ.get("FLASHVSR_FP8_GEMM", "0") != "0" +_FP8_SCOPE = frozenset( + s.strip() for s in + os.environ.get("FLASHVSR_FP8_GEMM_SCOPE", "qkv,o,ffn,lq").split(",") + if s.strip() +) +# Phase 4B: rowwise (2B-2) vs blockwise (DeepSeek-style 1x128 activation + +# 1x128 weight scales via torch._scaled_mm_v2). Blockwise captures per-128-K +# activation outliers (the post-GELU spikes that pinned rowwise at 40.7 dB), +# targeting the >=49 dB enable gate. Requires cublasLt >=12.9 + sm_90. +_FP8_MODE = os.environ.get("FLASHVSR_FP8_GEMM_MODE", "rowwise").lower() +# Quality-probe switch: eager group-quant (correct but slow) so the >=49 dB +# gate can be measured before investing in the fused blockwise quant kernel. +_FP8_BLOCKWISE_EAGER = os.environ.get("FLASHVSR_FP8_BLOCKWISE_EAGER", "1") != "0" + +try: + from torch.nn.functional import ScalingType as _ScalingType, SwizzleType as _SwizzleType + _HAS_V2 = hasattr(torch, "_scaled_mm_v2") +except Exception: + _HAS_V2 = False + +_FAILED = False # sticky: set on first unexpected error, forces eager fallback + +try: + import triton + import triton.language as tl + _HAS_TRITON = True +except Exception: # pragma: no cover + _HAS_TRITON = False + + +def _capability_ok(): + if not (_HAS_TRITON and hasattr(torch, "_scaled_mm")): + return False + try: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return (major, minor) >= (8, 9) # e4m3 matmul: sm_89+ + except Exception: + return False + + +_CAP_OK = None + + +def enabled(site: str) -> bool: + """Per-call-site gate. Reads the module flag at call time so the parity + harness can toggle `fp8_gemm._FP8_GEMM` on a live session.""" + global _CAP_OK + if not _FP8_GEMM or _FAILED or site not in _FP8_SCOPE: + return False + if _CAP_OK is None: + _CAP_OK = _capability_ok() + return _CAP_OK + + +if _HAS_TRITON: + + @triton.jit + def _rowquant_kernel(X, Y, S, K, + stride_xm, stride_ym, + APPLY_GELU: tl.constexpr, BLOCK_K: tl.constexpr): + """Per-row dynamic e4m3 quantization (optionally fused with GELU-tanh). + + One program per row. Two passes over the row tiles (amax, then cast); + the second pass re-reads through L2. All math in fp32. + Y[row, :] = clamp(x * 448/amax, +-448) as e4m3; S[row] = amax/448. + """ + row = tl.program_id(0).to(tl.int64) + xrow = X + row * stride_xm + yrow = Y + row * stride_ym + + amax = 0.0 + for k0 in range(0, K, BLOCK_K): + offs = k0 + tl.arange(0, BLOCK_K) + m = offs < K + v = tl.load(xrow + offs, mask=m, other=0.0).to(tl.float32) + if APPLY_GELU: + # torch GELU(approximate='tanh'): 0.5*v*(1+tanh(c0*(v+c1*v^3))) + # tanh via exp(-2|t|) (never overflows), sign restored after. + t = 0.7978845608028654 * (v + 0.044715 * v * v * v) + e = tl.exp(-2.0 * tl.abs(t)) + th = (1.0 - e) / (1.0 + e) + th = tl.where(t < 0.0, -th, th) + v = 0.5 * v * (1.0 + th) + amax = tl.maximum(amax, tl.max(tl.abs(v), axis=0)) + + amax = tl.maximum(amax, 1e-12) + inv = 448.0 / amax + for k0 in range(0, K, BLOCK_K): + offs = k0 + tl.arange(0, BLOCK_K) + m = offs < K + v = tl.load(xrow + offs, mask=m, other=0.0).to(tl.float32) + if APPLY_GELU: + t = 0.7978845608028654 * (v + 0.044715 * v * v * v) + e = tl.exp(-2.0 * tl.abs(t)) + th = (1.0 - e) / (1.0 + e) + th = tl.where(t < 0.0, -th, th) + v = 0.5 * v * (1.0 + th) + q = tl.minimum(tl.maximum(v * inv, -448.0), 448.0) + tl.store(yrow + offs, q.to(tl.float8e4nv), mask=m) + tl.store(S + row, amax / 448.0) + + +def quant(x2d: torch.Tensor, apply_gelu: bool = False): + """Shared activation quant. Rowwise (2B-2) or blockwise (4B) per mode.""" + if _FP8_MODE == "blockwise": + if x2d.stride(1) != 1: + x2d = x2d.contiguous() + return (_blockquant_eager_gelu(x2d) if apply_gelu + else _blockquant_eager(x2d)) + assert x2d.dim() == 2, "2D input required" + if x2d.stride(1) != 1: + # e.g. the LQ projector's rearrange returns a transposed view; one + # contiguous copy here is amortized over all GEMMs sharing the quant. + x2d = x2d.contiguous() + M, K = x2d.shape + y = torch.empty(M, K, dtype=torch.float8_e4m3fn, device=x2d.device) + s = torch.empty(M, 1, dtype=torch.float32, device=x2d.device) + _rowquant_kernel[(M,)]( + x2d, y, s, K, x2d.stride(0), y.stride(0), + APPLY_GELU=apply_gelu, BLOCK_K=1024, num_warps=4, + ) + return y, s + + +def _weight_fp8(lin): + """Lazy per-module weight cast: (N,K) e4m3 + (1,N) fp32 per-out-channel + scales, cached on the module. Weights are inference-static; the cache key + guards against reload/device moves. Works for AutoWrappedLinear too + (.weight/.bias exposed).""" + w = lin.weight + key = (w.data_ptr(), w.device, w.dtype, w.shape) + cache = getattr(lin, "_fp8_wcache", None) + if cache is not None and cache[0] == key: + return cache[1], cache[2] + with torch.no_grad(): + wf = w.detach().to(torch.float32) + s = wf.abs().amax(dim=1, keepdim=True).clamp_(min=1e-12) / 448.0 # (N,1) + w8 = (wf / s).clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) # (N,K) + sb = s.reshape(1, -1) # (1,N) fp32, contiguous view of (N,1) + lin._fp8_wcache = (key, w8, sb) + return w8, sb + + +def _blockquant_eager(x2d): + """Eager 1x128 group quant (quality probe). (M,K)->(e4m3 (M,K), + scale (M,K//128) laid out as .t() of contiguous (K//128,M)).""" + M, K = x2d.shape + G = K // 128 + xf = x2d.to(torch.float32).reshape(M, G, 128) + amax = xf.abs().amax(-1).clamp_(min=1e-12) # (M,G) + s = amax / 448.0 + x8 = (xf / s[..., None]).clamp_(-448.0, 448.0).to( + torch.float8_e4m3fn).reshape(M, K) + sa = s.t().contiguous().t() # (M,G) col-major + return x8, sa + + +def _blockquant_eager_gelu(h2d): + """GELU(tanh) fused into 1x128 group quant (ffn2 input).""" + v = h2d.to(torch.float32) + v = torch.nn.functional.gelu(v, approximate="tanh") + return _blockquant_eager(v) + + +def _weight_fp8_blockwise(lin): + """Lazy 1x128-along-K weight cast: b operand (K,N) col-major e4m3 + + scale (N,K//128) laid out as .t() of contiguous (K//128,N).""" + w = lin.weight + key = (w.data_ptr(), w.device, w.dtype, w.shape) + cache = getattr(lin, "_fp8_wcache_bw", None) + if cache is not None and cache[0] == key: + return cache[1], cache[2] + N, K = w.shape + G = K // 128 + with torch.no_grad(): + wf = w.detach().to(torch.float32).reshape(N, G, 128) + wamax = wf.abs().amax(-1).clamp_(min=1e-12) # (N,G) + ws = wamax / 448.0 + w8 = (wf / ws[..., None]).clamp_(-448.0, 448.0).to( + torch.float8_e4m3fn).reshape(N, K) + b8 = w8.t() # (K,N) col-major + sb = ws.t().contiguous().t() # (N,G) col-major + lin._fp8_wcache_bw = (key, b8, sb) + return b8, sb + + +def _mm_blockwise(x8, sa, b8, sb, bias, out_dtype): + return torch._scaled_mm_v2( + x8, b8, + [sa], [_ScalingType.BlockWise1x128], [_SwizzleType.NO_SWIZZLE], + [sb], [_ScalingType.BlockWise1x128], [_SwizzleType.NO_SWIZZLE], + bias, out_dtype) + + +def _linear_blockwise(lin, x, pre=None): + K = x.shape[-1] + x2 = x.reshape(-1, K) + if x2.stride(1) != 1: + x2 = x2.contiguous() + x8, sa = _blockquant_eager(x2) if pre is None else pre + b8, sb = _weight_fp8_blockwise(lin) + out = _mm_blockwise(x8, sa, b8, sb, lin.bias, x.dtype) + return out.reshape(*x.shape[:-1], b8.shape[1]) + + +def _ffn_blockwise(seq, x): + h = _linear_blockwise(seq[0], x) + h2 = h.reshape(-1, h.shape[-1]) + g8, gs = _blockquant_eager_gelu(h2) + b8, sb = _weight_fp8_blockwise(seq[2]) + out = _mm_blockwise(g8, gs, b8, sb, seq[2].bias, x.dtype) + return out.reshape(*x.shape[:-1], b8.shape[1]) + + +def _warn_and_disable(e): + global _FAILED + if not _FAILED: + _FAILED = True + print(f"[fp8_gemm] disabled after error, falling back to eager: {e!r}") + + +def linear(lin, x: torch.Tensor, pre=None) -> torch.Tensor: + """FP8 replacement for lin(x). `pre` = optional shared (x8, scales) from + `quant()` when several linears consume the same activation (qkv, LQ). + Falls back to eager on any error (sticky).""" + try: + if _FP8_MODE == "blockwise": + return _linear_blockwise(lin, x, pre=pre) + K = x.shape[-1] + x2 = x.reshape(-1, K) + x8, sa = quant(x2) if pre is None else pre + w8, sb = _weight_fp8(lin) + out = torch._scaled_mm(x8, w8.t(), scale_a=sa, scale_b=sb, + bias=lin.bias, out_dtype=x.dtype) + return out.reshape(*x.shape[:-1], w8.shape[0]) + except Exception as e: # pragma: no cover + _warn_and_disable(e) + return torch.nn.functional.linear(x, lin.weight, lin.bias) + + +def ffn(seq, x: torch.Tensor) -> torch.Tensor: + """FP8 replacement for Sequential(Linear, GELU(tanh), Linear)(x). + ffn1 out stays bf16; the GELU is fused into ffn2's input quantization + (fp32 gelu -> e4m3 directly, replacing the eager bf16 GELU pass).""" + try: + if _FP8_MODE == "blockwise": + return _ffn_blockwise(seq, x) + h = linear(seq[0], x) + h2 = h.reshape(-1, h.shape[-1]) + g8, gs = quant(h2, apply_gelu=True) + w8, sb = _weight_fp8(seq[2]) + out = torch._scaled_mm(g8, w8.t(), scale_a=gs, scale_b=sb, + bias=seq[2].bias, out_dtype=x.dtype) + return out.reshape(*x.shape[:-1], w8.shape[0]) + except Exception as e: # pragma: no cover + _warn_and_disable(e) + return seq(x) + + +def ffn_partial(seq, x: torch.Tensor) -> torch.Tensor: + """Scope 'ffn1' (not in the default set): e4m3 ffn1 only, eager bf16 + GELU + ffn2. Quality/speed bisection point — ffn2's post-GELU activation + is outlier-heavy and is the dominant fp8 quality cost.""" + try: + h = linear(seq[0], x) + return seq[2](seq[1](h)) + except Exception as e: # pragma: no cover + _warn_and_disable(e) + return seq(x) diff --git a/diffsynth/models/stepvideo_text_encoder.py b/diffsynth/models/stepvideo_text_encoder.py index 598825a9..486fdccf 100644 --- a/diffsynth/models/stepvideo_text_encoder.py +++ b/diffsynth/models/stepvideo_text_encoder.py @@ -18,7 +18,11 @@ import torch.nn.functional as F from .stepvideo_dit import RMSNorm from safetensors.torch import load_file -from transformers.modeling_utils import PretrainedConfig, PreTrainedModel +try: + from transformers.modeling_utils import PretrainedConfig, PreTrainedModel +except ImportError: # newer transformers no longer re-exports PretrainedConfig from modeling_utils + from transformers.modeling_utils import PreTrainedModel + from transformers import PretrainedConfig from einops import rearrange import json from typing import List diff --git a/diffsynth/models/triton_block_sparse_attn.py b/diffsynth/models/triton_block_sparse_attn.py new file mode 100644 index 00000000..56d3fada --- /dev/null +++ b/diffsynth/models/triton_block_sparse_attn.py @@ -0,0 +1,326 @@ +"""Hopper (sm_90) WGMMA block-sparse FlashAttention kernel (Triton). + +FlashVSR's self-attention uses a block-sparse mask (block size 128) at density +~0.6. The bundled `block_sparse_attn` CUDA kernel is FlashAttention-2 style and +emits only Ampere `HMMA` tensor-core ops even on sm_90 (measured: 380928 HMMA, +0 WGMMA), reaching only ~33% of bf16 peak (~327 TFLOP/s). cuDNN's dense fused +attention reaches ~62% peak (605 TFLOP/s) using Hopper `WGMMA`, but it cannot +express FlashVSR's arbitrary 2D-spatial block mask (cuDNN `block_mask` is not +available on sm_90, and its 1D diagonal-band masks don't cover a 2D-local +pattern efficiently). + +This Triton kernel closes the gap: it honors the exact per-(q_block, kv_block) +boolean mask while compiling to Hopper `WGMMA` (verified in PTX/SASS), giving a +bit-for-bit-equivalent result (cos 0.99999 vs block_sparse_attn) at ~1.2x the +speed (6.0 vs 7.4 ms at the real 25344x12x128 / density-0.606 shape). + +Forward-only, bf16, no dropout. Used opt-in via FLASHVSR_ATTN_BACKEND=triton, +guarded to sm_90, with a silent fallback to the original block_sparse kernel. +""" +import math +import os + +import torch + +try: + import triton + import triton.language as tl + _TRITON_OK = True +except Exception: # pragma: no cover + _TRITON_OK = False + +# Optional TMA (Tensor Memory Accelerator) path: device TMA bulk loads of Q/K/V +# overlap with WGMMA, lifting peak from ~40% to ~44% on GH200 (5.6 vs 6.1 ms at +# the real shape). Requires Triton >= 3.x host TensorDescriptor + an allocator. +_TMA_OK = False +if _TRITON_OK: + try: + from triton.tools.tensor_descriptor import TensorDescriptor as _TensorDescriptor + triton.set_allocator( + lambda size, align, stream: torch.empty(size, dtype=torch.int8, device="cuda") + ) + _TMA_OK = True + except Exception: + _TMA_OK = False + +_USE_TMA = os.environ.get("FLASHVSR_ATTN_TMA", "1") != "0" + + +if _TRITON_OK: + + @triton.jit + def _bsfa_kernel( + Q, K, V, O, KVIdx, KVCnt, sm_scale, + sqh, sqm, sqk, skh, skn, skk, svh, svn, svk, soh, som, sok, + sih, sim, sic, sch, scm, H, N_Q, N_KV, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, + ): + start_m = tl.program_id(0) + off_h = tl.program_id(1) + qo = off_h * sqh + kvo = off_h * skh + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, HEAD_DIM) + q = tl.load(Q + qo + offs_m[:, None] * sqm + offs_k[None, :] * sqk, + mask=offs_m[:, None] < N_Q, other=0.0) + qs = (q * sm_scale).to(q.dtype) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + cnt = tl.load(KVCnt + off_h * sch + start_m * scm) + base = KVIdx + off_h * sih + start_m * sim + for j in range(0, cnt): + kvb = tl.load(base + j * sic) + n = kvb * BLOCK_N + offs_n + k = tl.load(K + kvo + n[None, :] * skn + offs_k[:, None] * skk, + mask=n[None, :] < N_KV, other=0.0) + qk = tl.dot(qs, k) + qk = tl.where(n[None, :] < N_KV, qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2((qk - m_ij[:, None]) * 1.44269504) + alpha = tl.math.exp2((m_i - m_ij) * 1.44269504) + l_i = l_i * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + vv = tl.load(V + kvo + n[:, None] * svn + offs_k[None, :] * svk, + mask=n[:, None] < N_KV, other=0.0) + acc += tl.dot(p.to(vv.dtype), vv) + m_i = m_ij + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + tl.store(O + qo + offs_m[:, None] * som + offs_k[None, :] * sok, + acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_Q) + + +if _TMA_OK: + + # ----------------------------------------------------------------------- + # Phase 2A-3 strided-IO TMA kernel (FLASHVSR_ATTN_STRIDED_IO). + # + # Identical math/masking to _bsfa_tma_kernel below; only the addressing + # differs. Q/K/V stay in the glue's natural token-major layout + # (S, H, D) == a contiguous 2D matrix (S, H*D): the descriptor covers that + # 2D view and each (head, block) tile is loaded at [row0, head*HEAD_DIM] + # instead of flattening to (H*S, D) — which is what forced the three + # (S,n,d)->(n,S,d) .contiguous() copies (11.6 ms/chunk, ANALYSIS §1.2/2.3). + # The output is written with explicit strides straight into an (S, H, D) + # tensor, killing the output transpose too. + # ----------------------------------------------------------------------- + @triton.jit + def _bsfa_tma_kernel_snd( + q_desc, k_desc, v_desc, O, KVIdx, KVCnt, sm_scale, + soh, som, sok, sih, sim, sic, sch, scm, H, N_Q, N_KV, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, + ): + start_m = tl.program_id(0) + off_h = tl.program_id(1) + col0 = off_h * HEAD_DIM + q = q_desc.load([start_m * BLOCK_M, col0]) # TMA tile @ (row, head-col) + qs = (q * sm_scale).to(q.dtype) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + offs_n = tl.arange(0, BLOCK_N) + cnt = tl.load(KVCnt + off_h * sch + start_m * scm) + base = KVIdx + off_h * sih + start_m * sim + for j in range(0, cnt): + kvb = tl.load(base + j * sic) + kk = k_desc.load([kvb * BLOCK_N, col0]) + qk = tl.dot(qs, kk.T) + n = kvb * BLOCK_N + offs_n + qk = tl.where(n[None, :] < N_KV, qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2((qk - m_ij[:, None]) * 1.44269504) + alpha = tl.math.exp2((m_i - m_ij) * 1.44269504) + l_i = l_i * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + vv = v_desc.load([kvb * BLOCK_N, col0]) + acc += tl.dot(p.to(vv.dtype), vv) + m_i = m_ij + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = tl.arange(0, HEAD_DIM) + tl.store(O + off_h * soh + offs_m[:, None] * som + offs_k[None, :] * sok, + acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_Q) + + @triton.jit + def _bsfa_tma_kernel( + q_desc, k_desc, v_desc, O, KVIdx, KVCnt, sm_scale, + soh, som, sok, sih, sim, sic, sch, scm, H, N_Q, N_KV, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, + ): + start_m = tl.program_id(0) + off_h = tl.program_id(1) + row0 = off_h * N_Q + start_m * BLOCK_M + q = q_desc.load([row0, 0]) # TMA bulk-load Q tile + qs = (q * sm_scale).to(q.dtype) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + offs_n = tl.arange(0, BLOCK_N) + cnt = tl.load(KVCnt + off_h * sch + start_m * scm) + base = KVIdx + off_h * sih + start_m * sim + kvbase = off_h * N_KV + for j in range(0, cnt): + kvb = tl.load(base + j * sic) + krow = kvbase + kvb * BLOCK_N + kk = k_desc.load([krow, 0]) # TMA bulk-load K tile + qk = tl.dot(qs, kk.T) + n = kvb * BLOCK_N + offs_n + qk = tl.where(n[None, :] < N_KV, qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2((qk - m_ij[:, None]) * 1.44269504) + alpha = tl.math.exp2((m_i - m_ij) * 1.44269504) + l_i = l_i * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + vv = v_desc.load([krow, 0]) # TMA bulk-load V tile + acc += tl.dot(p.to(vv.dtype), vv) + m_i = m_ij + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = tl.arange(0, HEAD_DIM) + tl.store(O + off_h * soh + offs_m[:, None] * som + offs_k[None, :] * sok, + acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_Q) + + +def _make_csr(bm): + """bm: (H, Nqb, Nkvb) bool -> (idx int32 (H,Nqb,Nkvb), cnt int32 (H,Nqb)).""" + cnt = bm.sum(-1).to(torch.int32) + idx = torch.argsort(bm.int(), dim=-1, descending=True, stable=True).to(torch.int32) + return idx.contiguous(), cnt.contiguous() + + +def _bsfa_tma(q, k, v, bm, sm_scale, BLOCK_M, BLOCK_N, num_warps, num_stages): + """TMA fast path. Requires contiguous (H,N,D); flattens to (H*N, D).""" + H, Nq, D = q.shape + Nkv = k.shape[1] + Nqb = triton.cdiv(Nq, BLOCK_M) + idx, cnt = _make_csr(bm) + o = torch.empty_like(q) + qf = q.reshape(H * Nq, D).contiguous() + kf = k.reshape(H * Nkv, D).contiguous() + vf = v.reshape(H * Nkv, D).contiguous() + q_desc = _TensorDescriptor.from_tensor(qf, [BLOCK_M, D]) + k_desc = _TensorDescriptor.from_tensor(kf, [BLOCK_N, D]) + v_desc = _TensorDescriptor.from_tensor(vf, [BLOCK_N, D]) + grid = (Nqb, H) + _bsfa_tma_kernel[grid]( + q_desc, k_desc, v_desc, o, idx, cnt, sm_scale, + o.stride(0), o.stride(1), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + H, Nq, Nkv, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, + num_warps=num_warps, num_stages=num_stages, + ) + return o + + +def triton_block_sparse_attention_snd(q, k, v, block_mask, sm_scale=None, + BLOCK_M=128, BLOCK_N=128, num_warps=8, + num_stages=2): + """Strided-IO WGMMA block-sparse attention (Phase 2A-3). + + q: (Nq, H, D), k/v: (Nkv, H, D) — the glue's natural token-major layout, + contiguous, WITHOUT per-head transpose copies. block_mask: (H, Nqb, Nkvb) + bool (True = compute). Returns (Nq, H, D) contiguous. + + Math (tile schedule, masking, accumulation order) is identical to + triton_block_sparse_attention; only tile addressing differs, so outputs + are expected bit-identical for the same inputs. + """ + assert _TRITON_OK, "triton not available" + Nq, H, D = q.shape + Nkv = k.shape[0] + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(D) + Nqb = triton.cdiv(Nq, BLOCK_M) + Nkvb = triton.cdiv(Nkv, BLOCK_N) + bm = block_mask[..., :Nqb, :Nkvb] + idx, cnt = _make_csr(bm) + o = torch.empty_like(q) # (Nq, H, D) — final layout + grid = (Nqb, H) + + if _USE_TMA and _TMA_OK and q.is_contiguous() and k.is_contiguous() \ + and v.is_contiguous(): + try: + # 2D (S, H*D) views are free on the contiguous (S, H, D) tensors. + q2 = q.view(Nq, H * D) + k2 = k.view(Nkv, H * D) + v2 = v.view(Nkv, H * D) + q_desc = _TensorDescriptor.from_tensor(q2, [BLOCK_M, D]) + k_desc = _TensorDescriptor.from_tensor(k2, [BLOCK_N, D]) + v_desc = _TensorDescriptor.from_tensor(v2, [BLOCK_N, D]) + _bsfa_tma_kernel_snd[grid]( + q_desc, k_desc, v_desc, o, idx, cnt, sm_scale, + o.stride(1), o.stride(0), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + H, Nq, Nkv, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, + num_warps=num_warps, num_stages=max(num_stages, 3), + ) + return o + except Exception: + torch.cuda.empty_cache() # fall through to the stride-general kernel + + # Non-TMA path: _bsfa_kernel is fully stride-parameterized, so transposed + # *views* (no copies) express the (H, N, D) indexing over (N, H, D) data. + qt, kt, vt, ot = (t.transpose(0, 1) for t in (q, k, v, o)) + _bsfa_kernel[grid]( + qt, kt, vt, ot, idx, cnt, sm_scale, + qt.stride(0), qt.stride(1), qt.stride(2), + kt.stride(0), kt.stride(1), kt.stride(2), + vt.stride(0), vt.stride(1), vt.stride(2), + ot.stride(0), ot.stride(1), ot.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + H, Nq, Nkv, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, + num_warps=num_warps, num_stages=num_stages, + ) + return o + + +def triton_block_sparse_attention(q, k, v, block_mask, sm_scale=None, + BLOCK_M=128, BLOCK_N=128, num_warps=8, num_stages=2): + """WGMMA block-sparse attention. + + q: (H, Nq, D), k/v: (H, Nkv, D), block_mask: (H, Nqb, Nkvb) bool (True=compute). + Block size is 128 (matches FlashVSR's mask granularity). Returns (H, Nq, D). + """ + assert _TRITON_OK, "triton not available" + H, Nq, D = q.shape + Nkv = k.shape[1] + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(D) + Nqb = triton.cdiv(Nq, BLOCK_M) + Nkvb = triton.cdiv(Nkv, BLOCK_N) + bm = block_mask[..., :Nqb, :Nkvb] + + # TMA fast path (Hopper): overlaps Q/K/V bulk loads with WGMMA (~+5% at the + # isolated kernel level; ~+2% end-to-end where it overlaps with other work). + if _USE_TMA and _TMA_OK: + try: + # TMA descriptors need num_stages>=3; SMEM caps BLOCK_N at 128. + return _bsfa_tma(q, k, v, bm, sm_scale, BLOCK_M, BLOCK_N, num_warps, + max(num_stages, 3)) + except Exception: + torch.cuda.empty_cache() # fall back to the non-TMA kernel below + + idx, cnt = _make_csr(bm) + o = torch.empty_like(q) + grid = (Nqb, H) + _bsfa_kernel[grid]( + q, k, v, o, idx, cnt, sm_scale, + q.stride(0), q.stride(1), q.stride(2), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + o.stride(0), o.stride(1), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + H, Nq, Nkv, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, + num_warps=num_warps, num_stages=num_stages, + ) + return o diff --git a/diffsynth/models/triton_block_sparse_attn_v2.py b/diffsynth/models/triton_block_sparse_attn_v2.py new file mode 100644 index 00000000..ba4dc5c7 --- /dev/null +++ b/diffsynth/models/triton_block_sparse_attn_v2.py @@ -0,0 +1,483 @@ +"""Hopper (sm_90) warp-specialized block-sparse FlashAttention v2 (Gluon). + +Phase 3 of the GH200 acceleration campaign. The v1 Triton kernel +(`triton_block_sparse_attn.py::_bsfa_tma_kernel_snd`) is a single-CTA-per-SM +software-pipelined loop: all 8 warps synchronize at every stage barrier, the +scheduler has no eligible warp 68.6% of cycles, and the tensor pipe idles +~60% (ncu: 2.0 ms/call, 40% SOL at the reference shape q 8448 x kv 33792, +h12 d128, density ~0.42-0.45). + +v2 restructures the SAME math into an FA3-style producer/consumer pipeline +using Triton's Gluon dialect: + + - producer warp group (4 warps, low registers): scalar-loads the selected + kv-block index list (exact CSR of the boolean block mask, unchanged from + v1, next index software-prefetched) and streams K/V tiles into an + NBUF-deep shared-memory ring via TMA, guarded by mbarriers; + - two consumer warpgroups (4+4 warps) owning 64-row halves of the q block + ("pingpong": while one warpgroup occupies the WGMMA pipe the other runs + its softmax), each additionally deferring its P.V wgmma by one iteration + (QK(j) is issued BEFORE PV(j-1); wgmma groups retire in order, so + wait(pendings=1) completes QK(j) while PV(j-1) drains under the softmax). + +Mask semantics are bit-identical to v1: the (H, Nqb, Nkvb) boolean block mask +is consumed through the exact same `_make_csr` (stable argsort -> ascending +kv-block order, degenerate all-false rows produce zero output rows through +the l==0 -> l_safe=1 guard). Numerics: the softmax scale is folded into the +exponent (exp2((qk - m) * scale*log2e)) like the reference FA2-style +`block_sparse_attn` kernel, instead of v1's bf16 q-prescale; both are gated +at cosine >= 0.9999 against `block_sparse_attn`. + +Forward-only, bf16, D=128, no dropout. Opt-in via +FLASHVSR_ATTN_BACKEND=triton2, sm_90-guarded; ANY failure at import/compile/ +launch raises to the caller, which falls back to the v1 triton path (then to +`block_sparse_attn`), preserving the established fallback ladder. +""" +import math +import os + +import torch + +from ..perf_stats import record as _perf_record + +_V2_OK = False +_V2_ERR = None +try: + import triton # noqa: F401 + from triton.experimental import gluon + import triton.experimental.gluon.language as ttgl + from triton.experimental.gluon.language.nvidia import hopper as _hop + from triton.experimental.gluon.language.nvidia.hopper import ( + tma as _tma, mbarrier as _mb) + from triton.experimental.gluon.nvidia.hopper import ( + TensorDescriptor as _GluonTensorDescriptor) + _V2_OK = True +except Exception as e: # pragma: no cover + _V2_ERR = e + +from .triton_block_sparse_attn import _make_csr + +# --------------------------------------------------------------------------- +# Phase 3.5-2: fused CSR build (FLASHVSR_FUSED_CSR, default OFF). +# +# `_make_csr` uses torch.argsort(descending, stable) — a full radix sort +# (radixSortKVInPlace ~0.71 ms/chunk in-pipe) — only to move the True +# kv-blocks to the front in ascending order. The v2 kernel reads only +# idx[0:cnt], so a single-pass cumsum-scatter reproduces idx[0:cnt] and cnt +# bit-identically without a sort. Lossless (idx[0:cnt] + cnt equal to the +# argsort path; the unused tail idx[cnt:] is never read). +# --------------------------------------------------------------------------- +_FUSED_CSR = os.environ.get("FLASHVSR_FUSED_CSR", "0") != "0" + +if _V2_OK: + import triton + import triton.language as tl + + @triton.jit + def _csr_scan_kernel(BM, IDX, CNT, R, NKVB, + sbr, sbn, sir, sin_, scr, + BLOCK: tl.constexpr): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + valid = offs < NKVB + m = tl.load(BM + row * sbr + offs * sbn, mask=valid, other=0).to(tl.int32) + # inclusive prefix sum -> write slot of each True block; ascending offs + # -> ascending slots -> idx[0:cnt] holds True kv-blocks in ascending + # order (== argsort(descending, stable)[0:cnt]). + pos = tl.cumsum(m, axis=0) - 1 + cnt = tl.sum(m, axis=0) + keep = (m > 0) & valid + tl.store(IDX + row * sir + pos * sin_, offs.to(tl.int32), mask=keep) + tl.store(CNT + row * scr, cnt) + + +def _make_csr_fused(bm): + """Sort-free CSR: (H,Nqb,Nkvb) bool -> (idx int32, cnt int32). + + idx[0:cnt] and cnt are bit-identical to `_make_csr`; the tail idx[cnt:] + (never consumed by the kernel) is left uninitialized. + """ + H, Nqb, Nkvb = bm.shape + R = H * Nqb + bmf = bm.reshape(R, Nkvb) + idx = torch.empty(R, Nkvb, dtype=torch.int32, device=bm.device) + cnt = torch.empty(R, dtype=torch.int32, device=bm.device) + BLOCK = max(16, triton.next_power_of_2(Nkvb)) + bmi = bmf.to(torch.int8) + _csr_scan_kernel[(R,)]( + bmi, idx, cnt, R, Nkvb, + bmi.stride(0), bmi.stride(1), idx.stride(0), idx.stride(1), cnt.stride(0), + BLOCK=BLOCK) + return idx.view(H, Nqb, Nkvb), cnt.view(H, Nqb) + + +# Ring depth for the K/V shared-memory pipeline. NBUF=3 fits in 227 KB with +# BLOCK_M=BLOCK_N=D=128 (Q 32K + 3x(K 32K + V 32K)); measured fastest +# (sweep in PHASE_BENCH_LOG Phase 3). +_NBUF = max(2, min(4, int(os.environ.get("FLASHVSR_ATTN_V2_NBUF", "3")))) +# Producer warp-group register budget (multiple of 8 in [24, 256]). +_PROD_REGS = int(os.environ.get("FLASHVSR_ATTN_V2_PROD_REGS", "40")) +# KV tile width. Fixed at 128 (= the mask block granularity). A 64-wide +# exact-refinement variant (2-CTA/SM residency) was evaluated during Phase 3 +# and rejected: it needs two distinct wgmma layouts (qk N=64, pv N=128) with +# per-iteration layout conversions, for an expected perf LOSS vs the +# measured-passing 12-warp warp-specialized config. +_BLOCK_N = 128 + + +if _V2_OK: + + @gluon.jit + def _producer(q_desc, k_desc, v_desc, KVIdx, cnt, row_q, col0, + q_smem, k_bufs, v_bufs, q_full, k_full, k_empty, + v_full, v_empty, v_f0, nhw, nw, + BLOCK_N: ttgl.constexpr, HEAD_DIM: ttgl.constexpr, + NBUF: ttgl.constexpr, RASTER_V: ttgl.constexpr): + # Q tile once (two 64-column TMA boxes -> one barrier). + QBYTES: ttgl.constexpr = q_smem.type.shape[0] * HEAD_DIM * 2 + _mb.expect(q_full, QBYTES) + _tma.async_copy_global_to_shared( + q_desc, [row_q, col0], q_full, q_smem.slice(0, 64, dim=1)) + _tma.async_copy_global_to_shared( + q_desc, [row_q, col0 + 64], q_full, q_smem.slice(64, 64, dim=1)) + KVBYTES: ttgl.constexpr = BLOCK_N * HEAD_DIM * 2 + # software-prefetch the selected-block index list: issue the load for + # j+1 before iteration j's barrier waits so the dependent-load latency + # (the dominant long-scoreboard stall) hides behind the ring handshake + kvb = ttgl.load(KVIdx, mask=cnt > 0, other=0) + for j in range(cnt): + buf = j % NBUF + phase = (j // NBUF) & 1 + kvb_next = ttgl.load(KVIdx + j + 1, mask=j + 1 < cnt, other=0) + row = kvb * BLOCK_N + _mb.wait(k_empty.index(buf), phase) + _mb.expect(k_full.index(buf), KVBYTES) + kb = k_bufs.index(buf) + _tma.async_copy_global_to_shared( + k_desc, [row, col0], k_full.index(buf), kb.slice(0, 64, dim=1)) + _tma.async_copy_global_to_shared( + k_desc, [row, col0 + 64], k_full.index(buf), + kb.slice(64, 64, dim=1)) + _mb.wait(v_empty.index(buf), phase) + _mb.expect(v_full.index(buf), KVBYTES) + vb = v_bufs.index(buf) + if RASTER_V: + # 5A-v1 zero-copy: gather the (2,8,8) window straight from the + # raster (F,H,W,Hh*D) V arena; box order == partition order + # (bit-exact, probe-verified). + slot = kvb // nhw + r = kvb % nhw + f0 = v_f0 + slot * 2 + h0 = (r // nw) * 8 + w0 = (r % nw) * 8 + _tma.async_copy_global_to_shared( + v_desc, [f0, h0, w0, col0], v_full.index(buf), + vb.slice(0, 64, dim=1)) + _tma.async_copy_global_to_shared( + v_desc, [f0, h0, w0, col0 + 64], v_full.index(buf), + vb.slice(64, 64, dim=1)) + else: + _tma.async_copy_global_to_shared( + v_desc, [row, col0], v_full.index(buf), + vb.slice(0, 64, dim=1)) + _tma.async_copy_global_to_shared( + v_desc, [row, col0 + 64], v_full.index(buf), + vb.slice(64, 64, dim=1)) + kvb = kvb_next + + @gluon.jit + def _consumer(O, cnt, scale_log2e, row_q, col_o, som, sok, + q_smem, k_bufs, v_bufs, q_full, k_full, k_empty, + v_full, v_empty, nhw, nw, hrast, wrast, + row_off: ttgl.constexpr, + HALF_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, + HEAD_DIM: ttgl.constexpr, NBUF: ttgl.constexpr, + RASTER_OUT: ttgl.constexpr): + # One consumer WARPGROUP owning a HALF_M-row horizontal stripe of the + # q block (FA3 "pingpong": two such warpgroups run the same loop on + # different stripes; while one occupies the WGMMA pipe the other runs + # its softmax, hiding the f32/MUFU chain). + mma: ttgl.constexpr = ttgl.NVMMADistributedLayout( + version=[3, 0], warps_per_cta=[4, 1], + instr_shape=[16, HEAD_DIM, 16]) + dot_a: ttgl.constexpr = ttgl.DotOperandLayout( + operand_index=0, parent=mma, k_width=2) + q_half = q_smem.slice(row_off, HALF_M, dim=0) + # Track the running max in the ALREADY-SCALED (log2e) domain: ms = m*s. + # Since s>0, maximum(m_i, rowmax)*s == maximum(ms_i, rowmax*s), so this + # is algebraically identical but lets the exponent be a single FFMA + # (qk*s - ms) instead of a broadcast-sub-then-mul ((qk-m)*s, which the + # compiler cannot contract). 3.5-1a; gated cos>=0.9999 + PSNR>=49. + ms_i = ttgl.full([HALF_M], -float("inf"), ttgl.float32, + ttgl.SliceLayout(1, mma)) + l_i = ttgl.zeros([HALF_M], ttgl.float32, ttgl.SliceLayout(1, mma)) + acc = ttgl.zeros([HALF_M, HEAD_DIM], ttgl.float32, mma) + qk0 = ttgl.zeros([HALF_M, BLOCK_N], ttgl.float32, mma) + # p tile of the previous iteration whose P.V wgmma is deferred so it + # overlaps this iteration's softmax (wgmma groups retire in order, so + # issuing QK(j) BEFORE PV(j-1) lets wait(pendings=1) complete QK(j) + # while PV(j-1) is still in flight). + p_prev = ttgl.zeros([HALF_M, BLOCK_N], ttgl.bfloat16, dot_a) + _mb.wait(q_full, 0) + if cnt > 0: + # peeled iteration 0: QK only; its P.V is deferred into the loop + _mb.wait(k_full.index(0), 0) + qk_tok = _hop.warpgroup_mma( + q_half, k_bufs.index(0).permute((1, 0)), qk0, + use_acc=False, is_async=True) + qk = _hop.warpgroup_mma_wait(0, deps=[qk_tok]) + _mb.arrive(k_empty.index(0), count=1) + ms_ij = ttgl.maximum(ms_i, ttgl.max(qk, 1) * scale_log2e) + p = ttgl.exp2(qk * scale_log2e - ttgl.expand_dims(ms_ij, 1)) + l_i = ttgl.sum(p, 1) + p_prev = ttgl.convert_layout(p.to(ttgl.bfloat16), dot_a) + ms_i = ms_ij + for j in range(1, cnt): + buf = j % NBUF + phase = (j // NBUF) & 1 + pbuf = (j - 1) % NBUF + pphase = ((j - 1) // NBUF) & 1 + _mb.wait(k_full.index(buf), phase) + qk_tok = _hop.warpgroup_mma( + q_half, k_bufs.index(buf).permute((1, 0)), qk0, + use_acc=False, is_async=True) + _mb.wait(v_full.index(pbuf), pphase) + acc = _hop.warpgroup_mma(p_prev, v_bufs.index(pbuf), acc, + is_async=True) + qk, acc = _hop.warpgroup_mma_wait(1, deps=[qk_tok, acc]) + _mb.arrive(k_empty.index(buf), count=1) + ms_ij = ttgl.maximum(ms_i, ttgl.max(qk, 1) * scale_log2e) + p = ttgl.exp2(qk * scale_log2e - ttgl.expand_dims(ms_ij, 1)) + alpha = ttgl.exp2(ms_i - ms_ij) + l_i = l_i * alpha + ttgl.sum(p, 1) + acc = _hop.warpgroup_mma_wait(0, deps=[acc]) + _mb.arrive(v_empty.index(pbuf), count=1) + acc = acc * ttgl.expand_dims(alpha, 1) + p_prev = ttgl.convert_layout(p.to(ttgl.bfloat16), dot_a) + ms_i = ms_ij + if cnt > 0: + lbuf = (cnt - 1) % NBUF + lphase = ((cnt - 1) // NBUF) & 1 + _mb.wait(v_full.index(lbuf), lphase) + acc = _hop.warpgroup_mma(p_prev, v_bufs.index(lbuf), acc, + is_async=True) + acc = _hop.warpgroup_mma_wait(0, deps=[acc]) + _mb.arrive(v_empty.index(lbuf), count=1) + l_safe = ttgl.where(l_i == 0.0, 1.0, l_i) + acc = acc / ttgl.expand_dims(l_safe, 1) + if RASTER_OUT: + # 5A-v1: store rows straight to RASTER token order (kills the + # win_rev reverse-partition pass). q block qb covers window + # (slot, hq, wq); row r in [0,128) covers (fl, hl, wl). + qb = row_q // 128 + slot = qb // nhw + rr = qb % nhw + hq = (rr // nw) * 8 + wq = (rr % nw) * 8 + r = row_off + ttgl.arange(0, HALF_M, ttgl.SliceLayout(1, mma)) + fl = r // 64 + hl = (r % 64) // 8 + wl = r % 8 + tok = ((slot * 2 + fl) * hrast + hq + hl) * wrast + wq + wl + offs_k = ttgl.arange(0, HEAD_DIM, ttgl.SliceLayout(0, mma)) + ptrs = O + col_o + (ttgl.expand_dims(tok, 1) * som + + ttgl.expand_dims(offs_k, 0) * sok) + ttgl.store(ptrs, acc.to(O.dtype.element_ty)) + else: + offs_m = row_q + row_off + ttgl.arange( + 0, HALF_M, ttgl.SliceLayout(1, mma)) + offs_k = ttgl.arange(0, HEAD_DIM, ttgl.SliceLayout(0, mma)) + ptrs = O + col_o + (ttgl.expand_dims(offs_m, 1) * som + + ttgl.expand_dims(offs_k, 0) * sok) + ttgl.store(ptrs, acc.to(O.dtype.element_ty)) + + @gluon.jit + def _bsfa_v2_kernel(q_desc, k_desc, v_desc, O, KVIdx, KVCnt, + scale_log2e, soh, som, sok, sih, sim, sic, sch, scm, + v_f0, nhw, nw, hrast, wrast, + BLOCK_M: ttgl.constexpr, BLOCK_N: ttgl.constexpr, + HEAD_DIM: ttgl.constexpr, NBUF: ttgl.constexpr, + PROD_REGS: ttgl.constexpr, + RASTER_V: ttgl.constexpr, RASTER_OUT: ttgl.constexpr): + start_m = ttgl.program_id(0) + off_h = ttgl.program_id(1) + cnt = ttgl.load(KVCnt + off_h * sch + start_m * scm) + idx_base = KVIdx + off_h * sih + start_m * sim + row_q = start_m * BLOCK_M + col0 = off_h * HEAD_DIM + + smem: ttgl.constexpr = ttgl.NVMMASharedLayout( + swizzle_byte_width=128, element_bitwidth=16, rank=2) + q_smem = ttgl.allocate_shared_memory( + ttgl.bfloat16, [BLOCK_M, HEAD_DIM], smem) + k_bufs = ttgl.allocate_shared_memory( + ttgl.bfloat16, [NBUF, BLOCK_N, HEAD_DIM], smem) + v_bufs = ttgl.allocate_shared_memory( + ttgl.bfloat16, [NBUF, BLOCK_N, HEAD_DIM], smem) + bl: ttgl.constexpr = _mb.MBarrierLayout() + q_full = ttgl.allocate_shared_memory(ttgl.int64, [1], bl) + k_full = ttgl.allocate_shared_memory(ttgl.int64, [NBUF, 1], bl) + k_empty = ttgl.allocate_shared_memory(ttgl.int64, [NBUF, 1], bl) + v_full = ttgl.allocate_shared_memory(ttgl.int64, [NBUF, 1], bl) + v_empty = ttgl.allocate_shared_memory(ttgl.int64, [NBUF, 1], bl) + _mb.init(q_full, count=1) + for i in ttgl.static_range(NBUF): + _mb.init(k_full.index(i), count=1) + # ring slots are released by BOTH consumer warpgroups + _mb.init(k_empty.index(i), count=2) + _mb.init(v_full.index(i), count=1) + _mb.init(v_empty.index(i), count=2) + # mark all ring slots initially empty (completes phase 0) + _mb.arrive(k_empty.index(i), count=2) + _mb.arrive(v_empty.index(i), count=2) + + col_o = off_h * soh + HALF_M: ttgl.constexpr = BLOCK_M // 2 + ttgl.warp_specialize( + [(_consumer, (O, cnt, scale_log2e, row_q, col_o, som, sok, + q_smem, k_bufs, v_bufs, q_full, k_full, k_empty, + v_full, v_empty, nhw, nw, hrast, wrast, + 0, HALF_M, BLOCK_N, HEAD_DIM, + NBUF, RASTER_OUT)), + (_consumer, (O, cnt, scale_log2e, row_q, col_o, som, sok, + q_smem, k_bufs, v_bufs, q_full, k_full, k_empty, + v_full, v_empty, nhw, nw, hrast, wrast, + HALF_M, HALF_M, BLOCK_N, HEAD_DIM, + NBUF, RASTER_OUT)), + (_producer, (q_desc, k_desc, v_desc, idx_base, cnt, row_q, col0, + q_smem, k_bufs, v_bufs, q_full, k_full, k_empty, + v_full, v_empty, v_f0, nhw, nw, + BLOCK_N, HEAD_DIM, NBUF, RASTER_V))], + [4, 4], [232, PROD_REGS], + ) + + +def _build_call(q, k, v, bm, sm_scale, BLOCK_M=128, BLOCK_N=None): + """Shared setup for the pipeline wrapper and the kernel-only bench hook.""" + if BLOCK_N is None: + BLOCK_N = _BLOCK_N + assert _V2_OK, f"gluon unavailable: {_V2_ERR}" + Nq, H, D = q.shape + Nkv = k.shape[0] + assert D == 128, "v2 kernel requires head_dim == 128" + assert Nq % BLOCK_M == 0 and Nkv % BLOCK_N == 0, \ + "v2 kernel requires 128-aligned sequence lengths" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + Nqb = Nq // BLOCK_M + Nkvb = Nkv // BLOCK_N + bm = bm[..., :Nqb, :Nkv // 128] + assert bm.shape == (H, Nqb, Nkv // 128) + if BLOCK_N != 128: + # exact refinement: every 128-wide mask block is covered by + # 128/BLOCK_N consecutive tiles inheriting the same mask bit + bm = bm.repeat_interleave(128 // BLOCK_N, dim=2) + if _FUSED_CSR: + idx, cnt = _make_csr_fused(bm) + _perf_record("csr_fused") + else: + idx, cnt = _make_csr(bm) + _perf_record("csr_argsort") + assert idx.stride(2) == 1 # producer walks the index list contiguously + o = torch.empty_like(q) + scale_log2e = float(sm_scale) * 1.4426950408889634 + q2 = q.view(Nq, H * D) + k2 = k.view(Nkv, H * D) + v2 = v.view(Nkv, H * D) + layout = ttgl.NVMMASharedLayout(swizzle_byte_width=128, + element_bitwidth=16, rank=2) + q_desc = _GluonTensorDescriptor.from_tensor(q2, [BLOCK_M, 64], layout) + k_desc = _GluonTensorDescriptor.from_tensor(k2, [BLOCK_N, 64], layout) + v_desc = _GluonTensorDescriptor.from_tensor(v2, [BLOCK_N, 64], layout) + grid = (Nqb, H) + + def call(): + _bsfa_v2_kernel[grid]( + q_desc, k_desc, v_desc, o, idx, cnt, scale_log2e, + o.stride(1), o.stride(0), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + 0, 1, 1, 1, 1, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, NBUF=_NBUF, + PROD_REGS=_PROD_REGS, RASTER_V=False, RASTER_OUT=False, + num_warps=4) + return o + return call + + +def triton_block_sparse_attention_v2(q, k, v, block_mask, sm_scale=None): + """Warp-specialized block-sparse attention (Phase 3). + + q: (Nq, H, D), k/v: (Nkv, H, D) token-major contiguous (same interface as + the v1 strided-IO path). block_mask: (H, Nqb, Nkvb) bool, True = compute. + Returns (Nq, H, D). Raises on any unsupported input; caller falls back. + """ + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(q.shape[-1]) + return _build_call(q, k, v, block_mask, sm_scale)() + + +def bsfa_v2_kernel_only(q, k, v, bm, sm_scale=None): + """Bench hook: returns a zero-setup-cost closure launching only the kernel.""" + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(q.shape[-1]) + return _build_call(q, k, v, bm, sm_scale) + + +def triton_block_sparse_attention_v2_zc(q, k, v_buf, v_f0, hrast, wrast, + block_mask, sm_scale=None): + """Phase 5A-v1 zero-copy variant: V is gathered straight from the RASTER + (F_cap, hrast, wrast, Hh*D) arena buffer (window (2,8,8) per kv-block, + starting at frame `v_f0`) and the output is stored straight in RASTER + token order (the caller skips WindowPartition3D.reverse). + + q: (Nq, H, D) window-token-major (as today), k: (Nkv, H, D) window-major. + Returns o: (Nq, H, D) in RASTER token order. Raises on any unsupported + input; caller falls back to the partitioned path. + """ + assert _V2_OK, f"gluon unavailable: {_V2_ERR}" + BLOCK_M = 128 + BLOCK_N = 128 + Nq, H, D = q.shape + Nkv = k.shape[0] + assert D == 128 and Nq % 128 == 0 and Nkv % 128 == 0 + assert q.is_contiguous() and k.is_contiguous() and v_buf.is_contiguous() + assert v_buf.dim() == 4 and v_buf.shape[1] == hrast and v_buf.shape[2] == wrast + assert v_buf.shape[3] == H * D + nh, nw = hrast // 8, wrast // 8 + nhw = nh * nw + Nqb = Nq // 128 + Nkvb = Nkv // 128 + assert Nkvb % nhw == 0 and Nqb % nhw == 0 + assert v_f0 + (Nkvb // nhw) * 2 <= v_buf.shape[0] + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(D) + bm = block_mask[..., :Nqb, :Nkvb] + assert bm.shape == (H, Nqb, Nkvb) + if _FUSED_CSR: + idx, cnt = _make_csr_fused(bm) + _perf_record("csr_fused") + else: + idx, cnt = _make_csr(bm) + _perf_record("csr_argsort") + assert idx.stride(2) == 1 + o = torch.empty_like(q) # (Nq, H, D), RASTER token order + scale_log2e = float(sm_scale) * 1.4426950408889634 + layout = ttgl.NVMMASharedLayout(swizzle_byte_width=128, + element_bitwidth=16, rank=2) + q_desc = _GluonTensorDescriptor.from_tensor( + q.view(Nq, H * D), [BLOCK_M, 64], layout) + k_desc = _GluonTensorDescriptor.from_tensor( + k.view(Nkv, H * D), [BLOCK_N, 64], layout) + v_desc = _GluonTensorDescriptor.from_tensor( + v_buf, [2, 8, 8, 64], layout) + grid = (Nqb, H) + _bsfa_v2_kernel[grid]( + q_desc, k_desc, v_desc, o, idx, cnt, scale_log2e, + o.stride(1), o.stride(0), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + v_f0, nhw, nw, hrast, wrast, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D, NBUF=_NBUF, + PROD_REGS=_PROD_REGS, RASTER_V=True, RASTER_OUT=True, + num_warps=4) + return o diff --git a/diffsynth/models/triton_dit_rows.py b/diffsynth/models/triton_dit_rows.py new file mode 100644 index 00000000..daa2a2f5 --- /dev/null +++ b/diffsynth/models/triton_dit_rows.py @@ -0,0 +1,89 @@ +"""Row-wise DiT elementwise kernels used by the optional Phase-7 fast path.""" + +import torch + +try: + import triton + import triton.language as tl + _TRITON_OK = True +except Exception: # pragma: no cover + _TRITON_OK = False + + +if _TRITON_OK: + + @triton.jit + def _layer_norm_modulate_kernel( + X, SHIFT, SCALE, OUT, + N: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr, + ): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK) + mask = cols < N + x = tl.load(X + row * N + cols, mask=mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / N + centered = x - mean + variance = tl.sum(centered * centered, axis=0) / N + normalized = centered * tl.rsqrt(variance + EPS) + shift = tl.load(SHIFT + cols, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(SCALE + cols, mask=mask, other=0.0).to(tl.float32) + tl.store(OUT + row * N + cols, normalized * (1.0 + scale) + shift, mask=mask) + + + @triton.jit + def _gated_residual_kernel( + X, GATE, RESIDUAL, OUT, TOTAL, + N: tl.constexpr, BLOCK: tl.constexpr, + ): + block = tl.program_id(0) + offsets = block * BLOCK + tl.arange(0, BLOCK) + mask = offsets < TOTAL + cols = offsets % N + x = tl.load(X + offsets, mask=mask, other=0.0).to(tl.float32) + residual = tl.load(RESIDUAL + offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(GATE + cols, mask=mask, other=0.0).to(tl.float32) + tl.store(OUT + offsets, x + gate * residual, mask=mask) + + +def _check_bf16_rows(x, *vectors): + if not _TRITON_OK: + raise RuntimeError("Triton is unavailable") + if x.device.type != "cuda" or x.dtype != torch.bfloat16 or not x.is_contiguous(): + raise ValueError("expected contiguous CUDA bf16 rows") + if x.ndim < 2: + raise ValueError("expected a tensor with a row dimension") + width = x.shape[-1] + if width > 4096: + raise ValueError(f"unsupported row width {width}") + for vector in vectors: + if vector.device != x.device or vector.dtype != x.dtype or vector.numel() != width: + raise ValueError("modulation vector does not match the input rows") + return width + + +def layer_norm_modulate_triton(x, shift, scale, eps): + """Fused affine-free LayerNorm followed by AdaLN modulation.""" + width = _check_bf16_rows(x, shift, scale) + rows = x.numel() // width + out = torch.empty_like(x) + block = triton.next_power_of_2(width) + _layer_norm_modulate_kernel[(rows,)]( + x, shift.reshape(-1), scale.reshape(-1), out, + N=width, EPS=eps, BLOCK=block, num_warps=8, + ) + return out + + +def gated_residual_triton(x, gate, residual): + """Fused ``x + gate * residual`` for row-broadcast AdaLN gates.""" + width = _check_bf16_rows(x, gate) + if residual.shape != x.shape or residual.dtype != x.dtype or not residual.is_contiguous(): + raise ValueError("residual does not match the input rows") + out = torch.empty_like(x) + total = x.numel() + block = 2048 + _gated_residual_kernel[(triton.cdiv(total, block),)]( + x, gate.reshape(-1), residual, out, total, + N=width, BLOCK=block, num_warps=8, + ) + return out diff --git a/diffsynth/models/triton_rope.py b/diffsynth/models/triton_rope.py new file mode 100644 index 00000000..befb8cb2 --- /dev/null +++ b/diffsynth/models/triton_rope.py @@ -0,0 +1,81 @@ +"""Hand-written coalesced RoPE apply kernel (FLASHVSR_ROPE_KERNEL=triton). + +The Phase-2A-1b fused RoPE (torch.compile) is bit-exact vs eager but reaches +only ~484 GB/s (143 us/call at the 8448-token steady shape): it reads +`freqs.real/.imag` as stride-2 fp64 views of the complex128 tensor (each 8 B +element pulls a 16 B line) and inductor's codegen for the interleaved +(..., dc, 2) stack is scalar-ish. RoPE apply is a pure elementwise complex +multiply — no reductions — so a hand kernel can reproduce the EXACT same fp64 +operations per element (o_r = xr*fr - xi*fi; o_i = xr*fi + xi*fr, then a +single fp64->bf16 round, the same cast the inductor kernel performs) while +moving all data as packed 32-bit words: + + * x / out: bf16 pairs loaded/stored as ONE u32 per complex pair (fully + coalesced; bf16->fp32 by bit-shift is exact), + * freqs: read directly from the contiguous complex128 storage as adjacent + fp64 (re, im) pairs — no strided .real/.imag views. + +Ideal traffic ~61 MB/call -> ~20-30 us at HBM3 rates (vs 143 us). +Bit-exactness is gated (kernel + E2E max|diff| == 0 vs the FUSE_ROPE path). +Default OFF; any failure falls back to the existing fused/eager paths. +""" +import torch + +try: + import triton + import triton.language as tl + _TRITON_OK = True +except Exception: # pragma: no cover + _TRITON_OK = False + + +if _TRITON_OK: + + @triton.jit + def _rope_u32_kernel(XU, F, OU, TOTAL, S, + ND: tl.constexpr, DC: tl.constexpr, + BLOCK: tl.constexpr): + pid = tl.program_id(0) + p = pid * BLOCK + tl.arange(0, BLOCK) + m = p < TOTAL + row = p // ND # token row (b*S + s) + s = row % S # freqs row + c = p % DC # complex lane within a head + # x pair: one u32 = (bf16 imag << 16) | bf16 real + xv = tl.load(XU + p, mask=m, other=0).to(tl.uint32, bitcast=True) + xr = ((xv & 0xFFFF) << 16).to(tl.float32, bitcast=True).to(tl.float64) + xi = ((xv >> 16) << 16).to(tl.float32, bitcast=True).to(tl.float64) + # freqs pair: adjacent fp64 (re, im) in the complex128 storage + fo = (s * DC + c) * 2 + fr = tl.load(F + fo, mask=m, other=0.0) + fi = tl.load(F + fo + 1, mask=m, other=0.0) + # same fp64 expressions as the fused impl; torch casts fp64->bf16 + # THROUGH fp32 (c10::BFloat16 is constructed from float), so match + # that double-rounding chain exactly for bit-equality + o_r = (xr * fr - xi * fi).to(tl.float32).to(tl.bfloat16) + o_i = (xr * fi + xi * fr).to(tl.float32).to(tl.bfloat16) + r16 = o_r.to(tl.uint16, bitcast=True).to(tl.uint32) + i16 = o_i.to(tl.uint16, bitcast=True).to(tl.uint32) + tl.store(OU + p, (r16 | (i16 << 16)).to(tl.int32, bitcast=True), mask=m) + + +def rope_apply_triton(x, freqs, num_heads): + """Drop-in for rope_apply's fused path. x: (B,S,n*d) bf16 contiguous, + freqs: (S,1,d//2) complex128 contiguous. Raises on unsupported input; + the caller falls back to the fused/eager paths.""" + assert _TRITON_OK + B, S, D = x.shape + n = num_heads + dc = (D // n) // 2 + assert x.dtype == torch.bfloat16 and x.is_contiguous() + assert freqs.dtype == torch.complex128 and freqs.is_contiguous() + assert freqs.shape[0] == S and freqs.shape[-1] == dc + o = torch.empty_like(x) + xu = x.view(torch.int32) # (B,S,D/2) u32 pairs, same storage + ou = o.view(torch.int32) + fv = torch.view_as_real(freqs) # (S,1,dc,2) fp64 view of the storage + total = B * S * n * dc + BLOCK = 2048 + _rope_u32_kernel[(triton.cdiv(total, BLOCK),)]( + xu, fv, ou, total, S, ND=n * dc, DC=dc, BLOCK=BLOCK, num_warps=8) + return o diff --git a/diffsynth/models/wan_video_dit.py b/diffsynth/models/wan_video_dit.py index a69e2bb2..93ecd607 100644 --- a/diffsynth/models/wan_video_dit.py +++ b/diffsynth/models/wan_video_dit.py @@ -8,6 +8,10 @@ from typing import Tuple, Optional, List from einops import rearrange from .utils import hash_state_dict_keys +from ..nvtx_utils import nvtx_range +from ..perf_stats import record as _perf_record, record_error as _perf_error +# Phase 2B-2: FP8 GEMM infrastructure (FLASHVSR_FP8_GEMM, default OFF). +from . import fp8_gemm as _fp8g try: import flash_attn_interface @@ -31,6 +35,206 @@ from PIL import Image import numpy as np +# --------------------------------------------------------------------------- +# Hopper adaptive attention backend. +# +# Measured on GH200 @768x1408 (seq=25344, 12 heads, dim=128): the block-sparse +# self-attention runs at block-mask density ~0.606 and takes ~7.3 ms, while +# cuDNN's fused dense SDPA computes the FULL attention in ~6.5 ms (605 TFLOP/s). +# i.e. at high density the sparse kernel is slower than dense AND drops context. +# +# When the block mask is dense enough (density >= threshold) we route to cuDNN +# fused dense attention instead of the sparse kernel: faster and uses full +# context. Below the threshold the sparse kernel wins, so we keep it. +# +# Knob: FLASHVSR_ATTN_BACKEND = sparse | triton | triton2 | auto | dense +# sparse -> always block_sparse (DEFAULT = original behaviour, no quality change) +# triton -> Hopper WGMMA block-sparse kernel: exact same mask, output matches +# block_sparse very closely (~49.7 dB PSNR end-to-end; the only +# difference is WGMMA vs HMMA accumulation order). ~1.2x faster than +# block_sparse at the kernel level (~+9% end-to-end). sm_90 only; +# silently falls back to block_sparse elsewhere / on error. +# triton2 -> Phase-3 warp-specialized (producer/consumer + pingpong) Gluon +# kernel: exact same mask/CSR, ~1.5x the v1 kernel at the steady +# shape. sm_90 only; falls back v2 -> v1 triton -> block_sparse. +# auto -> density-adaptive: dense if density>=FLASHVSR_ATTN_DENSE_THRESH +# dense -> always cuDNN fused dense +# FLASHVSR_ATTN_DENSE_THRESH default 0.5 (the measured crossover point). +# +# NOTE: routing to dense changes the output (it uses FULL attention instead of +# the trained locality-constrained sparse pattern, which measurably lowers PSNR), +# and at the default density the E2E gain is negligible, so the default stays +# 'sparse'. The knob exists for future aggressive-sparsity experiments (lower +# topk -> lower density -> sparse wins big, see docs). +# --------------------------------------------------------------------------- +_ATTN_BACKEND = os.environ.get("FLASHVSR_ATTN_BACKEND", "sparse").lower() +_ATTN_DENSE_THRESH = float(os.environ.get("FLASHVSR_ATTN_DENSE_THRESH", "0.5")) +_DENSE_FAILED = False +_ATTN_V2_FAILED = False +_ATTN_V1_STRIDED_FAILED = False +_ATTN_V1_FAILED = False + +# --------------------------------------------------------------------------- +# Norm / elementwise fusion (Phase 3-B). +# +# The DiT block is dominated (~17% of denoise GPU time) by memory-bound +# elementwise kernels: RMSNorm (q/k, with an fp32 up/down cast), LayerNorm, +# modulate (x*(1+scale)+shift) and the gate (x + gate*residual). torch.compile +# fuses each chain into a single kernel; in isolation the fused RMSNorm is +# several times faster than the eager path (the exact ratio depends on the +# tensor shape), at near-identical precision (cos ~1.0). +# These functions contain no attention / custom kernels, so compiling them does +# not interact with the block_sparse path or the streaming cache. End-to-end the +# fusion is measured at ~49 dB PSNR vs the unfused path (see test_fuse_norm.py). +# +# Knob FLASHVSR_FUSE_NORM = 0 | 1 (default off; opt-in, parity-gated). +# --------------------------------------------------------------------------- +_FUSE_NORM = os.environ.get("FLASHVSR_FUSE_NORM", "0") != "0" + +# Phase 7-A: fuse the two affine-free LayerNorm -> AdaLN chains in every DiT +# block. The gate kernel has the same row-broadcast contract. Both are optional +# and fall back to the existing torch.compile path if a runtime shape is not +# supported. +_DIT_ROW_FUSION = os.environ.get("FLASHVSR_DIT_ROW_FUSION", "0") != "0" +_DIT_ROW_FUSION_FAILED = False +try: + from .triton_dit_rows import ( + layer_norm_modulate_triton as _DIT_LN_MODULATE, + gated_residual_triton as _DIT_GATED_RESIDUAL, + ) +except Exception: + _DIT_LN_MODULATE = None + _DIT_GATED_RESIDUAL = None + +# --------------------------------------------------------------------------- +# Lossless step-invariant caches (Phase B). +# +# Several per-block elementwise results are recomputed on every denoise step but +# depend only on fixed inputs (parameters + the once-computed timestep modulation +# t_mod), NOT on x / q / k. Caching them is bit-identical (max|diff|==0). +# +# B1 FLASHVSR_CACHE_MOD = 0 | 1 -> cache (modulation + t_mod).chunk(...) +# per DiTBlock / Head. +# B2 FLASHVSR_CACHE_MASK_BIAS = 0 | 1 -> cache the local_attn_mask additive bias +# (0/-inf) in generate_draft_block_mask. +# Both default OFF (opt-in), pure elementwise, silent fallback. They do not touch +# attention / the streaming cache, so they compose with every other knob. +# --------------------------------------------------------------------------- +_CACHE_MOD = os.environ.get("FLASHVSR_CACHE_MOD", "0") != "0" +_CACHE_MASK_BIAS = os.environ.get("FLASHVSR_CACHE_MASK_BIAS", "0") != "0" + +# --------------------------------------------------------------------------- +# Phase 2A-4: mask-generation allocation/sync cleanup (FLASHVSR_MASKGEN_LEAN, +# default OFF). Exact-semantics only — the produced boolean mask is identical: +# (a) threshold via torch.kthvalue(n-k) instead of topk(k+1).values[:,-1] +# — the same order statistic (ties included), computed by one +# radix-select kernel without materializing the (rows, k+1) values + +# int64 indices tensors (topk here selects ~45% of 17k elements/row); +# (b) drop the no-op `.repeat(1,1,1,1)` copy of the boolean mask (B==1 is +# asserted in generate_draft_block_mask); +# (c) sparse backend only: cache the cu_seqlens_q/k + head_mask_type int32 +# tensors keyed on (seqlen, seqlen_kv, heads, device) — their per-call +# `torch.tensor(..., device=...)` construction is a hidden H2D sync +# (ANALYSIS §3, sparse-backend idle). +# --------------------------------------------------------------------------- +_MASKGEN_LEAN = os.environ.get("FLASHVSR_MASKGEN_LEAN", "0") != "0" + +# Phase 7-B prototype: the global order-statistic dominates mask generation +# once the KV window has reached its steady shape. Reuse the previous chunk's +# exact threshold per block and shape; the cache is reset at each new video. +# This changes the selected mask, so it is quality-gated rather than lossless. +_MASKGEN_THRESHOLD_CACHE = ( + os.environ.get("FLASHVSR_MASKGEN_THRESHOLD_CACHE", "0") != "0") + +# --------------------------------------------------------------------------- +# Phase 5-P2: incremental pooled-K cache (FLASHVSR_POOLED_K_CACHE, default +# OFF). mask_gen re-pools ALL kv windows (torch.mean over 128 tokens) every +# DiT block every chunk (~1.6 ms/chunk), but old windows' K values are +# bit-identical across chunks (arena/cat copies) and torch.mean(dim=1) is +# per-row deterministic INDEPENDENT of the batch row count (verified +# bit-equal: (264,128,C) one-call == 4x(66,128,C) slices). So pool each +# window ONCE (when new) and carry the pooled rows in a per-block rolling +# buffer that mirrors the KV window's exact cat/trim sequence -> bit-exact +# (E2E max|diff|==0 gated). Self-heals on any length mismatch by re-pooling +# everything (new video, arena reset, shape change). +# --------------------------------------------------------------------------- +_POOLED_K_CACHE = os.environ.get("FLASHVSR_POOLED_K_CACHE", "0") != "0" + +# (c): persistent per-shape int32 tensors for the block_sparse_attn call. +_SPARSE_SEQLENS_CACHE = {} + + +def _maybe_compile(fn): + if not _FUSE_NORM: + return fn + try: + return torch.compile(fn, dynamic=True) + except Exception: + return fn + +try: + from torch.nn.attention import sdpa_kernel as _sdpa_kernel, SDPBackend as _SDPBackend + _CUDNN_SDPA_OK = True +except Exception: + _CUDNN_SDPA_OK = False + +# Triton WGMMA block-sparse attention kernel (optional; sm_90 fast path). +try: + from .triton_block_sparse_attn import triton_block_sparse_attention as _TRITON_BSA +except Exception: + _TRITON_BSA = None + +# --------------------------------------------------------------------------- +# Phase 2A-3: strided attention IO (FLASHVSR_ATTN_STRIDED_IO, default OFF). +# +# The triton-backend glue below performs three (S,n,d)->(n,S,d) .contiguous() +# transposes for q/k/v plus a transpose of the output — ~11.6 ms/chunk of +# SM-bound strided copies (ANALYSIS §1.2/§2.3). The strided-IO variant keeps +# tensors in the glue's natural token-major (S, n, d) layout end to end: +# TMA descriptors address per-head tiles inside the 2D (S, n*d) view, and the +# kernel stores the output directly in (S, n, d). Same tile schedule and +# accumulation order -> gated on max|diff| kernel-level + E2E PSNR >= 49 dB. +# On any failure the glue falls back to the contiguous triton path (NOT to +# the sparse backend), preserving the existing fallback ladder. +# --------------------------------------------------------------------------- +_ATTN_STRIDED_IO = os.environ.get("FLASHVSR_ATTN_STRIDED_IO", "0") != "0" +try: + from .triton_block_sparse_attn import triton_block_sparse_attention_snd as _TRITON_BSA_SND +except Exception: + _TRITON_BSA_SND = None + +# --------------------------------------------------------------------------- +# Phase 3: warp-specialized block-sparse attention v2 +# (FLASHVSR_ATTN_BACKEND=triton2, sm_90 + Gluon only). +# +# FA3-style producer/consumer rewrite of the block-sparse kernel: a TMA +# producer warp group streams the CSR-selected K/V blocks into a ring buffer +# while two 64-row consumer warpgroups pingpong softmax against WGMMA +# (12 warps/SM vs v1's barrier-locked 8). Exact same boolean mask, same CSR, +# same (S, n, d) glue interface as the strided-IO path. Fallback ladder on +# ANY failure: v2 -> v1 triton (strided -> contiguous) -> block_sparse_attn. +# --------------------------------------------------------------------------- +try: + from .triton_block_sparse_attn_v2 import ( + triton_block_sparse_attention_v2 as _TRITON_BSA_V2) +except Exception: + _TRITON_BSA_V2 = None + + +def _is_hopper_dev(device): + try: + return device.type == "cuda" and torch.cuda.get_device_capability(device) == (9, 0) + except Exception: + return False + + +def _dense_sdpa(q_bnsd): + """cuDNN fused dense attention on (B, heads, S, D) layout.""" + if _CUDNN_SDPA_OK: + with _sdpa_kernel(_SDPBackend.CUDNN_ATTENTION): + return F.scaled_dot_product_attention(*q_bnsd) + return F.scaled_dot_product_attention(*q_bnsd) + # ---------------------------- # Local / window masks @@ -112,13 +316,175 @@ def reverse(windows: torch.Tensor, win: Tuple[int, int, int], orig: Tuple[int, i return x.view(B, F, H, W, -1) +# --------------------------------------------------------------------------- +# Phase 2A-2: KV cache arena (FLASHVSR_KV_RINGBUF, default OFF). +# +# The streaming self-attention KV cache is maintained today as +# k_w = torch.cat([pre_cache_k, k_w_new], dim=0) # copies the WHOLE window +# cache_k = k_w[one_len:] # trim = slice view +# i.e. every chunk copies pre+new (~264 window-blocks @768) per tensor per +# block at HBM bandwidth (kv_cat = 3.4 ms/chunk, ANALYSIS §1.2) even though +# only `one_len` (~66) windows are new. +# +# The arena replaces this with a preallocated per-block buffer of +# (kv_len + 1 + SPARE) temporal slots × one_len windows: +# - new windows are partition-written directly into the tail (this replaces +# the .contiguous() materialization inside WindowPartition3D.partition, so +# it adds no traffic), +# - the live KV window is the contiguous slice buf[start : start+length] — +# same values, same temporal order, same contiguity as the cat result, +# - trimming the oldest slot advances `start` (no copy), +# - when the tail reaches capacity the live window is compacted to offset 0 +# (amortized once every SPARE chunks; overlap-safe). +# Net: ~264 → ~66+198/(SPARE+1) window-copies per chunk per tensor, at the +# cost of SPARE extra slots of retained memory per tensor. +# +# Values and ordering are bit-identical to the cat path (copies only, no +# arithmetic), so this is gated by max|diff|==0 across a full multi-chunk +# clip (exercises rotation + compaction). Works with every attention backend +# (the consumed view is contiguous, exactly like the cat result). +# --------------------------------------------------------------------------- +_KV_RINGBUF = os.environ.get("FLASHVSR_KV_RINGBUF", "0") != "0" +_KV_RINGBUF_SPARE = max(1, int(os.environ.get("FLASHVSR_KV_RINGBUF_SPARE", "2"))) + +# --------------------------------------------------------------------------- +# Phase 5A-v1: zero-copy V windowing (FLASHVSR_ATTN_ZEROCOPY, default OFF, +# requires FLASHVSR_ATTN_BACKEND=triton2). +# +# V never feeds the mask-gen pooling, so its window-partition permute-copy +# (arena append) and the output's reverse-partition (win_rev) are pure layout +# work. The v2 kernel gathers each (2,8,8) V window straight from a RASTER +# (F,h,w,C) ring arena via a rank-4 TMA box (bit-exact vs the partition path, +# probe-verified) and stores the attention output straight in raster token +# order. K and Q keep the partitioned layout (their pooled means feed the +# exact-mask top-k; re-pooling from raster would change reduction order). +# On ANY zero-copy failure the raster live view is partitioned on the fly and +# the standard ladder (triton2 -> v1 -> sparse) runs; V's cache slot then +# stays a _VRasterArena for the rest of the video (correct, slightly slower). +# Gated lossless: E2E max|diff| == 0 incl. ring rotation + compaction. +# --------------------------------------------------------------------------- +_ATTN_ZEROCOPY = os.environ.get("FLASHVSR_ATTN_ZEROCOPY", "0") != "0" +try: + from .triton_block_sparse_attn_v2 import ( + triton_block_sparse_attention_v2_zc as _ZC_ATTN) +except Exception: + _ZC_ATTN = None + + +class _VRasterArena: + """Raster (frames, h, w, C) sliding ring for V (5A-v1 zero-copy).""" + __slots__ = ("buf", "start", "length") + + def __init__(self, kv_len, h, w, C, dtype, device): + cap = (kv_len + 1 + _KV_RINGBUF_SPARE) * 2 # frames (2 per slot) + self.buf = torch.empty(cap, h, w, C, dtype=dtype, device=device) + self.start = 0 + self.length = 0 + + def append(self, x): + """x: (1, f, h, w, C) raster. Contiguous frame copy (no permute).""" + f = x.shape[1] + cap = self.buf.shape[0] + if self.start + self.length + f > cap: + _perf_record("v_arena_compaction") + with nvtx_range("kv_cat"): + if self.start < self.length: + tmp = self.buf[self.start:self.start + self.length].clone() + self.buf[:self.length].copy_(tmp) + else: + self.buf[:self.length].copy_( + self.buf[self.start:self.start + self.length]) + self.start = 0 + self.buf[self.start + self.length: + self.start + self.length + f].copy_(x[0]) + self.length += f + + def trim(self, frames=2): + self.start += frames + self.length -= frames + + @property + def live(self): + return self.buf[self.start:self.start + self.length] + + +class _KVArena: + """Sliding-window KV arena; flows through the pre_cache_k/v slots.""" + __slots__ = ("buf", "start", "length", "one_len") + + def __init__(self, kv_len, one_len, block_s, dim, dtype, device): + cap = (kv_len + 1 + _KV_RINGBUF_SPARE) * one_len + self.buf = torch.empty(cap, block_s, dim, dtype=dtype, device=device) + self.start = 0 + self.length = 0 + self.one_len = one_len + + def append_partition(self, x, win): + """Window-partition `x` (B=1, f, h, w, C) directly into the arena tail + and return the live contiguous view buf[start : start+length]. + + The write performs exactly the permute-copy that + WindowPartition3D.partition's .contiguous() would perform, with the + arena slice as destination -> bit-identical values/order.""" + B, F_, H_, W_, C = x.shape + wf, wh, ww = win + n_new = (F_ // wf) * (H_ // wh) * (W_ // ww) + cap = self.buf.shape[0] + if self.start + self.length + n_new > cap: + # Compact the live window to offset 0 (kv_cat residual, amortized). + _perf_record("kv_arena_compaction") + with nvtx_range("kv_cat"): + if self.start < self.length: + # overlapping ranges: stage through a temp (copy_ on + # overlapping storage is undefined). Only possible with + # very tight SPARE settings. + tmp = self.buf[self.start:self.start + self.length].clone() + self.buf[:self.length].copy_(tmp) + else: + self.buf[:self.length].copy_( + self.buf[self.start:self.start + self.length]) + self.start = 0 + dst = self.buf[self.start + self.length: + self.start + self.length + n_new] + src = x.view(B, F_ // wf, wf, H_ // wh, wh, W_ // ww, ww, C) \ + .permute(0, 1, 3, 5, 2, 4, 6, 7) + dst.view(B, F_ // wf, H_ // wh, W_ // ww, wf, wh, ww, C).copy_(src) + self.length += n_new + return self.buf[self.start:self.start + self.length] + + def trim(self, n_windows): + """Drop the oldest n_windows (pointer advance, no data movement).""" + self.start += n_windows + self.length -= n_windows + + +# B2: process-wide cache for the geometry-only additive attention bias. +_MASK_BIAS_CACHE = {} + + +def _build_mask_bias(local_attn_mask, repeat_head, repeat_len, repeat_num): + """Build the (repeat_head, S, S) 0/-inf additive bias from the boolean local + block mask. Exact re-expression of the original inline code (lines below) so + cached and uncached paths are bit-identical.""" + m = local_attn_mask.unsqueeze(1).unsqueeze(0).repeat(repeat_len, 1, repeat_num, 1) + m = rearrange(m, 'x a y b -> (x a) (y b)') + m = m.unsqueeze(0).repeat(repeat_head, 1, 1) + m = m.to(torch.float32) + m = m.masked_fill(m == False, -float('inf')) + m = m.masked_fill(m == True, 0) + return m + + @torch.no_grad() def generate_draft_block_mask(batch_size, nheads, seqlen, - q_w, k_w, topk=10, local_attn_mask=None): + q_w, k_w, topk=10, local_attn_mask=None, + avgpool_k=None, cached_thresholds=None, + return_thresholds=False): assert batch_size == 1, "Only batch_size=1 supported for now" assert local_attn_mask is not None, "local_attn_mask must be provided" avgpool_q = torch.mean(q_w, dim=1) - avgpool_k = torch.mean(k_w, dim=1) + if avgpool_k is None: + avgpool_k = torch.mean(k_w, dim=1) avgpool_q = rearrange(avgpool_q, 's (h d) -> s h d', h=nheads) avgpool_k = rearrange(avgpool_k, 's (h d) -> s h d', h=nheads) q_heads = avgpool_q.permute(1, 0, 2) @@ -129,13 +495,22 @@ def generate_draft_block_mask(batch_size, nheads, seqlen, repeat_head = scores.shape[0] repeat_len = scores.shape[1] // local_attn_mask.shape[0] repeat_num = scores.shape[2] // local_attn_mask.shape[1] - local_attn_mask = local_attn_mask.unsqueeze(1).unsqueeze(0).repeat(repeat_len, 1, repeat_num, 1) - local_attn_mask = rearrange(local_attn_mask, 'x a y b -> (x a) (y b)') - local_attn_mask = local_attn_mask.unsqueeze(0).repeat(repeat_head, 1, 1) - local_attn_mask = local_attn_mask.to(torch.float32) - local_attn_mask = local_attn_mask.masked_fill(local_attn_mask == False, -float('inf')) - local_attn_mask = local_attn_mask.masked_fill(local_attn_mask == True, 0) - scores = scores + local_attn_mask + + # B2 lossless cache: the additive bias (0 / -inf) depends only on the geometry + # (local_attn_mask + repeat factors), not on q/k, so it is identical across + # every block and every step. Recomputing it (repeat + 2x masked_fill + cast) + # each call is pure overhead. Cache it keyed on those shape-only inputs. + bias = None + if _CACHE_MASK_BIAS: + key = (id(local_attn_mask), repeat_head, repeat_len, repeat_num, + local_attn_mask.device, scores.shape[1], scores.shape[2]) + bias = _MASK_BIAS_CACHE.get(key) + if bias is None: + bias = _build_mask_bias(local_attn_mask, repeat_head, repeat_len, repeat_num) + _MASK_BIAS_CACHE[key] = bias + if bias is None: + bias = _build_mask_bias(local_attn_mask, repeat_head, repeat_len, repeat_num) + scores = scores + bias attn_map = torch.softmax(scores, dim=-1) attn_map = rearrange(attn_map, 'h (it s1) s2 -> (h it) s1 s2', it=seqlen) @@ -143,14 +518,37 @@ def generate_draft_block_mask(batch_size, nheads, seqlen, flat = attn_map.reshape(loop_num, -1) n = flat.shape[1] apply_topk = min(flat.shape[1]-1, topk) - thresholds = torch.topk(flat, k=apply_topk + 1, dim=1, largest=True).values[:, -1] - thresholds = thresholds.unsqueeze(1) + use_cached_thresholds = ( + cached_thresholds is not None + and cached_thresholds.shape == (loop_num,) + and cached_thresholds.device == flat.device + and cached_thresholds.dtype == flat.dtype + ) + if use_cached_thresholds: + threshold_values = cached_thresholds + _perf_record("mask_threshold_cached") + elif _MASKGEN_LEAN: + # 2A-4(a): (apply_topk+1)-th largest == (n-apply_topk)-th smallest. + # Identical exact value (order statistic, ties and all); single + # radix-select kernel, no (rows, k+1) values/indices materialization. + threshold_values = torch.kthvalue(flat, n - apply_topk, dim=1).values + _perf_record("mask_threshold_kthvalue") + else: + threshold_values = torch.topk( + flat, k=apply_topk + 1, dim=1, largest=True).values[:, -1] + _perf_record("mask_threshold_topk") + thresholds = threshold_values.unsqueeze(1) mask_new = (flat > thresholds).reshape(loop_num, s1, s2) mask_new = rearrange(mask_new, '(h it) s1 s2 -> h (it s1) s2', it=seqlen) # keep shape note # 修正：上行变量名统一 # mask_new = rearrange(attn_map, 'h (it s1) s2 -> h (it s1) s2', it=seqlen) * 0 + mask_new - mask = mask_new.unsqueeze(0).repeat(batch_size, 1, 1, 1) - return mask + if _MASKGEN_LEAN and batch_size == 1: + # 2A-4(b): batch_size==1 (asserted above) -> repeat(1,1,1,1) is a + # full copy with no semantic effect; a view is enough. + mask = mask_new.unsqueeze(0) + else: + mask = mask_new.unsqueeze(0).repeat(batch_size, 1, 1, 1) + return (mask, threshold_values) if return_thresholds else mask @torch.no_grad() @@ -172,17 +570,104 @@ def generate_causal_block_mask(batch_size, nheads, seqlen, local_num, window_siz # Attention kernels # ---------------------------- def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False, attention_mask=None, return_KV=False): + global _DENSE_FAILED, _ATTN_V2_FAILED, _ATTN_V1_STRIDED_FAILED, _ATTN_V1_FAILED if attention_mask is not None: seqlen = q.shape[1] seqlen_kv = k.shape[1] q = rearrange(q, "b s (n d) -> (b s) n d", n=num_heads) k = rearrange(k, "b s (n d) -> (b s) n d", n=num_heads) v = rearrange(v, "b s (n d) -> (b s) n d", n=num_heads) - cu_seqlens_q = torch.tensor([0, seqlen], device=q.device, dtype=torch.int32) - cu_seqlens_k = torch.tensor([0, seqlen_kv], device=q.device, dtype=torch.int32) - head_mask_type = torch.tensor([1]*num_heads, device=q.device, dtype=torch.int32) - streaming_info = None base_blockmask = attention_mask + + # Adaptive backend: route dense when the sparse mask is too dense to pay off. + use_dense = False + if _ATTN_BACKEND == "dense": + use_dense = True + elif _ATTN_BACKEND == "auto" and _is_hopper_dev(q.device): + try: + density = base_blockmask.float().mean().item() + except Exception: + density = 1.0 + use_dense = density >= _ATTN_DENSE_THRESH + + if use_dense and not _DENSE_FAILED: + try: + # (S, n, d) -> (1, n, S, d) for fused dense SDPA, then back. + qd = q.unsqueeze(0).transpose(1, 2) + kd = k.unsqueeze(0).transpose(1, 2) + vd = v.unsqueeze(0).transpose(1, 2) + xd = _dense_sdpa((qd, kd, vd)) # (1, n, S, d) + x = xd.transpose(1, 2) # (1, S, n, d) + _perf_record("attn_dense") + return rearrange(x, "b s n d -> b s (n d)", n=num_heads) + except Exception as exc: + _DENSE_FAILED = True + _perf_error("attn_dense", exc) + torch.cuda.empty_cache() # fall back to sparse + + # Triton WGMMA block-sparse backend (opt-in, Hopper-guarded, exact mask). + if _ATTN_BACKEND in ("triton", "triton2") and _is_hopper_dev(q.device) and _TRITON_BSA is not None: + try: + # block mask -> (H, Nqb, Nkvb) bool + bm = base_blockmask + if bm.dim() == 4: + bm = bm[0] + bm = bm.bool() + # Phase 3: warp-specialized v2 kernel. Falls back to the v1 + # triton paths below on any error (then to block_sparse). + if (_ATTN_BACKEND == "triton2" and _TRITON_BSA_V2 is not None + and not _ATTN_V2_FAILED): + try: + xh = _TRITON_BSA_V2(q, k, v, bm) # (S, n, d) + _perf_record("attn_v2") + return xh.reshape(1, xh.shape[0], -1) + except Exception as exc: + _ATTN_V2_FAILED = True + _perf_error("attn_v2", exc) + torch.cuda.empty_cache() # fall back to v1 triton + # 2A-3: strided IO — q/k/v stay (S, n, d), output comes back + # (S, n, d); zero transpose/contiguous copies in the glue. + if (_ATTN_STRIDED_IO and _TRITON_BSA_SND is not None + and not _ATTN_V1_STRIDED_FAILED): + try: + xh = _TRITON_BSA_SND(q, k, v, bm) # (S, n, d) + _perf_record("attn_v1_strided") + return xh.reshape(1, xh.shape[0], -1) + except Exception as exc: + _ATTN_V1_STRIDED_FAILED = True + _perf_error("attn_v1_strided", exc) + torch.cuda.empty_cache() # fall back to contiguous triton + # contiguous path: q/k/v (S,n,d) -> (n,S,d) + if not _ATTN_V1_FAILED: + qh = q.transpose(0, 1).contiguous() + kh = k.transpose(0, 1).contiguous() + vh = v.transpose(0, 1).contiguous() + xh = _TRITON_BSA(qh, kh, vh, bm) # (n, S, d) + x = xh.transpose(0, 1).contiguous().unsqueeze(0) # (1, S, n, d) + _perf_record("attn_v1_contiguous") + return rearrange(x, "b s n d -> b s (n d)", n=num_heads) + except Exception as exc: + if not _ATTN_V1_FAILED: + _perf_error("attn_v1_contiguous", exc) + _ATTN_V1_FAILED = True + torch.cuda.empty_cache() # fall back to sparse + + if _MASKGEN_LEAN: + # 2A-4(c): these int32 tensors depend only on shapes; building + # them per call is a hidden H2D sync on the sparse path. + skey = (seqlen, seqlen_kv, num_heads, q.device.index) + ent = _SPARSE_SEQLENS_CACHE.get(skey) + if ent is None: + ent = (torch.tensor([0, seqlen], device=q.device, dtype=torch.int32), + torch.tensor([0, seqlen_kv], device=q.device, dtype=torch.int32), + torch.tensor([1]*num_heads, device=q.device, dtype=torch.int32)) + _SPARSE_SEQLENS_CACHE[skey] = ent + cu_seqlens_q, cu_seqlens_k, head_mask_type = ent + else: + cu_seqlens_q = torch.tensor([0, seqlen], device=q.device, dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, seqlen_kv], device=q.device, dtype=torch.int32) + head_mask_type = torch.tensor([1]*num_heads, device=q.device, dtype=torch.int32) + streaming_info = None max_seqlen_q_ = seqlen max_seqlen_k_ = seqlen_kv p_dropout = 0.0 @@ -200,6 +685,7 @@ def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads exact_streaming=False, return_attn_probs=False, ).unsqueeze(0) + _perf_record("attn_sparse") x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) elif compatibility_mode: q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) @@ -236,10 +722,27 @@ def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads return x -def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): +def _modulate_impl(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): return (x * (1 + scale) + shift) +modulate = _maybe_compile(_modulate_impl) + + +def layer_norm_modulate(x, norm, shift, scale, eps): + global _DIT_ROW_FUSION_FAILED + if (_DIT_ROW_FUSION and _DIT_LN_MODULATE is not None + and not _DIT_ROW_FUSION_FAILED): + try: + out = _DIT_LN_MODULATE(x, shift, scale, eps) + _perf_record("dit_row_ln_modulate") + return out + except Exception as exc: + _DIT_ROW_FUSION_FAILED = True + _perf_error("dit_row_ln_modulate", exc) + return modulate(norm(x), shift, scale) + + def sinusoidal_embedding_1d(dim, position): sinusoid = torch.outer(position.type(torch.float64), torch.pow( 10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2))) @@ -262,17 +765,106 @@ def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0): return freqs_cis +# --------------------------------------------------------------------------- +# Phase 2A-1b: fused RoPE apply (FLASHVSR_FUSE_ROPE, default OFF). +# +# The eager rope_apply materializes an fp64 copy of x (~104 MB @768), a +# complex128 product (~104 MB) and a bf16 down-cast per call — 60 calls per +# steady chunk ≈ 10 ms of pure memory traffic (ANALYSIS §1.2 "rope apply"). +# The fused path computes the *same* fp64 complex multiply in real arithmetic +# ((a+bi)(c+di) = (ac−bd)+(ad+bc)i — exactly what the eager complex kernel +# does) inside one torch.compile-generated kernel: reads bf16 x + fp64 freqs, +# writes bf16 out, no fp64 intermediates hit DRAM. freqs.real / freqs.imag are +# strided fp64 views of the complex tensor (no copy). Numerics: identical +# operations in fp64; any FMA-contraction difference is far below the bf16 +# output quantum — gated at PSNR ≥ 49 dB vs OFF (measured max|diff| reported +# in PHASE_BENCH_LOG.md). +# Compiled lazily on first use so the flag can be toggled at runtime; any +# compile/runtime failure falls back to the eager path silently. +# --------------------------------------------------------------------------- +_FUSE_ROPE = os.environ.get("FLASHVSR_FUSE_ROPE", "0") != "0" + +# --------------------------------------------------------------------------- +# Phase 5-P1: hand-written coalesced RoPE kernel (FLASHVSR_ROPE_KERNEL=triton, +# default OFF). Same fp64 elementwise math as the fused path (bit-exact, +# max|diff|==0 gated) but ~5x the effective bandwidth: the fused kernel reads +# freqs through stride-2 fp64 views of the complex tensor and stores through +# an interleaved stack (measured ~484 GB/s, 143 us/call); the hand kernel +# moves bf16 pairs as packed u32 words and reads the complex storage directly. +# Fallback ladder: triton kernel -> fused (torch.compile) -> eager. +# --------------------------------------------------------------------------- +_ROPE_KERNEL = os.environ.get("FLASHVSR_ROPE_KERNEL", "").lower() +_ROPE_TRITON_FAILED = False +_ROPE_FUSED_FAILED = False +try: + from .triton_rope import rope_apply_triton as _ROPE_TRITON +except Exception: + _ROPE_TRITON = None + +_rope_fused_fn = None + + +def _rope_apply_fused_impl(x, f_real, f_imag, num_heads): + B, S, D = x.shape + xv = x.reshape(B, S, num_heads, -1, 2) + xr = xv[..., 0].to(torch.float64) + xi = xv[..., 1].to(torch.float64) + # f_real/f_imag: (S, 1, dc) fp64 views -> broadcast over (B, S, n, dc) + o_r = xr * f_real - xi * f_imag + o_i = xr * f_imag + xi * f_real + out = torch.stack((o_r, o_i), dim=-1) # (B, S, n, dc, 2) + return out.flatten(2).to(x.dtype) # (B, S, n*d), interleaved pairs + + +def _get_rope_fused(): + global _rope_fused_fn + if _rope_fused_fn is None: + try: + _rope_fused_fn = torch.compile(_rope_apply_fused_impl, dynamic=True) + except Exception: + _rope_fused_fn = _rope_apply_fused_impl + return _rope_fused_fn + + def rope_apply(x, freqs, num_heads): + global _ROPE_TRITON_FAILED, _ROPE_FUSED_FAILED + if (_ROPE_KERNEL == "triton" and _ROPE_TRITON is not None + and not _ROPE_TRITON_FAILED): + try: + out = _ROPE_TRITON(x, freqs, num_heads) + _perf_record("rope_triton") + return out + except Exception as exc: + _ROPE_TRITON_FAILED = True + _perf_error("rope_triton", exc) + if _FUSE_ROPE and not _ROPE_FUSED_FAILED: + try: + out = _get_rope_fused()(x, freqs.real, freqs.imag, num_heads) + _perf_record("rope_fused") + return out + except Exception as exc: + _ROPE_FUSED_FAILED = True + _perf_error("rope_fused", exc) x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) x_out = torch.view_as_complex(x.to(torch.float64).reshape( x.shape[0], x.shape[1], x.shape[2], -1, 2)) x_out = torch.view_as_real(x_out * freqs).flatten(2) + _perf_record("rope_eager") return x_out.to(x.dtype) # ---------------------------- # Norms & Blocks # ---------------------------- +def _rmsnorm_impl(x, weight, eps): + dtype = x.dtype + out = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return out.to(dtype) * weight + + +_rmsnorm_fused = _maybe_compile(_rmsnorm_impl) + + class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-5): super().__init__() @@ -283,6 +875,8 @@ def norm(self, x): return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) def forward(self, x): + if _FUSE_NORM: + return _rmsnorm_fused(x, self.weight, self.eps) dtype = x.dtype return self.norm(x.float()).to(dtype) * self.weight @@ -318,68 +912,219 @@ def forward(self, x, freqs, f=None, h=None, w=None, local_num=None, topk=None, train_img=False, block_id=None, kv_len=None, is_full_block=False, is_stream=False, pre_cache_k=None, pre_cache_v=None, local_range = 9): B, L, D = x.shape + if is_stream and f == 6: + # A six-frame call begins a new video; never carry thresholds from + # its predecessor into the next stream. + self._mask_threshold_cache = None if is_stream and pre_cache_k is not None and pre_cache_v is not None: assert f==2, "f must be 2" if is_stream and (pre_cache_k is None or pre_cache_v is None): assert f==6, " start f must be 6" assert L == f * h * w, "Sequence length mismatch with provided (f,h,w)." - q = self.norm_q(self.q(x)) - k = self.norm_k(self.k(x)) - v = self.v(x) - q = rope_apply(q, freqs, self.num_heads) - k = rope_apply(k, freqs, self.num_heads) + with nvtx_range("qkv_norm"): + if _fp8g.enabled("qkv"): + # 2B-2: q/k/v consume the same activation -> quantize x once + # (per-row scales) and run three e4m3 scaled_mm GEMMs. + pre = _fp8g.quant(x.reshape(-1, x.shape[-1])) + q = self.norm_q(_fp8g.linear(self.q, x, pre=pre)) + k = self.norm_k(_fp8g.linear(self.k, x, pre=pre)) + v = _fp8g.linear(self.v, x, pre=pre) + else: + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + with nvtx_range("rope"): + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) win = (2, 8, 8) - q = q.view(B, f, h, w, D) - k = k.view(B, f, h, w, D) - v = v.view(B, f, h, w, D) + seqlen = f//win[0] + use_arena = _KV_RINGBUF and is_stream and B == 1 + use_zc = (_ATTN_ZEROCOPY and is_stream and B == 1 + and _ATTN_BACKEND == "triton2" and _ZC_ATTN is not None) + varena = None + if use_zc: + # 5A-v1: V skips window partitioning entirely — raster ring arena + # + rank-4 TMA gather in the kernel. K/Q keep the window layout + # (their pooled means feed the exact-mask top-k). + with nvtx_range("win_part"): + q = q.view(B, f, h, w, D) + k = k.view(B, f, h, w, D) + v = v.view(B, f, h, w, D) + q_w = WindowPartition3D.partition(q, win) + one_len = (h // win[1]) * (w // win[2]) + block_tokens = win[0] * win[1] * win[2] + varena = pre_cache_v if isinstance(pre_cache_v, _VRasterArena) \ + else _VRasterArena(kv_len, h, w, D, x.dtype, x.device) + varena.append(v) + if use_arena: + arena_k = pre_cache_k if isinstance(pre_cache_k, _KVArena) else \ + _KVArena(kv_len, one_len, block_tokens, D, x.dtype, x.device) + k_w = arena_k.append_partition(k, win) + else: + k_w = WindowPartition3D.partition(k, win) + if pre_cache_k is not None: + with nvtx_range("kv_cat"): + k_w = torch.cat([pre_cache_k, k_w], dim=0) + v_w = None + elif use_arena: + # 2A-2 arena path: partition K/V straight into the preallocated + # cache; k_w/v_w are the live contiguous views (pre + new, in + # temporal order) — bit-identical to the cat path below. + with nvtx_range("win_part"): + q = q.view(B, f, h, w, D) + k = k.view(B, f, h, w, D) + v = v.view(B, f, h, w, D) + q_w = WindowPartition3D.partition(q, win) + one_len = (h // win[1]) * (w // win[2]) + block_tokens = win[0] * win[1] * win[2] # tokens per window (=128) + arena_k = pre_cache_k if isinstance(pre_cache_k, _KVArena) else \ + _KVArena(kv_len, one_len, block_tokens, D, x.dtype, x.device) + arena_v = pre_cache_v if isinstance(pre_cache_v, _KVArena) else \ + _KVArena(kv_len, one_len, block_tokens, D, x.dtype, x.device) + k_w = arena_k.append_partition(k, win) + v_w = arena_v.append_partition(v, win) + else: + with nvtx_range("win_part"): + q = q.view(B, f, h, w, D) + k = k.view(B, f, h, w, D) + v = v.view(B, f, h, w, D) - q_w = WindowPartition3D.partition(q, win) - k_w = WindowPartition3D.partition(k, win) - v_w = WindowPartition3D.partition(v, win) + q_w = WindowPartition3D.partition(q, win) + k_w = WindowPartition3D.partition(k, win) + v_w = WindowPartition3D.partition(v, win) - seqlen = f//win[0] - one_len = k_w.shape[0] // B // seqlen - if pre_cache_k is not None and pre_cache_v is not None: - k_w = torch.cat([pre_cache_k, k_w], dim=0) - v_w = torch.cat([pre_cache_v, v_w], dim=0) + one_len = k_w.shape[0] // B // seqlen + if pre_cache_k is not None and pre_cache_v is not None: + with nvtx_range("kv_cat"): + k_w = torch.cat([pre_cache_k, k_w], dim=0) + v_w = torch.cat([pre_cache_v, v_w], dim=0) block_n = q_w.shape[0] // B block_s = q_w.shape[1] block_n_kv = k_w.shape[0] // B - reorder_q = rearrange(q_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n, block_s=block_s) - reorder_k = rearrange(k_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n_kv, block_s=block_s) - reorder_v = rearrange(v_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n_kv, block_s=block_s) + with nvtx_range("reorder"): + reorder_q = rearrange(q_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n, block_s=block_s) + reorder_k = rearrange(k_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n_kv, block_s=block_s) + reorder_v = None if use_zc else rearrange(v_w, '(b block_n) (block_s) d -> b (block_n block_s) d', block_n=block_n_kv, block_s=block_s) window_size = win[0]*h*w//128 - if self.local_attn_mask is None or self.local_attn_mask_h!=h//8 or self.local_attn_mask_w!=w//8 or self.local_range!=local_range: - self.local_attn_mask = build_local_block_mask_shifted_vec_normal_slide(h//8, w//8, local_range, local_range, include_self=True, device=k_w.device) - self.local_attn_mask_h = h//8 - self.local_attn_mask_w = w//8 - self.local_range = local_range - attention_mask = generate_draft_block_mask(B, self.num_heads, seqlen, q_w, k_w, topk=topk, local_attn_mask=self.local_attn_mask) - - x = self.attn(reorder_q, reorder_k, reorder_v, attention_mask) - - cur_block_n, cur_block_s, _ = k_w.shape - cache_num = cur_block_n // one_len - if cache_num > kv_len: - cache_k = k_w[one_len:, :, :] - cache_v = v_w[one_len:, :, :] - else: - cache_k = k_w - cache_v = v_w - - x = rearrange(x, 'b (block_n block_s) d -> (b block_n) (block_s) d', block_n=block_n, block_s=block_s) - x = WindowPartition3D.reverse(x, win, (f, h, w)) - x = x.view(B, f*h*w, D) - + with nvtx_range("mask_gen"): + if self.local_attn_mask is None or self.local_attn_mask_h!=h//8 or self.local_attn_mask_w!=w//8 or self.local_range!=local_range: + self.local_attn_mask = build_local_block_mask_shifted_vec_normal_slide(h//8, w//8, local_range, local_range, include_self=True, device=k_w.device) + self.local_attn_mask_h = h//8 + self.local_attn_mask_w = w//8 + self.local_range = local_range + pooled_k = None + if _POOLED_K_CACHE and B == 1 and is_stream: + # 5-P2: pool only the windows appended this call; carry old + # rows (bit-identical values) in a rolling buffer that mirrors + # the kv window's cat/trim sequence exactly. + n_new = one_len * seqlen + old_rows = k_w.shape[0] - n_new + prev = getattr(self, "_pk_cache", None) + if (prev is not None and old_rows > 0 + and prev.shape[0] == old_rows + and prev.dtype == k_w.dtype + and prev.shape[1] == k_w.shape[2]): + pooled_k = torch.cat( + [prev, torch.mean(k_w[old_rows:], dim=1)], dim=0) + _perf_record("pooled_k_incremental") + else: # fresh video / reset / mismatch -> re-pool everything + pooled_k = torch.mean(k_w, dim=1) + _perf_record("pooled_k_rebuild") + # mirror cache_trim below: drop the oldest temporal slot when + # more than kv_len slots are cached + if (k_w.shape[0] // one_len) > kv_len: + self._pk_cache = pooled_k[one_len:] + else: + self._pk_cache = pooled_k + if _MASKGEN_THRESHOLD_CACHE and B == 1 and is_stream: + attention_mask, thresholds = generate_draft_block_mask( + B, self.num_heads, seqlen, q_w, k_w, topk=topk, + local_attn_mask=self.local_attn_mask, avgpool_k=pooled_k, + cached_thresholds=getattr(self, "_mask_threshold_cache", None), + return_thresholds=True) + self._mask_threshold_cache = thresholds + else: + attention_mask = generate_draft_block_mask( + B, self.num_heads, seqlen, q_w, k_w, topk=topk, + local_attn_mask=self.local_attn_mask, avgpool_k=pooled_k) + + zc_done = False + with nvtx_range("attn_core"): + if use_zc: + if not getattr(self, "_zc_failed", False): + try: + bmb = attention_mask[0] if attention_mask.dim() == 4 else attention_mask + hd = D // self.num_heads + xh = _ZC_ATTN( + reorder_q.view(-1, self.num_heads, hd), + reorder_k.view(-1, self.num_heads, hd), + varena.buf, varena.start, h, w, bmb.bool()) + x = xh.reshape(B, f * h * w, D) # already RASTER order + zc_done = True + _perf_record("attn_zc_v2") + except Exception as exc: + self._zc_failed = True + _perf_error("attn_zc_v2", exc) + torch.cuda.empty_cache() + if not zc_done: + # materialize windowed V from the raster arena and take the + # standard ladder (triton2 -> v1 triton -> block_sparse) + v_w = WindowPartition3D.partition( + varena.live.unsqueeze(0), win) + reorder_v = rearrange( + v_w, '(b block_n) (block_s) d -> b (block_n block_s) d', + block_n=block_n_kv, block_s=block_s) + if not zc_done: + x = self.attn(reorder_q, reorder_k, reorder_v, attention_mask) + + with nvtx_range("cache_trim"): + if use_zc: + if use_arena: + if arena_k.length // one_len > kv_len: + arena_k.trim(one_len) + cache_k = arena_k + else: + cache_k = k_w[one_len:, :, :] \ + if (k_w.shape[0] // one_len) > kv_len else k_w + if varena.length // 2 > kv_len: + varena.trim(2) + cache_v = varena + elif use_arena: + # same semantics as the slice-trim below: drop the oldest + # temporal slot once more than kv_len slots are cached. + if arena_k.length // one_len > kv_len: + arena_k.trim(one_len) + arena_v.trim(one_len) + cache_k = arena_k + cache_v = arena_v + else: + cur_block_n, cur_block_s, _ = k_w.shape + cache_num = cur_block_n // one_len + if cache_num > kv_len: + cache_k = k_w[one_len:, :, :] + cache_v = v_w[one_len:, :, :] + else: + cache_k = k_w + cache_v = v_w + + if not zc_done: + with nvtx_range("win_rev"): + x = rearrange(x, 'b (block_n block_s) d -> (b block_n) (block_s) d', block_n=block_n, block_s=block_s) + x = WindowPartition3D.reverse(x, win, (f, h, w)) + x = x.view(B, f*h*w, D) + + with nvtx_range("o_proj"): + # 2B-2: e4m3 o-projection (per-row activation scales). + out = _fp8g.linear(self.o, x) if _fp8g.enabled("o") else self.o(x) if is_stream: - return self.o(x), cache_k, cache_v - return self.o(x) + return out, cache_k, cache_v + return out class CrossAttention(nn.Module): @@ -429,11 +1174,30 @@ def forward(self, x: torch.Tensor, y: torch.Tensor, is_stream: bool = False): return self.o(x) +def _gate_impl(x, gate, residual): + return x + gate * residual + + +_gate_fused = _maybe_compile(_gate_impl) + + class GateModule(nn.Module): def __init__(self,): super().__init__() def forward(self, x, gate, residual): + global _DIT_ROW_FUSION_FAILED + if (_DIT_ROW_FUSION and _DIT_GATED_RESIDUAL is not None + and not _DIT_ROW_FUSION_FAILED): + try: + out = _DIT_GATED_RESIDUAL(x, gate, residual) + _perf_record("dit_row_gate") + return out + except Exception as exc: + _DIT_ROW_FUSION_FAILED = True + _perf_error("dit_row_gate", exc) + if _FUSE_NORM: + return _gate_fused(x, gate, residual) return x + gate * residual @@ -443,6 +1207,7 @@ def __init__(self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): self.dim = dim self.num_heads = num_heads self.ffn_dim = ffn_dim + self.eps = eps self.self_attn = SelfAttention(dim, num_heads, eps) self.cross_attn = CrossAttention(dim, num_heads, eps) @@ -454,22 +1219,47 @@ def __init__(self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): approximate='tanh'), nn.Linear(ffn_dim, dim)) self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) self.gate = GateModule() + # B1 lossless cache: (modulation + t_mod).chunk(6) is step-invariant. + self._mod_cache = None + self._mod_cache_key = None def forward(self, x, context, t_mod, freqs, f, h, w, local_num=None, topk=None, train_img=False, block_id=None, kv_len=None, is_full_block=False, is_stream=False, pre_cache_k=None, pre_cache_v=None, local_range = 9): - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) - input_x = modulate(self.norm1(x), shift_msa, scale_msa) - self_attn_output, self_attn_cache_k, self_attn_cache_v = self.self_attn( - input_x, freqs, f, h, w, local_num, topk, train_img, block_id, - kv_len=kv_len, is_full_block=is_full_block, is_stream=is_stream, - pre_cache_k=pre_cache_k, pre_cache_v=pre_cache_v, local_range = local_range) - - x = self.gate(x, gate_msa, self_attn_output) - x = x + self.cross_attn(self.norm3(x), context, is_stream=is_stream) - input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) - x = self.gate(x, gate_mlp, self.ffn(input_x)) + with nvtx_range("mod1"): + if _CACHE_MOD: + key = (id(t_mod), t_mod.dtype, t_mod.device) + if self._mod_cache_key != key: + self._mod_cache = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod + ).chunk(6, dim=1) + self._mod_cache_key = key + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._mod_cache + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + input_x = layer_norm_modulate(x, self.norm1, shift_msa, scale_msa, self.eps) + with nvtx_range("self_attn"): + self_attn_output, self_attn_cache_k, self_attn_cache_v = self.self_attn( + input_x, freqs, f, h, w, local_num, topk, train_img, block_id, + kv_len=kv_len, is_full_block=is_full_block, is_stream=is_stream, + pre_cache_k=pre_cache_k, pre_cache_v=pre_cache_v, local_range = local_range) + + with nvtx_range("gate1"): + x = self.gate(x, gate_msa, self_attn_output) + with nvtx_range("xattn"): + x = x + self.cross_attn(self.norm3(x), context, is_stream=is_stream) + with nvtx_range("ffn"): + input_x = layer_norm_modulate(x, self.norm2, shift_mlp, scale_mlp, self.eps) + if _fp8g.enabled("ffn"): + # 2B-2: e4m3 ffn1+ffn2; GELU is fused into ffn2's input + # quantization kernel (fp32 gelu -> e4m3). + x = self.gate(x, gate_mlp, _fp8g.ffn(self.ffn, input_x)) + elif _fp8g.enabled("ffn1"): + # 2B-2 bisection scope: e4m3 ffn1 only (ffn2 stays bf16). + x = self.gate(x, gate_mlp, _fp8g.ffn_partial(self.ffn, input_x)) + else: + x = self.gate(x, gate_mlp, self.ffn(input_x)) if is_stream: return x, self_attn_cache_k, self_attn_cache_v return x @@ -503,9 +1293,21 @@ def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + # B1 lossless cache: (modulation + t_mod).chunk(2) is step-invariant. + self._mod_cache = None + self._mod_cache_key = None def forward(self, x, t_mod): - shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) + if _CACHE_MOD: + key = (id(t_mod), t_mod.dtype, t_mod.device) + if self._mod_cache_key != key: + self._mod_cache = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod + ).chunk(2, dim=1) + self._mod_cache_key = key + shift, scale = self._mod_cache + else: + shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) x = (self.head(self.norm(x) * (1 + scale) + shift)) return x diff --git a/diffsynth/nvtx_utils.py b/diffsynth/nvtx_utils.py new file mode 100644 index 00000000..3a9e494f --- /dev/null +++ b/diffsynth/nvtx_utils.py @@ -0,0 +1,35 @@ +"""Knob-gated NVTX profiling helpers (FLASHVSR_NVTX=1, default OFF). + +When FLASHVSR_NVTX is unset/0, ``nvtx_range`` returns a shared null context +manager: zero GPU effect and negligible CPU cost, so the default path is +behaviourally identical to the un-instrumented code. + +When FLASHVSR_NVTX=1, ``nvtx_range(name)`` emits an NVTX range visible in +Nsight Systems / Nsight Compute (used for phase attribution and kernel +filtering, e.g. ``ncu --nvtx --nvtx-include``). +""" +import os + +import torch + +NVTX_ENABLED = os.environ.get("FLASHVSR_NVTX", "0") != "0" + + +class _NullCtx: + __slots__ = () + + def __enter__(self): + return None + + def __exit__(self, *exc): + return False + + +_NULL = _NullCtx() + +if NVTX_ENABLED: + def nvtx_range(name: str): + return torch.cuda.nvtx.range(name) +else: + def nvtx_range(name: str): # noqa: ARG001 - keep signature identical + return _NULL diff --git a/diffsynth/perf_stats.py b/diffsynth/perf_stats.py new file mode 100644 index 00000000..9dc78107 --- /dev/null +++ b/diffsynth/perf_stats.py @@ -0,0 +1,45 @@ +"""Low-overhead runtime route and fallback telemetry for FlashVSR. + +Enable with ``FLASHVSR_TELEMETRY=1``. The profiling target enables it +automatically when ``FLASHVSR_REQUIRE_FASTPATHS=1`` is requested. +""" + +import os +from collections import Counter + + +_ENABLED = ( + os.environ.get("FLASHVSR_TELEMETRY", "0") != "0" + or os.environ.get("FLASHVSR_REQUIRE_FASTPATHS", "0") != "0" +) +_COUNTS = Counter() +_ERRORS = {} + + +def enabled(): + return _ENABLED + + +def record(name, count=1): + if _ENABLED: + _COUNTS[name] += count + + +def record_error(name, error): + if not _ENABLED: + return + _COUNTS[f"{name}_error"] += 1 + _ERRORS.setdefault(name, f"{type(error).__name__}: {error}") + + +def reset(preserve_errors=False): + _COUNTS.clear() + if not preserve_errors: + _ERRORS.clear() + + +def snapshot(): + return { + "counts": dict(sorted(_COUNTS.items())), + "errors": dict(sorted(_ERRORS.items())), + } diff --git a/diffsynth/pipelines/flashvsr_tiny.py b/diffsynth/pipelines/flashvsr_tiny.py index 47a4eb16..9d83547d 100644 --- a/diffsynth/pipelines/flashvsr_tiny.py +++ b/diffsynth/pipelines/flashvsr_tiny.py @@ -16,9 +16,128 @@ from ..models.wan_video_dit import WanModel, RMSNorm, sinusoidal_embedding_1d from ..models.wan_video_vae import WanVideoVAE, RMS_norm, CausalConv3d, Upsample from ..schedulers.flow_match import FlowMatchScheduler +from ..nvtx_utils import nvtx_range +from ..perf_stats import record as _perf_record, record_error as _perf_error from .base import BasePipeline +# --------------------------------------------------------------------------- +# Phase 2A-1a: lossless RoPE freqs cache (FLASHVSR_CACHE_ROPE_FREQS, default OFF). +# +# The eager path assembles the per-chunk RoPE freqs tensor on the CPU from +# `dit.freqs` (three complex128 tables that live on the CPU) and then moves the +# ~(f*h*w, 1, 64)-complex result to the GPU — every chunk. Profiling (ANALYSIS +# §4 item 5) shows this as tens of ms of per-chunk CPU wall plus an ~8.6 MB H2D +# copy @768x1408. The assembly is pure slice/expand/cat (no arithmetic), so +# performing exactly the same copies on-device from device-resident base tables +# is bit-identical. +# +# Cache layout (bounded memory; shape/dtype/device aware): +# dit._rope_base_dev : {device_str: (f_tab, h_tab, w_tab) on device} +# one-time H2D copy of the small per-axis freq tables. +# dit._rope_freqs_buf: {(f, h, w, device_str): entry} +# entry = {"buf": (f*h*w, 1, D) complex buffer on device, +# "hw_done": bool, # h/w columns written (invariant per key) +# "f_start": int} # temporal offset currently in the f columns +# +# Cache key semantics: the assembled tensor depends only on (f, h, w, f_start, +# device). The h/w columns are invariant for a given (f, h, w); only the f +# columns depend on the chunk's temporal offset f_start (= 0 for chunk 0, else +# 4 + 2*idx), so they are rewritten in place when f_start changes. The buffer +# is consumed strictly inside the current chunk (rope_apply) and never retained +# by any cache (pre_cache_k/v hold post-RoPE K/V), so in-place reuse is safe. +# --------------------------------------------------------------------------- +_CACHE_ROPE_FREQS = os.environ.get("FLASHVSR_CACHE_ROPE_FREQS", "0") != "0" + + +def _rope_freqs_cached(dit, f, h, w, f_start, device): + """Device-side cached assembly of the per-chunk RoPE freqs tensor. + + Bit-identical to the eager CPU path: identical source values, identical + layout; only copy operations (slice/expand/copy_), no arithmetic. + """ + dev_key = str(device) + base_map = getattr(dit, "_rope_base_dev", None) + if base_map is None: + base_map = {} + dit._rope_base_dev = base_map + base = base_map.get(dev_key) + if base is None: + base = tuple(t.to(device) for t in dit.freqs) + base_map[dev_key] = base + f_tab, h_tab, w_tab = base + fd, hd, wd = f_tab.shape[1], h_tab.shape[1], w_tab.shape[1] + + buf_map = getattr(dit, "_rope_freqs_buf", None) + if buf_map is None: + buf_map = {} + dit._rope_freqs_buf = buf_map + key = (f, h, w, dev_key) + ent = buf_map.get(key) + if ent is None: + buf = torch.empty(f * h * w, 1, fd + hd + wd, dtype=f_tab.dtype, device=device) + ent = {"buf": buf, "hw_done": False, "f_start": None} + buf_map[key] = ent + buf = ent["buf"] + v = buf.view(f, h, w, fd + hd + wd) + if not ent["hw_done"]: + v[..., fd:fd + hd].copy_(h_tab[:h].view(1, h, 1, hd).expand(f, h, w, hd)) + v[..., fd + hd:].copy_(w_tab[:w].view(1, 1, w, wd).expand(f, h, w, wd)) + ent["hw_done"] = True + if ent["f_start"] != f_start: + v[..., :fd].copy_(f_tab[f_start:f_start + f].view(f, 1, 1, fd).expand(f, h, w, fd)) + ent["f_start"] = f_start + return buf + + +# --------------------------------------------------------------------------- +# Phase 2B-1: decoder overlap on a side CUDA stream +# (FLASHVSR_DECODER_OVERLAP, default OFF). +# +# The serialized path decodes ONCE after the whole denoise loop, so the +# TCDecoder is a fully serialized tail (17-23% of E2E, ANALYSIS §1.1/§3 H6). +# With the flag ON, each chunk's latents are decoded on a dedicated side +# stream as soon as they are finalized, so decode N runs concurrently with +# denoise chunks N+1.. on the main stream. This is a pure scheduling change: +# the TCDecoder is streaming-capable by construction (TAEHV mem-blocks carry +# per-timestep state across decode_video calls; `flashvsr_tiny_long.py` +# decodes per chunk with the same LQ_pre_idx:LQ_cur_idx cond slices), and the +# per-chunk split feeds the decoder the exact same per-timestep inputs in the +# exact same order as the one-shot decode -> bit-identical output. +# +# Stream / event / lifetime contract: +# * main stream : denoise chunks, final `torch.cat` assembly, color fix. +# * decode stream: every TCDecoder op (incl. pixel_shuffle(cond), the +# channels_last weight conversion on first use, and the +# stateful TAEHV.mem updates). Decode calls are enqueued in +# chunk order on ONE stream, so mem-block state transitions +# are identical to the serialized path. +# * ready event (per chunk, recorded on main stream): protects the +# main->decode handoff. Recorded only after `cur_latents = cur_latents - +# noise_pred` is enqueued, i.e. after the decoder input is final (later +# iterations only rebind `cur_latents`, they never mutate it in place). +# * done event (per chunk, recorded on decode stream): protects the +# decode->main handoff. The main stream waits on all done events right +# before output assembly (GPU-side wait_event, no CPU sync in the loop). +# * Lifetime: decode-stream reads `cur_latents` (kept alive in +# `latents_total` for the whole call), `LQ_video` (caller-owned, alive and +# read-only for the whole call) and the decoder weights/mem (only ever +# touched by the decode stream). `record_stream` is additionally called on +# the cross-stream tensors as defense in depth so the caching allocator +# inserts event guards even if a future edit drops the references early. +# * Ordering: decoded chunks land in a list indexed by chunk id and are +# concatenated in that order -> output ordering is identical to the +# serialized path by construction (never completion order). +# --------------------------------------------------------------------------- +_DECODER_OVERLAP = os.environ.get("FLASHVSR_DECODER_OVERLAP", "0") != "0" + +# Write each overlapped decoder chunk directly into its final temporal slice. +# This removes the serialized post-wait torch.cat and skips stacking the three +# causal warm-up frames that are trimmed from chunk 0. Copy-only, opt-in. +_DECODER_DIRECT_OUTPUT = os.environ.get( + "FLASHVSR_TCDECODER_DIRECT_OUTPUT", "0") != "0" + + # ----------------------------- # 基础工具：ADAIN 所需的统计量（保留以备需要；管线默认用 wavelet） # ----------------------------- @@ -337,6 +456,8 @@ def __call__( latents = noise process_total_num = (num_frames - 1) // 8 - 2 + if process_total_num < 1: + raise ValueError("FlashVSR Tiny requires at least 25 input frames") is_stream = True # 清理可能存在的 LQ_proj_in cache @@ -348,8 +469,60 @@ def __call__( LQ_pre_idx = 0 LQ_cur_idx = 0 + # Phase 2B-1 (FLASHVSR_DECODER_OVERLAP): per-chunk decode on a side + # stream. See the module-level comment above `_DECODER_OVERLAP` for + # the full stream/event/lifetime contract. + use_decoder_overlap = ( + _DECODER_OVERLAP and LQ_video is not None and LQ_video.is_cuda + ) + if _DECODER_OVERLAP and not use_decoder_overlap: + _perf_record("decoder_overlap_unavailable") + if use_decoder_overlap: + if getattr(self, "_decode_stream", None) is None: + # One side stream per pipeline instance, reused across calls + # (warmup + measured) so allocator pools stay stable. + self._decode_stream = torch.cuda.Stream(device=LQ_video.device) + main_stream = torch.cuda.current_stream(LQ_video.device) + # Slots indexed by chunk id -> assembly is always in logical chunk + # order, regardless of decode completion order. + direct_decoder_output = _DECODER_DIRECT_OUTPUT + if direct_decoder_output: + # Chunk 0 yields 21 RGB frames after causal trim; every later + # two-latent chunk yields eight. For accepted 8n+5 inputs this + # deliberately matches the existing cat path, which ignores + # the final four input frames rather than returning unwritten + # output storage. + decoded_frame_count = 21 + 8 * (process_total_num - 1) + try: + frames = torch.empty( + (1, 3, decoded_frame_count, height, width), + dtype=self.torch_dtype, device=LQ_video.device) + frames.record_stream(self._decode_stream) + except Exception as exc: + _perf_error("decoder_direct_output", exc) + if os.environ.get("FLASHVSR_REQUIRE_FASTPATHS", "0") != "0": + raise + torch.cuda.empty_cache() + direct_decoder_output = False + if not direct_decoder_output: + decoded_chunks = [None] * process_total_num + decode_done_events = [None] * process_total_num + + # Profiling window (FLASHVSR_NVTX tooling): cudaProfilerStart at chunk + # PROF_START, cudaProfilerStop after chunk PROF_STOP-1 (i.e. window is + # [start, stop)). Read at call time so a warmup call can run with the + # window disabled and the measured call can enable it via os.environ. + # Default -1/-1 -> disabled, zero behaviour change. + _prof_start = int(os.environ.get("FLASHVSR_PROFILER_START_CHUNK", "-1")) + _prof_stop = int(os.environ.get("FLASHVSR_PROFILER_STOP_CHUNK", "-1")) + with torch.no_grad(): - for cur_process_idx in tqdm(range(process_total_num)): + for cur_process_idx in progress_bar_cmd(range(process_total_num)): + if _prof_start >= 0 and cur_process_idx == _prof_start: + torch.cuda.synchronize() + torch.cuda.profiler.start() + nvtx_chunk = nvtx_range(f"chunk{cur_process_idx}") + nvtx_chunk.__enter__() if cur_process_idx == 0: pre_cache_k = [None] * len(self.dit.blocks) pre_cache_v = [None] * len(self.dit.blocks) @@ -386,48 +559,133 @@ def __call__( cur_latents = latents[:, :, 4+cur_process_idx*2:6+cur_process_idx*2, :, :] # 推理（无 motion_controller / vace） - noise_pred_posi, pre_cache_k, pre_cache_v = model_fn_wan_video( - self.dit, - x=cur_latents, - timestep=self.timestep, - context=None, - tea_cache=None, - use_unified_sequence_parallel=False, - LQ_latents=LQ_latents, - is_full_block=is_full_block, - is_stream=is_stream, - pre_cache_k=pre_cache_k, - pre_cache_v=pre_cache_v, - topk_ratio=topk_ratio, - kv_ratio=kv_ratio, - cur_process_idx=cur_process_idx, - t_mod=self.t_mod, - t=self.t, - local_range = local_range, - ) + with nvtx_range("dit_forward"): + noise_pred_posi, pre_cache_k, pre_cache_v = model_fn_wan_video( + self.dit, + x=cur_latents, + timestep=self.timestep, + context=None, + tea_cache=None, + use_unified_sequence_parallel=False, + LQ_latents=LQ_latents, + is_full_block=is_full_block, + is_stream=is_stream, + pre_cache_k=pre_cache_k, + pre_cache_v=pre_cache_v, + topk_ratio=topk_ratio, + kv_ratio=kv_ratio, + cur_process_idx=cur_process_idx, + t_mod=self.t_mod, + t=self.t, + local_range = local_range, + ) # 更新 latent cur_latents = cur_latents - noise_pred_posi + # NOTE: in overlap mode `latents_total` doubles as the + # lifetime guard that keeps each chunk's decoder input alive + # until the final decode sync (do not drop this append). latents_total.append(cur_latents) - LQ_pre_idx = LQ_cur_idx - latents = torch.cat(latents_total, dim=2) + if use_decoder_overlap: + # ---- Phase 2B-1: hand chunk `cur_process_idx` to the ---- + # ---- decode stream and keep denoising on main. ---- + # `cur_latents` is final here (nothing after this point + # writes to it; next iteration rebinds the name). The + # ready event fences all main-stream work that produced it. + ready_event = torch.cuda.Event() + ready_event.record(main_stream) + with torch.cuda.stream(self._decode_stream): + self._decode_stream.wait_event(ready_event) + with nvtx_range(f"decode{cur_process_idx}"): + # Same call/cond slicing as the serialized decode, + # split per chunk (semantics identical to + # flashvsr_tiny_long.py); mul_/sub_ run on the + # decode stream and only touch decode-owned memory. + if direct_decoder_output: + frame_start = 0 if cur_process_idx == 0 \ + else 21 + (cur_process_idx - 1) * 8 + frame_count = 21 if cur_process_idx == 0 else 8 + output = frames[:, :, frame_start: + frame_start + frame_count] + output_ntchw = output.transpose(1, 2) + else: + output_ntchw = None + try: + dec = self.TCDecoder.decode_video( + cur_latents.transpose(1, 2), + parallel=False, + show_progress_bar=False, + cond=LQ_video[:, :, LQ_pre_idx:LQ_cur_idx, :, :], + out=output_ntchw, + ).transpose(1, 2).mul_(2).sub_(1) + except Exception as exc: + if direct_decoder_output: + _perf_error("decoder_direct_output", exc) + raise + _perf_record("decoder_overlap_chunks") + if direct_decoder_output: + _perf_record("decoder_direct_output_chunks") + done_event = torch.cuda.Event() + done_event.record(self._decode_stream) + # Defense in depth: tell the caching allocator these + # main-stream allocations are consumed by the decode + # stream, so any future free is event-guarded even if the + # Python references above were ever dropped early. + cur_latents.record_stream(self._decode_stream) + LQ_video.record_stream(self._decode_stream) + if not direct_decoder_output: + decoded_chunks[cur_process_idx] = dec + decode_done_events[cur_process_idx] = done_event - # Decode - frames = self.TCDecoder.decode_video(latents.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ_video[:,:,:LQ_cur_idx,:,:]).transpose(1, 2).mul_(2).sub_(1) + LQ_pre_idx = LQ_cur_idx + nvtx_chunk.__exit__(None, None, None) + if _prof_stop >= 0 and cur_process_idx == _prof_stop - 1: + torch.cuda.synchronize() + torch.cuda.profiler.stop() + + if use_decoder_overlap: + # ---- Phase 2B-1: final (and only) decode synchronization ---- + # GPU-side ordering only: the main stream waits on the decode + # done events; no torch.cuda.synchronize() and no CPU block. + # All decodes were already enqueued inside the chunk loop. + with nvtx_range("decode_wait"): + for done_event in decode_done_events: + main_stream.wait_event(done_event) + if direct_decoder_output: + _perf_record("decoder_direct_output_complete") + else: + for dec in decoded_chunks: + # Decode-stream allocations are read by the + # main-stream cat below; event-guard their free. + dec.record_stream(main_stream) + # Assemble strictly in chunk-id order. + frames = torch.cat(decoded_chunks, dim=2) + _perf_record("decoder_final_cat") + else: + latents = torch.cat(latents_total, dim=2) + + # Decode + with nvtx_range("decode"): + frames = self.TCDecoder.decode_video(latents.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ_video[:,:,:LQ_cur_idx,:,:]).transpose(1, 2).mul_(2).sub_(1) + _perf_record("decoder_serialized_calls") # 颜色校正（wavelet） try: if color_fix: - frames = self.ColorCorrector( - frames.to(device=LQ_video.device), - LQ_video[:, :, :frames.shape[2], :, :], - clip_range=(-1, 1), - chunk_size=16, - method='adain' - ) - except: - pass + with nvtx_range("color_fix"): + frames = self.ColorCorrector( + frames.to(device=LQ_video.device), + LQ_video[:, :, :frames.shape[2], :, :], + clip_range=(-1, 1), + chunk_size=16, + method='adain' + ) + _perf_record("color_fix_success") + except Exception as exc: + _perf_error("color_fix", exc) + if os.environ.get("FLASHVSR_REQUIRE_FASTPATHS", "0") != "0": + raise return frames[0] @@ -507,7 +765,8 @@ def model_fn_wan_video( **kwargs, ): # patchify - x, (f, h, w) = dit.patchify(x) + with nvtx_range("patchify"): + x, (f, h, w) = dit.patchify(x) win = (2, 8, 8) seqlen = f // win[0] @@ -518,18 +777,23 @@ def model_fn_wan_video( kv_len = int(kv_ratio) # RoPE 位置（分段） - if cur_process_idx == 0: - freqs = torch.cat([ - dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), - dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) - else: - freqs = torch.cat([ - dit.freqs[0][4 + cur_process_idx*2:4 + cur_process_idx*2 + f].view(f, 1, 1, -1).expand(f, h, w, -1), - dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + with nvtx_range("rope_freqs"): + if _CACHE_ROPE_FREQS: + # 2A-1a: on-device cached assembly (bit-identical, no CPU work/H2D). + f_start = 0 if cur_process_idx == 0 else 4 + cur_process_idx * 2 + freqs = _rope_freqs_cached(dit, f, h, w, f_start, x.device) + elif cur_process_idx == 0: + freqs = torch.cat([ + dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + else: + freqs = torch.cat([ + dit.freqs[0][4 + cur_process_idx*2:4 + cur_process_idx*2 + f].view(f, 1, 1, -1).expand(f, h, w, -1), + dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) # TeaCache（默认不启用） tea_cache_update = tea_cache.check(dit, x, t_mod) if tea_cache is not None else False @@ -548,27 +812,30 @@ def model_fn_wan_video( x = tea_cache.update(x) else: for block_id, block in enumerate(dit.blocks): - if LQ_latents is not None and block_id < len(LQ_latents): - x = x + LQ_latents[block_id] - x, last_pre_cache_k, last_pre_cache_v = block( - x, context, t_mod, freqs, f, h, w, - local_num, topk, - block_id=block_id, - kv_len=kv_len, - is_full_block=is_full_block, - is_stream=is_stream, - pre_cache_k=pre_cache_k[block_id] if pre_cache_k is not None else None, - pre_cache_v=pre_cache_v[block_id] if pre_cache_v is not None else None, - local_range = local_range, - ) - if pre_cache_k is not None: pre_cache_k[block_id] = last_pre_cache_k - if pre_cache_v is not None: pre_cache_v[block_id] = last_pre_cache_v + with nvtx_range(f"blk{block_id}"): + if LQ_latents is not None and block_id < len(LQ_latents): + x = x + LQ_latents[block_id] + x, last_pre_cache_k, last_pre_cache_v = block( + x, context, t_mod, freqs, f, h, w, + local_num, topk, + block_id=block_id, + kv_len=kv_len, + is_full_block=is_full_block, + is_stream=is_stream, + pre_cache_k=pre_cache_k[block_id] if pre_cache_k is not None else None, + pre_cache_v=pre_cache_v[block_id] if pre_cache_v is not None else None, + local_range = local_range, + ) + if pre_cache_k is not None: pre_cache_k[block_id] = last_pre_cache_k + if pre_cache_v is not None: pre_cache_v[block_id] = last_pre_cache_v - x = dit.head(x, t) + with nvtx_range("head"): + x = dit.head(x, t) if use_unified_sequence_parallel: import torch.distributed as dist from xfuser.core.distributed import get_sp_group if dist.is_initialized() and dist.get_world_size() > 1: x = get_sp_group().all_gather(x, dim=1) - x = dit.unpatchify(x, (f, h, w)) + with nvtx_range("unpatchify"): + x = dit.unpatchify(x, (f, h, w)) return x, pre_cache_k, pre_cache_v diff --git a/examples/WanVSR/probe_attention_shapes.py b/examples/WanVSR/probe_attention_shapes.py new file mode 100644 index 00000000..d155357b --- /dev/null +++ b/examples/WanVSR/probe_attention_shapes.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Capture real self-attention shapes + sparsity in the v1.1 Tiny DiT @768x1408. + +Hooks block_sparse_attn_func / generate_draft_block_mask to record q/k/v shapes, +block counts, and the fraction of KV blocks actually attended (sparsity). + +Run from examples/WanVSR/ : + python probe_attention_shapes.py +""" +import os, importlib.util +import numpy as np +from PIL import Image +import imageio +import torch + +os.environ["FLASHVSR_CONV3D_BACKEND"] = "gemm" +import utils.utils as wanutils; wanutils._CONV3D_BACKEND = "gemm" + +_here = os.path.dirname(os.path.abspath(__file__)) +_spec = importlib.util.spec_from_file_location("infer_v1_1_tiny", os.path.join(_here, "infer_flashvsr_v1.1_tiny.py")) +_infer = importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_infer) +init_pipeline = _infer.init_pipeline; largest_8n1_leq = _infer.largest_8n1_leq + +import diffsynth.models.wan_video_dit as dit + +REF_W, REF_H, SCALE = 768, 1408, 4 +SRC_W, SRC_H = REF_W // SCALE, REF_H // SCALE + +records = [] +_orig_bsa = dit.block_sparse_attn_func +def bsa_hook(q, k, v, cu_q, cu_k, head_mask_type, streaming_info, base_blockmask, *a, **kw): + m = base_blockmask + rec = { + "q": tuple(q.shape), "k": tuple(k.shape), "v": tuple(v.shape), + "mask": tuple(m.shape) if hasattr(m, "shape") else None, + "mask_density": float(m.float().mean().item()) if hasattr(m, "float") else None, + } + if len(records) < 6: + records.append(rec) + return _orig_bsa(q, k, v, cu_q, cu_k, head_mask_type, streaming_info, base_blockmask, *a, **kw) +dit.block_sparse_attn_func = bsa_hook + + +def build_lq(src, device="cuda", dtype=torch.bfloat16): + rdr = imageio.get_reader(src); total = rdr.count_frames() + idx = (list(range(total)) + [total - 1] * 4); F = largest_8n1_leq(len(idx)); idx = idx[:F] + frames = [] + for i in idx: + img = Image.fromarray(rdr.get_data(i)).convert("RGB").resize((SRC_W, SRC_H), Image.BICUBIC).resize((REF_W, REF_H), Image.BICUBIC) + t = torch.from_numpy(np.asarray(img, np.uint8)).to(device=device, dtype=torch.float32) + frames.append((t.permute(2, 0, 1) / 255.0 * 2.0 - 1.0).to(dtype)) + rdr.close() + return torch.stack(frames, 0).permute(1, 0, 2, 3).unsqueeze(0), F + + +def main(): + pipe = init_pipeline() + LQ, F = build_lq("./inputs/example0.mp4") + th, tw = REF_H, REF_W + with torch.no_grad(): + pipe(prompt="", negative_prompt="", cfg_scale=1.0, num_inference_steps=1, seed=0, + LQ_video=LQ, num_frames=F, height=th, width=tw, is_full_block=False, if_buffer=True, + topk_ratio=2.0 * 768 * 1280 / (th * tw), kv_ratio=3.0, local_range=11, color_fix=True) + print("\n=== Captured self-attention block_sparse calls @768x1408 ===") + for i, r in enumerate(records): + print(f"[{i}] q={r['q']} k={r['k']} mask={r['mask']} density={r['mask_density']}") + + +if __name__ == "__main__": + main() diff --git a/examples/WanVSR/profile_e2e_bottlenecks.py b/examples/WanVSR/profile_e2e_bottlenecks.py new file mode 100644 index 00000000..2bbd4547 --- /dev/null +++ b/examples/WanVSR/profile_e2e_bottlenecks.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""E2E kernel-level bottleneck profiler @ 768x1408 with the gemm conv3d backend. + +Goal: after the conv3d acceleration (now 1.86x A100), find where the remaining +denoise time goes so we can push past A100x3. Runs the v1.1 Tiny pipeline under +torch.profiler and aggregates CUDA self-time into categories (attention, GEMM +[linear/qkv/ffn], conv3d-gemm, norm/elementwise, layout/copy, decoder, other). + +Run from examples/WanVSR/ : + python profile_e2e_bottlenecks.py +""" +import os, time, importlib.util, re +import numpy as np +from PIL import Image +import imageio +import torch + +os.environ["FLASHVSR_CONV3D_BACKEND"] = "gemm" +import utils.utils as wanutils +wanutils._CONV3D_BACKEND = "gemm" + +_here = os.path.dirname(os.path.abspath(__file__)) +_spec = importlib.util.spec_from_file_location("infer_v1_1_tiny", os.path.join(_here, "infer_flashvsr_v1.1_tiny.py")) +_infer = importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_infer) +init_pipeline = _infer.init_pipeline +largest_8n1_leq = _infer.largest_8n1_leq + +REF_W, REF_H, SCALE = 768, 1408, 4 +SRC_W, SRC_H = REF_W // SCALE, REF_H // SCALE + + +def build_lq(src, device="cuda", dtype=torch.bfloat16): + rdr = imageio.get_reader(src); total = rdr.count_frames() + idx = (list(range(total)) + [total - 1] * 4) + F = largest_8n1_leq(len(idx)); idx = idx[:F] + frames = [] + for i in idx: + img = Image.fromarray(rdr.get_data(i)).convert("RGB").resize((SRC_W, SRC_H), Image.BICUBIC).resize((REF_W, REF_H), Image.BICUBIC) + t = torch.from_numpy(np.asarray(img, np.uint8)).to(device=device, dtype=torch.float32) + t = t.permute(2, 0, 1) / 255.0 * 2.0 - 1.0 + frames.append(t.to(dtype)) + rdr.close() + return torch.stack(frames, 0).permute(1, 0, 2, 3).unsqueeze(0), F + + +def categorize(name): + n = name.lower() + # attention kernels (incl. Triton WGMMA block-sparse kernel _bsfa[_tma]_kernel + # and the bundled block_sparse_attn CUDA kernel) + if any(k in n for k in ["fmha", "flash", "attention", "softmax", "scaled_dot", "mha", + "block_sparse", "_bsfa", "bsfa_"]): + return "attention" + # cuDNN convolution (TCDecoder conv2d, etc.), NOT the gemm-conv path. + # Note sm90_xmma_fprop_implicit_gemm is cuDNN conv (TCDecoder), not a linear GEMM. + if ("cudnn_convolution" in n or "implicit_gemm" in n or "xmma_fprop" in n + or "xmma_wgrad" in n or ("conv" in n and "convolution" in n)): + return "conv (cudnn, TCDecoder)" + # explicit GEMM (linear/qkv/ffn + our im2col conv -> addmm/nvjet/cublas) + if any(k in n for k in ["nvjet", "gemm", "cutlass", "wgmma", "cublas", "addmm", "matmul", "linear"]): + return "gemm (linear/ffn/im2col-conv)" + # layout conversions (big cost for conv2d in TCDecoder) + if any(k in n for k in ["nchwtonhwc", "nhwctonchw", "tonhwc", "tonchw"]): + return "layout (nchw<->nhwc)" + # normalization / elementwise / activation + if any(k in n for k in ["norm", "rms", "silu", "gelu", "relu", "elementwise", "mul", "add", "div", "layer_norm", "reduce", "sigmoid"]): + return "norm/elementwise/act" + # copy / cat / pad / transpose + if any(k in n for k in ["copy", "cat", "memcpy", "pad", "transpose", "permute", "contiguous"]): + return "copy/cat/pad" + # upsample / interpolation (decoder) + if any(k in n for k in ["upsample", "interpolate", "grid_sample"]): + return "decoder-resample" + return "other" + + +def main(): + src = os.environ.get("FLASHVSR_TEST_INPUT", "./inputs/example0.mp4") + pipe = init_pipeline() + LQ, F = build_lq(src) + th, tw = REF_H, REF_W + kwargs = dict(prompt="", negative_prompt="", cfg_scale=1.0, num_inference_steps=1, seed=0, + LQ_video=LQ, num_frames=F, height=th, width=tw, is_full_block=False, if_buffer=True, + topk_ratio=2.0 * 768 * 1280 / (th * tw), kv_ratio=3.0, local_range=11, color_fix=True) + + # warmup + with torch.no_grad(): + pipe(**kwargs) + torch.cuda.synchronize() + + from torch.profiler import profile, ProfilerActivity + with torch.no_grad(): + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + t0 = time.perf_counter() + pipe(**kwargs) + torch.cuda.synchronize() + wall = time.perf_counter() - t0 + + evs = prof.key_averages() + cats = {} + total_cuda = 0.0 + for e in evs: + st = e.self_device_time_total # us + if st <= 0: + continue + total_cuda += st + cats.setdefault(categorize(e.key), 0.0) + cats[categorize(e.key)] += st + + print(f"\n=== E2E bottleneck profile @ {tw}x{th}, gemm backend ===") + print(f"wall denoise: {wall*1e3:.0f} ms total CUDA self-time: {total_cuda/1e3:.0f} ms\n") + print(f"{'category':28s} {'CUDA ms':>10s} {'% of GPU':>9s}") + for k in sorted(cats, key=lambda x: -cats[x]): + print(f"{k:28s} {cats[k]/1e3:10.1f} {100*cats[k]/total_cuda:8.1f}%") + + print("\n=== Top 20 individual CUDA kernels ===") + print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=20)) + + +if __name__ == "__main__": + main() diff --git a/examples/WanVSR/profiling/ANALYSIS.md b/examples/WanVSR/profiling/ANALYSIS.md new file mode 100644 index 00000000..af869d80 --- /dev/null +++ b/examples/WanVSR/profiling/ANALYSIS.md @@ -0,0 +1,174 @@ +# FlashVSR v1.1 Tiny — GH200 Deep Profiling Analysis (Phase 1–3) + +Campaign date: 2026-07-08 · GH200 480GB (sm_90, 132 SM, 96GB HBM3) · driver 595.58.03 +CUDA 13.2 · torch 2.12 (NGC 26.05) · triton 3.7 · cuDNN 9.22 · nsys 2026.2.1 · ncu 2026.1.1 +Baseline config = all PR knobs ON (`gemm + NHWC + fuse_norm + triton attn + TMA + caches`). +Clocks locked at 1980 MHz (`nvidia-smi -lgc`), ncu run with `--clock-control none`. + +## 0. Headline numbers (untraced references, single runs) + +| Config | FPS | steady chunk | px-norm FPS | peak mem | +|---|---|---|---|---| +| 768x1408 full-knobs | **38.55** | 156.2 ms | 38.55 | 12.6 GiB | +| 1024x1920 full-knobs | 21.77 | 274.1 ms | 39.58 | 20.2 GiB | +| 1536x2560 full-knobs | 11.01 | 531.5 ms | 40.05 | 37.4 GiB | +| 768x1408 sparse attn | 35.35 | 176.2 ms | 35.35 | 12.6 GiB | + +px-norm FPS *rises* with resolution → the GPU is slightly under-fed at 768 but the +bottleneck structure is scale-invariant (attention ~47% everywhere). + +## 1. Where the time goes + +### 1.1 E2E wall budget (768x1408, F=81, traced shares match untraced totals) + +| Segment | time | share | +|---|---|---| +| chunk0+chunk1 (warm chunks) | ~516 ms | 26% | +| chunks 2..7 (steady, 6×156 ms) | ~937 ms | 47% | +| TCDecoder decode | ~343–430 ms | 17–21% | +| color_fix + misc python | ~10 ms | <1% | + +Decode is **fully serialized** after the denoise loop (single stream). Same shape at +1024/1536: decode = 22–23% of (chunks+decode). + +### 1.2 Steady chunk GPU ledger (@768, per chunk ≈156 ms wall, ≈151 ms GPU busy, idle 3.1%) + +| Phase | ms/chunk | % GPU | evidence | +|---|---|---|---| +| attn kernel `_bsfa_tma_kernel` (30×2.04 ms) | 58.7 | 39% | ncu: SM 40%, tensor 40%, occ 12.5% | +| attn transposes/copies (attn_core − kernel) | 11.6 | 7.7% | (S,n,d)→(n,S,d) `.contiguous()` ×3 + out | +| ffn (2 GEMM + gelu + gate) | 22.4 | 15% | GEMMs healthy (82% SOL) | +| rope apply (q,k) | 10.1 | 6.8% | elementwise, fusable | +| lq_conv1+conv2 (im2col+GEMM) | 11.7 | 7.8% | im2col copy = 29% of path | +| xattn (q/o GEMM + FA2 kernel) | 7.7 | 5.1% | `flash_fwd_kernel` 85 µs | +| mask_gen (pool+einsum+softmax+topk chain) | 7.4 | 4.9% | gatherTopK: SM 5%, 0.1 wave | +| qkv_norm (3 GEMM + RMSNorm) | 5.8 | 3.9% | | +| win_part / reorder / win_rev / cache_trim | 5.7 | 3.7% | copies at 66–81% DRAM BW | +| kv_cat (KV cache concat) | 3.4 | 2.3% | 3235 GB/s — at BW limit | +| mod1 + gate1 | 3.0 | 2.0% | fused kernels healthy (88% mem SOL) | +| head/patchify/unpatchify/lq_linears | ~1.5 | 1% | | +| **GPU idle within chunk** | **4.8** | **3.1%** | not launch-bound | + +## 2. Kernel deep-dives (ncu) + +### 2.1 `_bsfa_tma_kernel` — the #1 target (39% of GPU) +``` +duration 2.03 ms · SM SOL 40.2% · tensor pipe 40.2% (active 46.4%) +DRAM 134 GB/s (3.3%) · L2 hit 92.5% · achieved WGMMA math ≈ 393 TFLOP/s +occupancy 12.5% (1 block/SM: 178 reg/thr + 229 KB smem) · 8 warps/SM +stalls/issue: barrier 2.39 · wait 1.00 · short_sb 0.55 (of 6.37 total) +scheduler: 68.6% cycles with NO eligible warp +``` +Diagnosis: single-block-per-SM Triton pipeline; all 8 warps hit the same stage +barriers, nothing else to schedule → tensor pipe idles 60% of the time. +TMA=0 variant: 2.20 ms, long_scoreboard 0.87 (TMA removed it), 235 regs. + +Reference points at the exact shape (q 8448 × kv 25344, h12 d128): +| Kernel | time | note | +|---|---|---| +| cuDNN dense SDPA (full attention) | **1.86 ms** | 707 TF/s, computes 1.65× the FLOPs | +| ideal sparse = dense × 0.606 | **1.13 ms** | efficiency ceiling | +| `_bsfa_tma` (ours) | 2.03 ms | **56% of ideal** | +| FlexAttention + BlockMask (torch 2.12) | 1.92 ms | not the answer (59% of ideal) | + +→ A warp-specialized / 2-CTA / deeper-pipelined block-sparse kernel (FA3-style, +CUTLASS FMHA or hand-tuned Triton WS) has **~0.9 ms/call** headroom ⇒ ~26 ms/chunk. + +### 2.2 GEMMs — already near bf16 ceiling, FP8 is the lever +| kernel (nvjet) | avg | SM SOL | tensor | note | +|---|---|---|---|---| +| 256x128 coopA (qkv/o/ffn) | 88.7 µs | 82.7% | 82.7% | waves=1.0, perfect fit | +| 192x192 coopB (im2col conv) | 302.9 µs | 85.0% | 85.0% | 928 GB/s | + +FP8 microbench (torch._scaled_mm, M=8448): qkv/o ×1.59 · ffn1 ×1.55 · ffn2 ×1.72 +· lq_linear ×1.55 (1006–1377 TF/s). GEMM total ≈ 33 ms/chunk → FP8 saves ~13 ms/chunk. + +### 2.3 Elementwise & copies +Two classes: (a) already BW-bound (kv_cat 3235 GB/s = 81% HBM3; big copies 3268 GB/s) +→ only fix is *not doing the work*; (b) inefficient strided transposes +(65 µs × 8/chunk, SM-bound 71%, only 980 GB/s) — these are the triton-attn-path +`(S,n,d)→(n,S,d)` contiguous() calls → killable via kernel-side strides. + +### 2.4 mask_gen topk chain +gatherTopK 94.6 µs at **SM 5.2%, 0.1 waves** + 7–9 radix mini-kernels per call. +Pure latency, single fused kernel could do it in ~1 ms/chunk (now 7.4 ms). + +### 2.5 LQ projector conv (im2col+GEMM) +Path split @ conv2 shape: pad+cat 8% · **im2col copy 29%** · addmm 65% (707 TF/s). +cuDNN 9.22 direct conv3d at this shape: **152 ms vs 8.4 ms** (18× slower) — no new +engine on Hopper; GEMM path remains mandatory; fusing im2col into the GEMM +(CUTLASS conv or Triton fused) reclaims ~4 ms/chunk. + +### 2.6 Decoder +NHWC convs healthy; decode-window tensor pipe 37.8%, DRAM 16%. Main finding is +architectural: decode is 100% serial tail (see §3 H6). + +## 3. Hypothesis results + +| # | Hypothesis | Result | Ceiling @768 | +|---|---|---|---| +| H1 | launch-bound → CUDA Graphs | idle only 3.1%/chunk (0.1% @1024) | ≤4.8 ms/chunk | +| H2 | attn kernel inefficiency | 56% of ideal-sparse; barrier-stalled 1-CTA/SM | ~26 ms/chunk (bf16); more with FP8 attn | +| H3 | im2col overhead / cuDNN engine | im2col 29% of path; cuDNN still 18× slower | ~4 ms/chunk | +| H4 | GEMM efficiency / FP8 | bf16 at 82–85% SOL (no tuning left); FP8 ×1.55–1.72 | ~13 ms/chunk | +| H5 | elementwise BW | fused kernels near BW; transposes+rope+win copies removable | ~20 ms/chunk (transposes 11.6 + rope ~8) | +| H6 | decode serialization | decode = 17–23% of E2E, zero overlap today | +21–29% FPS if hidden | +| H7 | kv_cat + mask_gen | 3.4 + 7.4 ms/chunk | ~9 ms/chunk (ring buffer + fused topk) | +| H8 | power/clock residency | **700W platform cap** (900W denied); under load 649–687W, clocks sag to 1635–1815 MHz (84–92%) | efficiency wins compound; brute-force math won't | + +Extra: sparse backend (`block_sparse_attn`) steady chunk 176.2 ms vs triton 156.2 ms; +sparse also shows 8.2% idle — its per-call `torch.tensor(..., device=...)` cu_seqlens +creation is a hidden H2D sync the triton path avoids. + +## 4. Ranked Phase-2 roadmap (gain × confidence / effort) + +| # | Optimization | est. gain @768 (chunk 156 ms) | conf. | effort | +|---|---|---|---|---| +| 1 | **Attention kernel v2** (warp-specialized/2-CTA pipelined block-sparse; CUTLASS FMHA base or Triton WS; keep exact mask) | −26 ms | high (ref kernels prove it) | high | +| 2 | **Decoder overlap** (stream decode per-chunk on side stream, TCDecoder already streaming-capable) | +21–29% FPS E2E | high | med | +| 3 | **Kill attn-path transposes** (strided kernel IO or fold layout into kernel) | −11.6 ms | high | low-med | +| 4 | **FP8 GEMMs** (qkv/o/ffn/lq_linears via _scaled_mm, per-tensor scales, PSNR-gate) | −13 ms | med-high | med | +| 5 | **RoPE fusion** (single kernel for q+k, or fold into attn prologue; also cache per-chunk freqs — 44 ms/chunk CPU-side waste seen in traces) | −8 ms | high | low | +| 6 | **Fused topk mask_gen** (one kernel replaces gatherTopK+radix chain) | −5 ms | med | med | +| 7 | **Fused im2col-GEMM conv** (CUTLASS conv3d or Triton) | −4 ms | med | med | +| 8 | **KV ring buffer** (preallocated, no cat) | −3 ms | high | low | +| 9 | **CUDA Graphs on steady chunk** | −4.8 ms | med (shapes static per chunk) | med | +| 10 | FP8 attention (QK^T/PV in e4m3, needs quality gate) | −15–25 ms extra | low-med | high | + +Stacked (1,3,4,5,6,7,8,9 conservative): chunk 156 → ~90–100 ms ⇒ steady denoise +~80–89 FPS; with decoder overlap E2E ≈ **75–90 FPS** (~2–2.3× over 38.5) before +FP8-attention. Power cap (H8) will claw back some of this — efficiency-first +ordering maximizes what survives. + +## 5. Methodology notes +- Traced idle% is an upper bound; corrected against untraced per-chunk walls + (`[chunks]` line from `run_pipe_target.py`). nsys minimal trace (`cuda,nvtx`) + used for gap analysis; rich trace only for API attribution. +- ncu occupancy/SOL under `--clock-control none` with global 1980 MHz lock; + absolute TF/s ceilings should assume ~1650–1815 MHz effective under power cap. +- **ncu child-injection deadlock**: triton's `libcuda_dirs()` spawns `ldconfig`; + ncu tree-injection intermittently deadlocks that handshake (main python stuck in + `subprocess.communicate`). Fix baked into `ncu_run.sh`: `TRITON_LIBCUDA_PATH` + env + `--target-processes application-only`. +- nsys python-sampling produced no SAMPLING_CALLCHAINS on this aarch64 build; + irrelevant given idle ≈3%. + +## 6. Artifacts +``` +profiling/ + run_pipe_target.py env-driven target (W/H/F, steady window, per-chunk timing) + nsys_run.sh / ncu_run.sh / ncu_batch.sh drivers (dmon logging, deadlock fix) + analyze_gaps.py sqlite analyzer -> analysis.md/kernels.csv/phases.csv/gaps.csv/summary.json + ncu_extract.py .ncu-rep -> compact metric table (SOL/tensor/occ/stalls) + bench_ceilings.py H2/H3/H4 microbenches + reports/ + r1a_768_fullknobs/ nsys + gpu-metrics, full measured call (+analysis.md) + r1b_768_richapi/ nsys rich API trace, chunks 2-6 + r2_768_pysample/ (python sampling unavailable on aarch64) + r3_768_sparse/ sparse-attn single-diff + r4_1024_fullknobs/ r5_1536_fullknobs/ resolution shift + ncu/ attn_bsfa_tma{0,1}, gemms, elemwise, masktopk, decoder (+csv) +``` +NVTX instrumentation: `FLASHVSR_NVTX=1` (default OFF, zero-effect), steady window +via `FLASHVSR_PROFILER_START_CHUNK/STOP_CHUNK` (consumed per-call in +`flashvsr_tiny.py`). One behavioural fix: pipeline now honours `progress_bar_cmd`. diff --git a/examples/WanVSR/profiling/PHASE_BENCH_LOG.md b/examples/WanVSR/profiling/PHASE_BENCH_LOG.md new file mode 100644 index 00000000..41cfc2fa --- /dev/null +++ b/examples/WanVSR/profiling/PHASE_BENCH_LOG.md @@ -0,0 +1,1056 @@ +# FlashVSR Phase-2+ Benchmark Log + +Chronological, append-only log of every optimization attempt (successes, +failures, and reverts alike). Governed by the rules in +[`PHASE_ROADMAP.md`](./PHASE_ROADMAP.md) §0.5 and §5. + +**The gate:** do not start the next optimization until the current one has +(a) an entry in this file and (b) a 2–3 sentence interpretation. + +Measurement rules (short form): +- Untraced `run_pipe_target.py` runs only; 3 runs back-to-back, log the median. +- Before/after measured at the same commit, flag OFF vs flag ON. +- @768x1408 F=81 mandatory; @1536x2560 only at phase closure. +- Lossless claims: `max|diff| == 0`. Numeric-neutral claims: PSNR ≥ 49 dB. +- Traced (nsys/ncu) runs are attribution-only — never headline numbers. + +--- + +## Step 0 — Fresh Phase-2 baseline (MUST be filled before the first change) + +Command (from `examples/WanVSR/`): + +```bash +FLASHVSR_CONV3D_BACKEND=gemm FLASHVSR_TCDECODER_CHANNELS_LAST=1 \ +FLASHVSR_FUSE_NORM=1 FLASHVSR_ATTN_BACKEND=triton \ +FLASHVSR_CACHE_MOD=1 FLASHVSR_CACHE_MASK_BIAS=1 FLASHVSR_PROF_STEADY=off \ +/root/FlashVSR/venv/bin/python profiling/run_pipe_target.py +``` + +| Field | Value | +|---|---| +| Date/time | 2026-07-08 ~08:30 | +| Commit | df94d94 (phase1 instrumentation committed on top of 613bf9f) | +| GPU clocks locked | y (`nvidia-smi -lgc 1980,1980`) | +| Resolution / frames | 768x1408 / F=81 | +| Run 1 / Run 2 / Run 3 FPS | 38.585 / 38.525 / 38.594 | +| **Median FPS (= Phase-2 baseline)** | **38.585** | +| Median steady chunk ms | **156.28** (156.28 / 156.50 / 156.18) | +| Peak memory GiB | 12.6 | +| Reference (Phase-1 campaign, single run) | 38.55 FPS · 156.24 ms · 12.6 GiB | +| Notes | GPU otherwise idle, no compute apps; logs: `profiling/runs/phase2a/step0_baseline_run{1..3}.log`. Matches Phase-1 reference within noise. | + +--- + +## Per-change entries + + + +| Date | Phase | Optimization | Flag | FPS Before | FPS After | Delta | Steady Chunk Before | Steady Chunk After | Peak Mem Before | Peak Mem After | Correctness | Decision | +|------|-------|--------------|------|------------|-----------|-------|---------------------|--------------------|-----------------|----------------|-------------|----------| +| 2026-07-08 | 2A-1a | RoPE freqs device cache | `FLASHVSR_CACHE_ROPE_FREQS` | 38.585 | 38.592 | +0.02% | 156.28 ms | 156.30 ms | 12.60 GiB | 12.63 GiB | max\|diff\|=0 (bit-identical) | keep-behind-flag | +| 2026-07-08 | 2A-1b | Fused RoPE apply | `FLASHVSR_FUSE_ROPE` | 38.585 | 39.023 | +1.14% | 156.28 ms | 153.62 ms | 12.60 GiB | 12.60 GiB | max\|diff\|=0 (bit-identical) | keep-enabled | +| 2026-07-08 | 2A-2 | KV cache arena (no per-chunk cat) | `FLASHVSR_KV_RINGBUF` | 39.023 | 39.429 | +1.04% | 153.62 ms | 150.76 ms | 12.60 GiB | 15.50 GiB | max\|diff\|=0 over 9 chunks | keep-enabled | +| 2026-07-08 | 2A-3 | Attention strided IO (no transposes) | `FLASHVSR_ATTN_STRIDED_IO` | 39.429 | 41.099 | +4.24% | 150.76 ms | 141.14 ms | 15.50 GiB | 15.50 GiB | kernel + E2E max\|diff\|=0 | keep-enabled | +| 2026-07-08 | 2A-4 | mask_gen lean (kthvalue + no repeat + seqlens cache) | `FLASHVSR_MASKGEN_LEAN` | 41.099 | 41.477 | +0.92% | 141.14 ms | 139.00 ms | 15.50 GiB | 15.50 GiB | mask equality + E2E max\|diff\|=0 | keep-enabled | +| 2026-07-08 | 2A-5 | LQ projector lean (pad fold + no clones) | `FLASHVSR_LQPROJ_LEAN` | 41.477 | 41.580 | +0.25% | 139.00 ms | 138.46 ms | 15.50 GiB | 15.62 GiB | E2E max\|diff\|=0 (after dropping sub-item b) | keep-enabled | +| 2026-07-08 | 2A-6 | CUDA Graphs on steady chunk | (not implemented) | 41.580 | — | — | 138.46 ms | — | 15.62 GiB | — | n/a | postpone (go/no-go gate failed) | +| 2026-07-08 | 2B-1 | Decoder overlap on side CUDA stream | `FLASHVSR_DECODER_OVERLAP` | 41.662 | 42.373 | +1.71% | 138.30 ms | 190.58 ms (absorbs decode) | 15.62 GiB | 15.16 GiB | E2E max\|diff\|=0 (full+short clip, 3x repeats) | keep-behind-flag | +| 2026-07-08 | 2B-2 | FP8 GEMM infra (e4m3 `_scaled_mm` + Triton rowwise quant) | `FLASHVSR_FP8_GEMM` | 41.704 | 42.986 | +3.07% | 137.98 ms | 132.96 ms | 15.62 GiB | 16.65 GiB | PSNR 40.70 dB (full scope; NOT lossless by design — Phase-4 gate) | keep-behind-flag (Phase-4 enable gate) | +| 2026-07-08 | 2B-3 | Fused mask-gen threshold-select (Triton radix select+compare) | `FLASHVSR_FUSED_MASKGEN` (removed) | 42.00 | 41.84 | −0.4% (harness E2E) | — | — | 15.62 GiB | 15.62 GiB | mask + E2E max\|diff\|=0 (exact) | **revert** (bit-exact but performance-negative) | +| 2026-07-08 | 3 | Warp-specialized block-sparse attention v2 (Gluon producer/consumer + pingpong) | `FLASHVSR_ATTN_BACKEND=triton2` | 41.692 | 45.466 | **+9.05%** | 137.96 ms | 122.06 ms | 15.62 GiB | 15.62 GiB | kernel cos ≥0.999995 vs block_sparse (all densities + degenerates); E2E PSNR 50.03 dB; repeats bit-identical; default paths max\|diff\|=0 | keep-behind-flag (ncu gate 3/4; residency 12 warps/SM WS vs ≥16 letter) | +| 2026-07-08 | 3.5-1a | Attention v2 FFMA-form softmax (scaled-max domain) | (in `triton2`) | 45.466 | 46.112 | +1.42% | 122.06 ms | 119.92 ms | 15.62 GiB | 15.62 GiB | kernel cos 0.999996 all densities+degenerate; E2E PSNR 50.08 dB; repeat bit-identical | keep (in triton2) | +| 2026-07-08 | 3.5-1b | alpha==1 bit-exact rescale skip | (experiment) | 46.112 | — | — | — | — | — | — | bit-exact but −7% isolated kernel | **drop** (branch+reduce > overlapped FMULs) | +| 2026-07-08 | 3.5-1c | explicit pingpong named-barrier | (experiment) | 46.112 | — | — | — | — | — | — | n/a | **drop** (deadlock risk; already ~83% peak, HGMMA only 13.6%) | +| 2026-07-08 | 3.5-2 | Fused CSR build (sort-free cumsum-scatter) | `FLASHVSR_FUSED_CSR` | 46.112 | 46.383 | +0.59% | 119.92 ms | 118.84 ms | 15.62 GiB | 15.62 GiB | idx[0:cnt]+cnt bit-eq vs argsort (all densities+degenerate); v2 output max\|diff\|=0 | keep-behind-flag (lossless; recommended-with-triton2) | +| 2026-07-08 | 3.5-3 | Decoder overlap × v2 remeasure | `FLASHVSR_DECODER_OVERLAP` | 46.383 | 47.239 | +1.85% | 118.84 ms | 172.9 ms (absorbs decode) | 15.62 GiB | 15.16 GiB | inherited lossless (2B-1, backend-orthogonal); tail 543→126 ms | keep-behind-flag (co-exec 31.7% vs 34.1%; v2 did NOT free the ceiling) | +| 2026-07-08 | 3.5-4 | RoPE freqs device cache × v2 remeasure | `FLASHVSR_CACHE_ROPE_FREQS` | 46.383 | 46.429 | +0.10% (noise) | 118.84 ms | 118.40 ms | 15.62 GiB | 15.65 GiB | lossless (2A-1a) | keep-behind-flag (still FPS-neutral @768 under v2) | +| 2026-07-08 | 4A | RoPE fp32 (fused apply) | `FLASHVSR_ROPE_FP32` (reverted) | — | — | ~0 (x1.01 isolated) | — | — | — | — | ≤1 bf16 ULP (numerically fine) | **drop** (no speedup — rope is interleave-access-bound at ~400 GB/s, not fp64-bound) | +| 2026-07-08 | 4B | FP8 GEMM blockwise (`_scaled_mm_v2` 1x128 act + 1x128 weight) | `FLASHVSR_FP8_GEMM_MODE=blockwise` | — | — | isolated GEMM ×1.13–1.34 | — | — | — | +1.0 GiB | full-scope PSNR **41.00 dB** (ffn 42.71 / qkv,o,lq 43.51 / lq 49.49) | keep-behind-flag, **NOT enable-eligible** (only speed-neutral `lq` clears ≥49) | +| 2026-07-08 | 4C | FP8 attention (e4m3 K/V) quality de-risk probe | `FLASHVSR_ATTN_FP8_PROBE` (reverted) | — | — | — | — | — | — | — | E2E PSNR **34.96 dB** (optimistic bound) | **drop** (gate unreachable — attention far more e4m3-sensitive than GEMMs; no kernel written) | +| 2026-07-08 | 5-P1 | Coalesced RoPE kernel (u32 pairs + direct complex128 reads) | `FLASHVSR_ROPE_KERNEL=triton` | 46.383 | 47.780 | **+3.01%** | 118.84 ms | 113.72 ms | 15.62 GiB | 15.62 GiB | kernel bit-eq (triton==fused==eager, both shapes); E2E max\|diff\|=0 | keep (recommended) | +| 2026-07-08 | 5-P2 | Incremental pooled-K cache (rolling mean mirror) | `FLASHVSR_POOLED_K_CACHE` | 47.780 | 47.858 | +0.16% | 113.72 ms | 113.08 ms | 15.62 GiB | 15.64 GiB | probe: batched mean bit-eq; E2E max\|diff\|=0 + repeat bit-eq | keep (recommended) | +| 2026-07-08 | 5-P3 | Zero-copy V windowing (rank-4 TMA raster arena + raster out) | `FLASHVSR_ATTN_ZEROCOPY` | 47.858 | 48.268 | +0.86% | 113.08 ms | 112.24 ms | 15.64 GiB | 15.64 GiB | kernel zc-vs-std bit-eq (incl. ring offset); E2E max\|diff\|=0 + repeat bit-eq | keep (recommended, requires triton2) | +| 2026-07-08 | 5-P4a | Decoder overlap on the P1-P3 stack | `FLASHVSR_DECODER_OVERLAP` | 48.268 | 49.237 | **+2.01%** | 112.24 ms | 164.88 ms (absorbs decode) | 15.64 GiB | 15.18 GiB | lossless (2B-1, unchanged code) | **promote** (production set; keep OFF for per-chunk attribution runs) | +| 2026-07-08 | 5-P4b | cudnn.benchmark=True (decoder conv autotune) | (no code) | 48.12 | 48.08 | −0.1% (noise) | — | — | — | — | output bit-identical (same algos picked) | **drop** (heuristics already optimal) | +| 2026-07-09 | 6-P1 | TCDecoder recurrent-state pointer rotation | `FLASHVSR_TCDECODER_POINTER_STATE` | 49.229 | 49.421 | +0.39% | overlap metric | overlap metric | 15.18 GiB | 15.18 GiB | max\|diff\|=0, F=25/F=89 | keep | +| 2026-07-09 | 6-P2 | Direct decoder output placement | `FLASHVSR_TCDECODER_DIRECT_OUTPUT` | 49.421 | 49.734 | +0.63% | overlap metric | overlap metric | 15.18 GiB | 15.44 GiB | max\|diff\|=0, F=25/F=29/F=89 | keep | +| 2026-07-09 | 6-P3 | Exact MemBlock pointwise fusion | `FLASHVSR_TCDECODER_FUSE_POINTWISE` | 49.734 | 50.945 | +2.43% | overlap metric | overlap metric | 15.44 GiB | 15.44 GiB | max\|diff\|=0 incl. BF16 special values | keep | +| 2026-07-09 | 6-P4 | Channels-last nearest upsample kernel | `FLASHVSR_TCDECODER_UPSAMPLE` | 50.945 | 52.302 | +2.66% | overlap metric | overlap metric | 15.44 GiB | 15.44 GiB | max\|diff\|=0; isolated 7.7–8.6x | keep | +| 2026-07-09 | 6-P5 | Exact LQ Conv3D im2col packer | `FLASHVSR_CONV3D_PACKER=triton` | 52.302 | 53.009 | +1.35% | overlap metric | overlap metric | 15.44 GiB | 15.44 GiB | exact patch matrix + E2E max\|diff\|=0 | keep | +| 2026-07-09 | 6-P6 | Channels-last recurrent concat kernel | `FLASHVSR_TCDECODER_CONCAT` | 53.009 | 53.423 | +0.78% | overlap metric | overlap metric | 15.44 GiB | 15.44 GiB | max\|diff\|=0; isolated 1.7–1.8x | keep | +| 2026-07-09 | 6b-P1 | Reordered TGrow->unpack+upsample fusion | `FLASHVSR_TCDECODER_TGROW_UP` | 53.40 | 54.61 | **+2.27%** | overlap metric | overlap metric | 15.44 GiB | 15.6 GiB | max\|diff\|=0 E2E (F=25/29/89 @768, F=25/29/41 @1536); isolated 2.2–5.4x; @1536 F=41 spot 14.36→14.69 (+2.30%) | keep | +| 2026-07-09 | 6b-P2 | cuDNN runtime-fused Conv+Bias(+Add)+ReLU | `FLASHVSR_TCDECODER_CUDNN_FUSED` | 54.58 | 55.91 | **+2.44%** | overlap metric | overlap metric | 15.6 GiB | 15.6 GiB | **quality-gated**: E2E 55.4–55.8 dB PSNR (gate 49); isolated 70 dB, ~10–14% of values ±1 ULP; @1536 F=41 spot 14.69→15.07 (+2.59%) | keep (quality-gated) | +| 2026-07-09 | 6b-P3 | Triton `tl.dot` RGB tail conv (Cout=3, N=16-pad) | (prototype only) | 54.58 | — | −4% isolated | — | — | — | — | 65 dB isolated; best 0.351 ms vs cuDNN 0.335 ms — L2-BW bound at 9x activation re-reads | **drop** (fails ≥15% isolated bar) | + +#### 2026-07-08 09:00 · Phase 2A-1a · RoPE freqs device cache + +- Commit / patch: on top of df94d94 (committed as phase2a rope freqs cache) +- Files changed: `diffsynth/pipelines/flashvsr_tiny.py` (+ new `test_phase2a_lossless.py` harness) +- Flag: `FLASHVSR_CACHE_ROPE_FREQS` (default OFF) +- Env vars used (full set): full-knob baseline + flag under test +- Exact benchmark command: §0.4 primary command + `FLASHVSR_CACHE_ROPE_FREQS=1` +- Resolution / frames: 768x1408 / F=81 (spot-check 1536: n) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 38.585 → 38.592 (+0.02%, noise) — 3-run medians +- Steady chunk before → after (Δ): 156.28 → 156.30 ms (noise) +- Peak mem before → after: 12.60 → 12.63 GiB (+30 MB: device freqs buffers f=6 and f=2) +- Correctness: `test_phase2a_lossless.py CACHE_ROPE_FREQS` → max|diff| == 0 (PASS) +- Nsight report path: n/a (untraced only) +- Decision: keep-behind-flag +- Interpretation: the ~44 ms/chunk `rope_freqs` cost seen in *traced* runs is CPU + wall that rides entirely under the 156 ms GPU chunk in untraced runs, so + removing it does not move FPS @768 — consistent with the roadmap note that + this is "CPU-side headroom", not GPU time. The change is bit-identical and + removes per-chunk CPU tensor construction + an 8.6 MB H2D, which matters for + CPU-loaded deployments and is a precondition for CUDA-graph capture (2A-6), + so it stays available behind its flag rather than being reverted. + +#### 2026-07-08 09:25 · Phase 2A-1b · Fused RoPE apply + +- Commit / patch: phase2a fused rope apply +- Files changed: `diffsynth/models/wan_video_dit.py` (`rope_apply` + compiled impl) +- Flag: `FLASHVSR_FUSE_ROPE` (default OFF) +- Env vars used (full set): full-knob baseline + `FLASHVSR_FUSE_ROPE=1` +- Exact benchmark command: §0.4 primary command + `FLASHVSR_FUSE_ROPE=1` +- Resolution / frames: 768x1408 / F=81 (spot-check 1536: at phase closure) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 38.585 → 39.023 (+1.14%) — 3-run medians +- Steady chunk before → after (Δ): 156.28 → 153.62 ms (−2.66 ms) +- Peak mem before → after: 12.60 → 12.60 GiB +- Correctness: kernel-level max|diff|=0 on real shape; E2E + `test_phase2a_lossless.py FUSE_ROPE` → max|diff| == 0 (exceeds the ≥49 dB gate) +- Isolated kernel: eager 0.181 ms/call → fused 0.128 ms/call (×1.41, 60 calls/chunk) +- Nsight report path: n/a (untraced only) +- Decision: keep-enabled (joins recommended set) +- Interpretation: got −2.7 ms of the −8 ms roadmap ceiling: the estimate double + counted freqs-side effects, and the fp64 multiply itself (not just the + materialized intermediates) is part of the cost, so the fused kernel is + compute-bound at ~0.13 ms/call rather than pure-BW ~0.02 ms. Output is + bit-identical since the same fp64 operations are performed in one kernel. + Remaining rope headroom (folding the apply into the attention prologue or + fp32 freqs) is Phase-4 territory (numerics change). + +#### 2026-07-08 09:55 · Phase 2A-2 · KV cache arena (kv_cat removal) + +- Commit / patch: phase2a kv cache arena +- Files changed: `diffsynth/models/wan_video_dit.py` (`_KVArena`, `SelfAttention.forward`) +- Flag: `FLASHVSR_KV_RINGBUF` (default OFF); tuning: `FLASHVSR_KV_RINGBUF_SPARE` (default 2) +- Env vars used (full set): full-knob baseline + `FLASHVSR_FUSE_ROPE=1` (kept set) + flag under test +- Exact benchmark command: §0.4 primary command + `FLASHVSR_FUSE_ROPE=1 FLASHVSR_KV_RINGBUF=1` +- Resolution / frames: 768x1408 / F=81 (spot-check 1536: at phase closure) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 39.023 → 39.429 (+1.04%) — 3-run medians +- Steady chunk before → after (Δ): 153.62 → 150.76 ms (−2.86 ms) +- Peak mem before → after: 12.60 → 15.50 GiB (+2.9 GiB = 2 spare arena slots × 60 tensors) +- Correctness: `test_phase2a_lossless.py KV_RINGBUF` → max|diff| == 0 over a + 9-chunk clip (exercises slot rotation AND tail compaction) +- Nsight report path: n/a (untraced only) +- Decision: keep-enabled (joins recommended set) +- Interpretation: recovers essentially the whole kv_cat budget (−2.9 of the + 3.4 ms attribution): new windows are partition-written straight into the + arena (no extra traffic) and only the amortized compaction remains — visible + as ~+2.5 ms on every 3rd chunk in the per-chunk trace. Bit-identical since + the live view preserves value order and contiguity exactly. Cost is +2.9 GiB + retained memory (documented; `FLASHVSR_KV_RINGBUF_SPARE` trades memory for + compaction frequency, and the flag restores the old path entirely). + +#### 2026-07-08 10:20 · Phase 2A-3 · Attention-path strided IO + +- Commit / patch: phase2a attention strided IO +- Files changed: `diffsynth/models/triton_block_sparse_attn.py` + (`_bsfa_tma_kernel_snd`, `triton_block_sparse_attention_snd`), + `diffsynth/models/wan_video_dit.py` (glue branch in `flash_attention`) +- Flag: `FLASHVSR_ATTN_STRIDED_IO` (default OFF) +- Env vars used (full set): full-knob baseline + kept set + (`FUSE_ROPE=1 KV_RINGBUF=1`) + flag under test +- Exact benchmark command: §0.4 primary + `FLASHVSR_FUSE_ROPE=1 FLASHVSR_KV_RINGBUF=1 FLASHVSR_ATTN_STRIDED_IO=1` +- Resolution / frames: 768x1408 / F=81 (spot-check 1536: at phase closure) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 39.429 → 41.099 (+4.24%) — 3-run medians +- Steady chunk before → after (Δ): 150.76 → 141.14 ms (−9.62 ms) +- Peak mem before → after: 15.50 → 15.50 GiB +- Correctness: kernel-level max|diff| == 0 vs contiguous path on the real + shape (both TMA and non-TMA variants); E2E + `test_phase2a_lossless.py ATTN_STRIDED_IO` → max|diff| == 0 (exceeds the + ≥49 dB gate) +- Isolated: contiguous glue+kernel 2.851 ms/call → strided 2.514 ms/call; + strided is even faster than the kernel alone on pre-copied inputs + (2.558 ms) — per-head tiles are adjacent in the (S, n*d) layout, so TMA + loads of neighbouring heads share L2 lines +- Nsight report path: n/a (untraced only) +- Decision: keep-enabled (joins recommended set) +- Interpretation: −9.6 of the −11.6 ms ceiling: all four transpose copies are + gone; the residual is the win_part/win_rev reorder that was counted in the + same bucket. Bit-identical because the tile schedule and accumulation order + are unchanged — only addressing moved from flattened (H·S, D) descriptors + to 2D (S, H·D) descriptors with per-head column offsets. TMA=0 fallback is + also strided (stride-general kernel), and any failure falls back to the + contiguous triton path, not to the sparse backend. + +#### 2026-07-08 10:45 · Phase 2A-4 · mask_gen allocation / sync cleanup + +- Commit / patch: phase2a mask_gen lean +- Files changed: `diffsynth/models/wan_video_dit.py` + (`generate_draft_block_mask`, `flash_attention` sparse branch) +- Flag: `FLASHVSR_MASKGEN_LEAN` (default OFF) +- Env vars used (full set): full-knob baseline + kept set + (`FUSE_ROPE=1 KV_RINGBUF=1 ATTN_STRIDED_IO=1`) + flag under test +- Exact benchmark command: §0.4 primary + kept set + `FLASHVSR_MASKGEN_LEAN=1` +- Resolution / frames: 768x1408 / F=81 +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 41.099 → 41.477 (+0.92%) — 3-run medians +- Steady chunk before → after (Δ): 141.14 → 139.00 ms (−2.14 ms) +- Peak mem before → after: 15.50 → 15.50 GiB +- Correctness: threshold + boolean-mask exact equality vs topk on random / + heavy-tie / softmax-distributed inputs at the real shape; E2E + `test_phase2a_lossless.py MASKGEN_LEAN` → max|diff| == 0 +- Isolated: topk(k=7920 of 17424) 0.187 ms → kthvalue 0.076 ms per call + (30 calls/chunk); plus removal of the boolean-mask repeat copy; sparse + backend additionally gets persistent cu_seqlens/head_mask_type (hidden H2D + sync removed — benefits the sparse fallback path, not the triton bench) +- Nsight report path: n/a (untraced only) +- Decision: keep-enabled (joins recommended set) +- Interpretation: −2.1 ms banked (roadmap estimated −1–2 ms for the lean pass; + the kthvalue swap alone projected −3.3 ms isolated but part of the chain is + latency that overlaps with neighbouring small kernels). Mask semantics are + provably unchanged — same order statistic, ties behave identically, strict + `>` compare untouched. The remaining ~5 ms of the mask_gen chain needs the + fused top-k kernel (Phase 2B-3), not more hygiene. + +#### 2026-07-08 11:10 · Phase 2A-5 · LQ projector allocation/layout cleanup + +- Commit / patch: phase2a lq projector lean +- Files changed: `examples/WanVSR/utils/utils.py` (`CausalConv3d._forward_lean`, + `_conv3d_gemm(contig_out=)`, `Causal_LQ4x_Proj.stream_forward`) +- Flag: `FLASHVSR_LQPROJ_LEAN` (default OFF) +- Env vars used (full set): full-knob baseline + kept set + (`FUSE_ROPE=1 KV_RINGBUF=1 ATTN_STRIDED_IO=1 MASKGEN_LEAN=1`) + flag under test +- Exact benchmark command: §0.4 primary + kept set + `FLASHVSR_LQPROJ_LEAN=1` +- Resolution / frames: 768x1408 / F=81 +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 41.477 → 41.580 (+0.25%) — 3-run medians +- Steady chunk before → after (Δ): 139.00 → 138.46 ms (−0.54 ms, ~noise floor) +- Peak mem before → after: 15.50 → 15.62 GiB (persistent pad buffers + view retention) +- Correctness: E2E `test_phase2a_lossless.py LQPROJ_LEAN` → max|diff| == 0 + across a full clip (streaming cache exercised across chunk boundaries) +- Sub-item REJECTED during development: skipping the GEMM output + `.contiguous()` (layout-only in values) changed `F.normalize`'s reduction + accumulation order → measured max|diff|=0.397 / PSNR 49.9 dB, which + violates this item's lossless gate. Dropped; documented in code. Could be + revisited under a Phase-4 numerics-tolerant gate if the projector becomes + hot again. +- Nsight report path: n/a (untraced only) +- Decision: keep-enabled (joins recommended set) +- Interpretation: the kept copies-only cleanup (fold cat+F.pad into one + buffer write, drop streaming-cache clones) recovers ~0.5 ms of the ~1–2 ms + estimate — pad+cat was only ~8% of the projector path, and the addmm (65%) + and im2col gather (29%) are untouched by design. Since 2A-5 recovered + <2 ms, the roadmap gate keeps Phase 2B-4 (fused im2col-GEMM) on the table. + +#### 2026-07-08 11:40 · Phase 2A-6 · CUDA Graphs go/no-go gate → NO-GO (postpone) + +- Commit / patch: none (gate evaluation only, per roadmap §2A-6) +- Nsight report path: `profiling/reports/p2a_stack_768/` (nsys minimal trace, + full 2A stack, steady window chunks 2–6; attribution-only, not headline FPS) +- Gate input 1 — idle: corrected steady-chunk idle with the full 2A stack is + 0.6–2.5% (avg ~1.8%, `analysis.md` busy/idle table) → below the ≥2% go + threshold; ceiling ≤ ~3 ms/chunk @768 and ~0 at higher resolutions. +- Gate input 2 — capture safety: the steady body is NOT capture-stable + today: (a) the 2A-2 KV arena advances a Python-int slice offset every chunk + and compacts on a data-dependent schedule (a captured graph would freeze + both), (b) LQ conditioning slices a different video range per chunk, + (c) the 2A-1a freqs buffer is rewritten per chunk (solvable — update + outside the graph — but only relevant if (a)/(b) were solved). Fixing (a) + requires a true ring buffer with device-side index remapping — exactly the + "invasive changes" this experiment is instructed not to expand into. +- Decision: postpone (documented no-go; do not implement in 2A) +- Interpretation: after the 2A cleanup the chunk is even less launch-bound + than the 3.1% Phase-1 measurement, so the graphs ceiling shrank while its + implementation cost grew (arena offsets). Attribution cross-checks the + stack wins: kv_cat memcpy 3.4 → ~0.5 ms/chunk, rope 10.1 → ~7.4, mask_gen + 7.4 → ~4.4, attn transposes gone, `_bsfa_tma_kernel_snd` unchanged at + 2.03 ms/call. Revisit graphs only if Phase 2B/3 work makes the steady body + static (or if deployment shows CPU contention). + +#### 2026-07-08 · Phase 2B-1 · Step 2 — fresh Phase-2A-stack baseline (2B-1 starting point) + +Re-measured at `45b2ddd` (pre-change), full 2A recommended stack ON, decoder +overlap absent. Command = §0.4 primary + `FLASHVSR_FUSE_ROPE=1 +FLASHVSR_KV_RINGBUF=1 FLASHVSR_ATTN_STRIDED_IO=1 FLASHVSR_MASKGEN_LEAN=1 +FLASHVSR_LQPROJ_LEAN=1`; clocks locked 1980 MHz. + +| Field | Value | +|---|---| +| Run 1 / 2 / 3 FPS | 41.780 / 41.759 / 41.686 | +| **Median FPS (= 2B-1 baseline)** | **41.759** | +| Median steady chunk ms | 137.46 (137.46 / 136.90 / 137.96) | +| Peak memory GiB | 15.62 | +| Reference (Phase 2A closure) | 41.580 FPS · 138.46 ms · 15.62 GiB | +| Notes | Matches the 2A closure within noise (+0.4% FPS). Logs: `profiling/runs/phase2b/step2_baseline_run{1..3}.log`. | + +#### 2026-07-08 12:40 · Phase 2B-1 · Decoder overlap on a side CUDA stream + +- Commit / patch: phase2b decoder overlap. Note: a concurrent duplicate + session committed an equivalent variant of this same 2B-1 change as + 7cecb65 mid-campaign; this entry and the follow-up commit supersede it + (same design, independently implemented and fully re-validated — the + numbers in THIS entry correspond to the tree at the follow-up commit). +- Files changed: `diffsynth/pipelines/flashvsr_tiny.py` (flag + overlap path), + `examples/WanVSR/profiling/run_pipe_target.py` (additive `[tail]` + post-loop-ms print), new `examples/WanVSR/profiling/test_decoder_overlap_lossless.py` +- Flag: `FLASHVSR_DECODER_OVERLAP` (default OFF; serialized end-of-loop decode + path is byte-identical code when OFF) +- Design: after each chunk's `cur_latents -= noise_pred` a ready-event is + recorded on the main stream; a persistent side stream waits on it and runs + `TCDecoder.decode_video` for that chunk (cond slice `LQ_pre_idx:LQ_cur_idx`, + same per-chunk semantics as `flashvsr_tiny_long.py`); per-chunk done-events + are waited on by the main stream only once, right before output assembly + (`wait_event`, GPU-side, no CPU sync); decoded chunks are assembled from a + chunk-id-indexed list → output order is structural, never completion order. + TAEHV mem-block state is only ever touched by the decode stream; decode + inputs stay alive in `latents_total`; `record_stream` guards added as + defense in depth. +- Env vars used (full set): full-knob baseline + 2A kept set + flag under test +- Exact benchmark command: §0.4 primary + 2A kept set + `FLASHVSR_DECODER_OVERLAP=1` +- Resolution / frames: 768x1408 / F=81 (spot-check 1536x2560: y, below) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 41.662 → 42.373 (+1.71%) — 3-run medians, same commit +- Steady chunk before → after (Δ): 138.30 → 190.58 ms (+52.3 ms — the chunk + wall now absorbs ~50 ms/chunk of decode GPU work + ~12 ms CPU enqueue; this + metric is no longer denoise-only when the flag is ON, see interpretation) +- Decode tail (post-loop) before → after: 561.2 → 125.1 ms (−436 ms hidden) +- Peak mem before → after: 15.62 → 15.16 GiB (−0.46 GiB: the one-shot + 20-frame decode working set and the full-latents cat are gone) +- Correctness: `test_decoder_overlap_lossless.py` → max|diff| == 0, + mean|diff| == 0, 85/85 frames bit-equal (ordering), shape/dtype/device + identical, on BOTH full clip (F=89, 9 chunks) and short clip (F=25, + 1 chunk, trim edge); overlap run 3x back-to-back → all repeats bit-identical + (no concurrency races). `test_phase2a_lossless.py ALL` re-run → PASS + (max|diff|=0, 2A stack unaffected). +- Nsight report path: `profiling/reports/phase2b_decoder_overlap/` + (`profile.nsys-rep`, `profile.sqlite`, `overlap_analysis.txt`). Evidence: + decoder kernels (1060 cuDNN convs, 498.7 ms busy) on side stream 13 vs + denoise on stream 7; decode active across the whole loop span, not as a + tail; 170.2 ms (34.1%) of decode busy time co-executes with denoise + kernels; `decode_wait` = 0.1 ms and `color_fix` = 1.6 ms at the end (no + per-chunk sync, final sync negligible); GPU idle 1.9% of span. +- Decision: keep-behind-flag (default OFF) +- Interpretation: the mechanics work exactly as designed — bit-identical + output, decode tail fully hidden (561 → 125 ms), zero hidden syncs, and + peak memory *drops* — but the E2E gain is +1.7% @768 / +2.0% @1536, far + under the roadmap's +21–29% ceiling. Nsight shows why: the GPU is + work-conserving (idle 1.9%), and only ~34% of decode kernel time can + co-execute with denoise because the `_bsfa` attention kernel owns a full + SM (229 KB smem, 1 CTA/SM) and the decoder's cuDNN conv grids are + SM-saturating themselves, so the hardware time-shares instead of + co-executing — moving the decode from the tail into the chunks conserves + total GPU work and only the co-executed 170 ms (+ warm-chunk gap fill) is + actually won. The roadmap estimate implicitly assumed decode could ride on + spare SM capacity that doesn't exist under the current attention kernel. + Re-evaluate after Phase 3: the planned 2-CTA / warp-specialized attention + kernel frees smem headroom, which should raise the co-execution fraction + substantially — the flag composes losslessly with everything, so it can be + promoted then. Default stays OFF also for campaign hygiene: with the flag + ON the `[chunks]` steady metric measures denoise+decode contention rather + than denoise alone, which would muddy per-chunk attribution for 2B-2/3/4. + +##### Decoder-overlap detail (768x1408 F=81 medians; 1536x2560 single runs) + +| Config | E2E FPS | Steady Chunk Time | Decode Tail Time | Peak Memory | Notes | +|--------|---------|-------------------|------------------|-------------|-------| +| 768 · 2A stack, overlap OFF | 41.662 | 138.30 ms | 561.2 ms | 15.62 GiB | serialized decode after loop | +| 768 · 2A stack, overlap ON | 42.373 | 190.58 ms | 125.1 ms | 15.16 GiB | decode rides inside chunks; tail = last-chunk decode remainder + color fix | +| 1536 · 2A stack, overlap OFF | 11.485 | 502.46 ms | 2046.2 ms | 48.39 GiB | matches 2A closure spot-check (11.489) | +| 1536 · 2A stack, overlap ON | 11.712 | 682.66 ms | 470.3 ms | 46.72 GiB | +2.0% FPS; decode share is larger at high res but co-execution stays SM-bound | + +#### 2026-07-08 13:55 · Phase 2B-2 · FP8 GEMM infrastructure (`torch._scaled_mm`) + +- Commit / patch: phase2b fp8 gemm infrastructure (on top of 41f1d8e) +- Files changed: new `diffsynth/models/fp8_gemm.py` (flag, Triton rowwise + quantizer, gelu-fused quantizer, weight pre-cast cache, `_scaled_mm` + wrappers with sticky eager fallback), `diffsynth/models/wan_video_dit.py` + (qkv shared-quant site, o site, ffn site), `examples/WanVSR/utils/utils.py` + (LQ per-layer linears, one shared quant for 30 GEMMs), new + `examples/WanVSR/test_fp8_gemm_quality.py`, new + `examples/WanVSR/profiling/sweep_fp8_scope.py` (per-site attribution tool + for the Phase-4 audit) +- Flag: `FLASHVSR_FP8_GEMM` (default OFF — **permanently until the Phase-4 + quality protocol clears it**); scope bisection via + `FLASHVSR_FP8_GEMM_SCOPE=qkv,o,ffn,lq` (+`ffn1` token, not in default set) +- Design: weights pre-cast once to e4m3 with per-out-channel scales (lazy, + cached on module, +1.03 GiB); activations quantized per call with per-row + scales by ONE Triton kernel (row amax + cast in a single launch — the + eager/compiled quantize chain measured ~5x off bandwidth, 113-272 µs, and + erased the GEMM win; the kernel does 15.1 µs @ 8448x1536); ffn2's input + quant is fused with GELU-tanh (fp32 gelu -> e4m3, replacing the eager bf16 + GELU pass — on the 8960-wide activation this flips ffn2 from net-loss to + net-win); bias fused in `_scaled_mm` epilogue; `use_fast_accum=False` +- Env vars used (full set): full-knob baseline + 2A kept set + flag under test +- Exact benchmark command: §0.4 primary + 2A kept set + `FLASHVSR_FP8_GEMM=1` +- Resolution / frames: 768x1408 / F=81 (1536 spot-check: deferred to phase + closure; flag is default-OFF anyway) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 +- FPS before → after (Δ): 41.704 → 42.986 (+3.07%) — 3-run medians +- Steady chunk before → after (Δ): 137.98 → 132.96 ms (−5.02 ms) +- Peak mem before → after: 15.62 → 16.65 GiB (+1.03 GiB e4m3 weight copies) +- Correctness / quality (Phase-4 protocol measurement, example0 @768): + full scope PSNR **40.70 dB** (max|d|=0.875, mean|d|=0.0114) → + revert-or-redesign tier for *enablement*. Per-site sweep + (`sweep_fp8_scope.py`): qkv 45.58 dB (and net-SLOWER in-pipe) · o 45.05 · + ffn 42.37 (the speed lever AND the quality offender) · ffn1-only 44.15 · + lq 49.08 (recommended tier, speed-neutral) · qkv,o,lq 43.42 · + o,ffn,lq 41.34 — every meaningful-speed combination lands < 45 dB. + Flag-OFF neutrality: `test_phase2a_lossless.py ALL` max|diff|=0 PASS, + `test_decoder_overlap_lossless.py` PASS (default paths untouched). +- Kernel evidence: torch.profiler + ncu on the live pipeline steady chunk + (`profiling/reports/ncu/phase2b2_fp8_kernels.ncu-rep`): e4m3 GEMMs dispatch + `nvjet_sm90_qqtst_128x128_128x6_..._ovscale_bias_TNT` (6/10 sampled nvjet + launches) vs bf16 `nvjet_sm90_tst_256x128_...` — FP8 kernels confirmed + selected. Isolated GEMM ratios (locked 1980 MHz, pre-quantized inputs): + qkv ×1.45-1.53, ffn1 ×1.26-1.38, ffn2 ×1.54, lq ×1.43. +- Decision: keep-behind-flag (mandatory per roadmap — 2B-2 ships default-OFF + regardless of speed; enable decision belongs to Phase 4) +- Interpretation: the infrastructure works and is net-faster (+3.07% E2E, + −5.0 ms/chunk — inside the revised −3.5..−6 ms estimate once dynamic + quantization costs are counted; the roadmap's −13 ms ceiling assumed + pre-quantized activations). But the distilled one-step model is confirmed + fp8-sensitive: 40.7 dB full-scope with errors compounding across the 30 + blocks and persisting through the KV cache, so **no fp8 configuration is + currently enable-eligible**. Phase-4 redesign ticket (in priority order): + (1) blockwise/1x128 activation scales (post-GELU outliers dominate the + rowwise amax), (2) smoothquant-style weight/activation rebalancing, + (3) first/last-block bf16 exclusion, (4) fast-accum A/B. The per-site + sweep tool and the `ffn1` scope token exist precisely for that audit. + +#### 2026-07-08 14:35 · Phase 2B-3 · Fused mask-gen threshold-select → REVERT + +- Commit / patch: developed on top of 3b4b23d, reverted before landing; only + this log entry is committed. (Run logs: `profiling/runs/phase2b2/test_fused_maskgen*.log`.) +- Files changed (during the attempt, all removed on revert): + `diffsynth/models/fused_topk_mask.py` (Triton radix-select+compare kernel, + single-tile n<=16384 fast path + tiled variant), + `diffsynth/models/wan_video_dit.py` (branch in `generate_draft_block_mask`), + `examples/WanVSR/profiling/test_fused_maskgen_lossless.py` +- Flag: `FLASHVSR_FUSED_MASKGEN` (removed with the revert) +- Scope decision (pre-registered): the exact-mask gate forbids fusing + mean-pool/einsum/softmax (any reduction-order change flips ties across the + 66 concatenated softmax rows whose quotients share a threshold), so the + exact-safe fusion is ONLY kthvalue+broadcast+compare → one kernel. The + selected threshold is a pure order statistic (a VALUE of the multiset), so + any correct selection is bit-identical — no tie-breaking ambiguity. +- Correctness (the gate PASSED): kernel-level mask equality vs kthvalue AND + topk formulations on randn / heavy-tie / softmax-with--inf / all-equal / + negative inputs at (12|36)x13068, 12x172800, edge k∈{1,2,n/3,n−1,n} — all + exact; E2E max|diff| == 0 vs eager, 2 ON repeats identical. +- Performance (the reason for the revert): isolated @12x13068, + k_smallest=5149: eager kthvalue+compare 59.7 µs vs fused 69.2 µs (x0.86; + best of num_warps sweep {4,8,16,32}: 171.6/78.5/69.0/79.7 µs; first tiled + version was 113 µs). Harness E2E @768 F=89: OFF 42.00 FPS → ON 41.84/41.62 + FPS (−0.4..−0.9%). Below the pre-registered keep gate (≥ +0.5%). +- Peak mem before → after: 15.62 → 15.62 GiB (unchanged) +- Nsight report path: n/a (kernel-level + harness E2E evidence sufficient + for a negative result) +- Decision: **revert** — code removed cleanly, no flag left behind +- Interpretation: the roadmap's −5 ms/chunk estimate assumed fusing the + whole pool→einsum→softmax→topk chain to ~1 ms, but the exact-mask gate + makes everything upstream of the select bitwise-untouchable, and the + remaining exact-safe slice (select+compare, ~60 µs/call eager) has no + headroom: torch's `kthvalue` is already a single efficient radix-select + kernel (the 2A-4 lean pass already banked the topk→kthvalue win), and at + rows=12 a one-CTA-per-row fused kernel is latency-bound by 4 dependent + histogram rounds (~69 µs floor ≈ eager). The launch-gap latency the fusion + removes was already hidden by neighbouring kernels (2A-4 interpretation). + Conclusion: mask_gen (~4.4 ms/chunk) is now attention-kernel and + numerics-gate bound — further gains require either folding mask + generation into the Phase-3 attention kernel v2 or accepting non-bitwise + mask changes under the Phase-4 E2E-neutral protocol. 2B-4 (fused + im2col-GEMM) remains the only open 2B item. + +#### 2026-07-08 · Phase 3 · Step 2 — fresh 2A-stack baseline + kernel isolation (Phase-3 starting point) + +Re-measured at the Phase-3 working tree (parent 25811e3), full 2A recommended +stack, clocks locked 1980 MHz. Command = §0.4 primary + 2A kept set. + +| Field | Value | +|---|---| +| Run 1 / 2 / 3 FPS (pre-change tree) | 42.115 / 42.048 / 41.830 | +| **Median FPS** | **42.048** (matches 2B-1 baseline 41.759 within noise) | +| Median steady chunk ms | 136.88 · peak 15.62 GiB | +| ncu `_bsfa_tma_kernel_snd` (steady, in-pipe) | 1997.4 µs · SM/tensor SOL 40.7% · occ 12.5% = 1 CTA/SM (178 reg + 229.4 KB smem) · grid 792 = 6.0 waves · L2 92% · stalls barrier 2.36 / wait 1.01 / short_sb 0.55 · no-eligible 68.6% | +| `bench_ceilings.py` H2 re-run | cuDNN dense 1.945 ms → ideal sparse 1.179 ms (ANALYSIS ref: 1.86 / 1.13) | +| Shape correction (measured, was undocumented) | steady attention consumes the UNTRIMMED kv window: q 8448 × **kv 33792** (264 blocks) at mask density **0.42–0.45** — FLOP-identical to ANALYSIS's "kv 25344 @ 0.606" framing (~120 active kv blocks/row either way). True ideal-sparse at the real shape/density: dense(33792) 2.576 ms × 0.4213 = **1.085 ms**. | +| Logs | `profiling/runs/phase3/step2_baseline_run{1..3}.log`, `profiling/reports/ncu/phase3_bsfa_before.{ncu-rep,csv}` | + +#### 2026-07-08 16:30 · Phase 3 · Warp-specialized block-sparse attention v2 (`triton2`) + +- Commit / patch: phase3 attention v2 (this commit; parent 25811e3) +- Files changed: new `diffsynth/models/triton_block_sparse_attn_v2.py` (Gluon + kernel + wrapper), `diffsynth/models/wan_video_dit.py` (triton2 branch in + `flash_attention`, knob doc), new `examples/WanVSR/test_attention_v2.py`, + new `examples/WanVSR/profiling/bench_attn_v2.py` (+ captured real mask + `profiling/cache/attn_mask_768_steady.pt`) +- Flag: `FLASHVSR_ATTN_BACKEND=triton2` (default backend remains `sparse`; + the recommended-set line keeps `triton` until this entry is independently + confirmed). Tuning knobs: `FLASHVSR_ATTN_V2_NBUF` (default 3, measured + best), `FLASHVSR_ATTN_V2_PROD_REGS` (default 40, measured best). +- Design (kernel): FA3-style Gluon warp specialization on sm_90 — + 1 TMA producer warp group (4 warps @ 40 regs via `worker_num_regs` + setmaxnreg reallocation) streams the CSR-selected K/V 128×128 tiles into an + NBUF=3-deep smem ring (mbarrier full/empty handshakes, 2×64-col TMA boxes + per tile, next-index software prefetch); 2 consumer warpgroups (4 warps + each) own 64-row halves of the q block ("pingpong": one WG's softmax + overlaps the other's WGMMA) and additionally defer each P·V wgmma by one + iteration (issue QK(j) before PV(j−1); in-order retirement lets + `wait(pendings=1)` complete QK while PV drains under the softmax). + Exact-mask preservation is structural: same `(H,Nqb,Nkvb)` boolean mask, + same `_make_csr` (stable argsort → ascending kv-block order), all-false + rows → l=0 → l_safe → zero rows. Softmax scale folded into the exponent + (matches block_sparse_attn's FA2 formulation; v1 pre-scaled q in bf16). +- Fallback ladder: triton2 → v1 strided (`_bsfa_tma_kernel_snd`) → v1 + contiguous → `block_sparse_attn`; any v2 import/compile/launch failure + falls through silently. Default `sparse`/`triton` code paths untouched. +- Route notes (prototyped and rejected): + (a) `tl.range(warp_specialize=True)` one-liner on the v1 kernel — the + hopper autoWS pass accepted the annotation but did NOT partition + (num_warps stayed 8; 1.751 ms = no-op). (b) occupancy tuning for + 2 CTA/SM (M64/N64/w4 variants) — all slower (1.83–3.32 ms); WGMMA at + M64 tiles + no WS loses more than residency gains. (c) BLOCK_N=64 + 2-CTA variant of v2 — needs two wgmma layouts with per-iteration + conversions; rejected (compile complexity, expected loss). +- Isolated kernel (bench_attn_v2.py, real captured mask d=0.4213, locked + clocks): v1 kernel-only 1.753 ms → v2 **1.139 ms** (×1.54, 649 TF/s + sparse-effective, 95% of the 1.085 ms ideal-sparse ceiling); NBUF/PROD_REGS + sweep: {2,3}×{24,40,56} → 3/40 best. +- ncu acceptance gate (in-pipe steady chunk, `phase3_bsfa_v2_after.ncu-rep`): + +| Metric | Gate | Before (`_bsfa_tma_kernel_snd`) | After (`_bsfa_v2_kernel`) | Verdict | +|---|---|---|---|---| +| duration @ reference shape | ≤ 1.3 ms | 1997.4 µs | **1260.5 µs** | PASS | +| tensor pipe (SM SOL) | ≥ 60% elapsed | 40.7% | **66.2%** | PASS | +| barrier stall /issue | < 1.0 | 2.36 | **0.71** | PASS | +| residency | ≥2 CTA/SM or WS ≥16 warps/SM | 1 CTA/SM · 8 warps | 1 CTA/SM · **12 warps WS** (4+4 consumers + 4 producer) | PARTIAL (see interpretation) | +| (info) no-eligible-warp cycles | — | 68.6% | 57.7% | — | +| (info) stalls long_sb / wait | — | 0.55 / 1.01 | 2.00 / 1.32 | — | + +- E2E @768x1408 F=81 (untraced, 3-run medians, same tree, back-to-back): + OFF (`triton`) 41.692 FPS / 137.96 ms / 15.62 GiB → ON (`triton2`) + **45.466 FPS / 122.06 ms / 15.62 GiB** = **+9.05% FPS, −15.90 ms/chunk, + peak unchanged**. Logs `profiling/runs/phase3/e2e_{off,on}_run{1..3}.log`. +- Quality/correctness: kernel cos vs `block_sparse_attn` ≥ 0.999995 at the + real shape across densities {real 0.42, 0.1, 0.3, 0.45, 0.9, all-true} + + degenerate all-false masks/rows/heads (all-false rows exactly zero) + + determinism (2 calls bit-identical); E2E `test_attention_v2.py`: + PSNR(sparse, triton2) = **50.03 dB** (gate ≥49, v1 reference ~49.7), + max|d| = 0.3135, triton2 ×2 repeats bit-identical (multi-chunk KV/cache + contract stable); `test_phase2a_lossless.py ALL` → max|diff| = 0 + (default paths byte-identical). +- nsys attribution (`profiling/reports/phase3_attn_v2/`): `_bsfa_v2_kernel` + n=240 (8 chunks × 30 blocks) fully replaces `_bsfa_tma_kernel_snd` (0 + occurrences); steady chunks 137.96 → 122.06 ms untraced. +- @1536x2560 spot-check (single runs): OFF 11.472 FPS / 503.22 ms / + 48.39 GiB → ON **12.436 FPS / 449.24 ms / 48.39 GiB** (+8.4%, + −53.98 ms/chunk) — the win holds/grows at scale as predicted + (attention share is scale-invariant). +- Nsight report paths: + `profiling/reports/ncu/phase3_bsfa_before.{ncu-rep,csv}`, + `profiling/reports/ncu/phase3_bsfa_v2_after.{ncu-rep,csv}`, + `profiling/reports/phase3_attn_v2/` (nsys). +- Decision: **keep-behind-flag** (pre-registered tier: the ncu gate is 3/4 — + the residency letter asks ≥2 CTA/SM or ≥16 warps/SM and v2 runs 12 + warps/SM WS at 1 CTA/SM). Presumptive recommended-set promotion + (`FLASHVSR_ATTN_BACKEND=triton2`) after one independent confirming entry, + per the two-entry promotion rule. +- Interpretation: the kernel hit the acceptance window on every *binding* + metric — 1.26 ms in-pipe (target ≤1.3; 1.14 ms isolated = 95% of the + ideal-sparse ceiling), tensor-active 66.2% (target ≥60), barrier stalls + 2.36 → 0.71 — confirming the Phase-1 diagnosis that the v1 kernel was + scheduling-bound, not math-bound. The residency sub-criterion was + deliberately not chased: every ≥16-warp/2-CTA configuration we could + construct measured slower (occ variants 1.83–3.32 ms; N=64 2-CTA needs + per-iteration layout conversions), i.e. 12-warp warp-specialization with + softmax/WGMMA pingpong is the empirically optimal structure here, and the + gate's intent (kill the no-eligible-warp starvation) is what the passing + metrics measure. E2E banks −15.9 ms/chunk vs the −22.1 ms one would get by + naively scaling the single-shape ncu delta (737 µs × 30): that −22.1 was an + extrapolation, not a target — the 30 attention calls/chunk span DiT blocks + at different densities and warm-cache states, so the measured −15.9 ms is + the truth and the gap is extrapolation slack, NOT power throttling. + Power/clock was directly measured (nvidia-smi dmon, OFF vs ON busy window): + the GPU is a 900 W-capable Hopper enforced to 700 W (the 1000 W Grace+Hopper + module budget minus ~300 W reserved for Grace/LPDDR5X), but BOTH backends + peak at only ~660 W (v1 666 W / v2 660 W) — never touching the 700 W cap — + and BOTH sag identically to ~1650–1800 MHz under the 1980 MHz lock (a + heavy-tensor DVFS operating point, not a power or thermal cap: temps 52–56 + °C, power under limit). Since the sag is backend-independent it cannot be + the OFF-vs-ON differential; v2 does the SAME work at the SAME ~660 W in less + time = pure efficiency. (An earlier draft of this entry wrongly attributed + the gap to an H8 power-cap clawback; corrected here against the dmon data.) + @768 lands at 45.5 FPS (+9.1%); the roadmap's 49–51 was itself built on the + same optimistic 0.9 ms/call extrapolation. Follow-ups: (1) re-measure `FLASHVSR_DECODER_OVERLAP` + co-execution on top of triton2 (v2 still occupies 229 KB smem/SM, so the + 34% co-execution ceiling probably persists — measure, don't assume); + (2) **CORRECTION (PC-sampling, 335k samples, source page):** the aggregate + `long_sb=2.00` label is a warp-state classification artifact of the producer + warps parked on mbarriers — it is NOT the real limiter. Per-PC sampling + attributes warp-time as ~59% softmax f32/MUFU chain (MUFU.EX2 18.9% + + FMNMX 14.9% + FADD 12.0% + FMUL 9.7% + F2FP 1.9% + SHFL 2.0%) and ~24% + wgmma (HGMMA issue 13.6% + WARPGROUP.DEPBAR 10.0%); producer/TMA/LDG/ + mbarrier PCs are <1%. The kernel is **softmax-math + wgmma-wait bound, NOT + TMA/ring-head bound** — deeper KV prefetch (NBUF) buys nothing; the levers + are FFMA-form exp2 (fold `qk*s - m*s` into one FFMA), an alpha==1 bit-exact + rescale skip, and explicit pingpong ordering so softmax(WG-A) overlaps + HGMMA(WG-B). Isolated 1.139 ms = 95% of the 1.085 ms ideal-sparse ceiling; + perfect softmax/wgmma overlap floor ≈ 0.95–1.0 ms in-pipe. (3) FP8 attention (Phase 4) now + has a WS substrate to build on. + +#### 2026-07-08 · Phase 3.5 · Exact-math efficiency pack (kernel FFMA + fused CSR + overlap/rope remeasures) + +Base = commit 1b4d5c6 (triton2 + 2A stack). Clocks locked 1980 MHz. Same-session +OFF (`triton`) reference (3-run median): **41.531 FPS / 138.94 ms / 15.62 GiB**. +All runs `profiling/runs/phase3_1/`. + +- **3.5-0 log attribution fix** (commit): PC-sampling (335k) showed the kernel is + softmax-math (~59%) + wgmma-wait (~24%) bound, NOT TMA/ring-head; the aggregate + `long_sb=2.00` was a producer-park artifact. No perf change. + +- **3.5-1a FFMA-form softmax** (commit, in `triton2`): track the running max in + the pre-scaled (log2e) domain (`ms = m*s`) so the exponent becomes one FFMA + `qk*s - ms` instead of the non-contractable `(qk-m)*s`. Isolated kernel + 1.155 → **1.047–1.087 ms** (−6% median, −9.4% best), in-pipe ncu **1260 → 1179 µs, + tensor SOL 66.2 → 71.1%** (barrier 0.71 → 0.87, still < 1.0). Correctness: kernel + cos 0.999996 across all densities + degenerate; E2E PSNR(sparse, triton2) + **50.08 dB**; triton2 repeat bit-identical. E2E 45.466 → **46.112 FPS** (+1.42%), + 122.06 → 119.92 ms. ncu report `reports/ncu/phase3_1_bsfa_v2.ncu-rep`. +- **3.5-1b alpha==1 rescale skip** (dropped): IEEE-exact (x*1.0≡x) but the + per-iteration full-reduce + branch cost **more** than the acc-rescale FMULs it + skips (those overlap the tensor pipe): isolated 1.047 → 1.126 ms. Reverted. +- **3.5-1c explicit pingpong named-barrier** (dropped): primitives exist + (`inline_asm_elementwise`, `ttgl.barrier`) but the two consumer partitions are + separate warp_specialize regions — coordinating a hardware named barrier across + them is high deadlock-risk for 14 dB; a real kernel (e4m3 P, e4m3 cache) would be + worse. Probe reverted (no code left). Stop-condition honored (PSNR <45 every + variant → revert, per roadmap 4C). +- Consequence for the roadmap trajectory: the FP8 lever (both GEMM and + attention) is closed at the ≥49 gate for this distilled model. The remaining + denoise-side headroom is structural/lossless (5A) or the deferred rope + interleave; the E2E ceiling is now decode-tail bound (%32 wall). + +#### 2026-07-08 · Phase 5A · Zero-copy windowing — FEASIBILITY CONFIRMED, deferred + +- Rank-4 gluon TMA descriptor `[2,8,8,128]` over raster `(F,H,W,Hh·D)` at + `[f0,h0,w0,head·D]` reproduces `WindowPartition3D.partition`'s token order + **bit-exactly** (match=True, max|diff|=0) into a 2D `[128,D]` NVMMA smem tile + → WGMMA-ready. So win_part/win_rev (5.9 ms/chunk of pure permutation copies) + can be removed with no math change. +- Deferred to a focused session: the integration (KV-arena raster ring + + mask_gen raster pooling + 4D output store + multi-chunk arena bit-stability) + touches the KV-cache bit-stability contract and must not be rushed. Full spec + in `.opencode/plans/hw-efficiency-roadmap.md` §5A. Expected −4..−6 ms/chunk, + lossless (max|diff|=0 gate). Flag `FLASHVSR_ATTN_ZEROCOPY`. + +#### 2026-07-08 · Phase 3.5–5 campaign closure (this session) + +- **Landed (kept):** 3.5-1a FFMA-form softmax + 3.5-2 fused CSR (both in/with + `triton2`). @768 OFF(triton) 41.53 → **46.38 FPS (+11.68%)**, 138.9 → + 118.8 ms; @1536 11.46 → **12.67 FPS (+10.5%)**. Quality: kernel cos 0.999996, + E2E PSNR 50.08 dB, CSR bit-exact. Peak unchanged 15.62 GiB. +- **Measured/kept-flag:** 3.5-3 decoder-overlap ×v2 +1.85% (co-exec 31.7%, the + v2 smem monopoly persists — a real Phase-5 dependency); 3.5-4 rope-freqs + cache +0.1% (noise); 4B FP8-GEMM blockwise infra (default-OFF, 41.0 dB). +- **Dropped with evidence (saved dead-end kernel work):** 3.5-1b alpha-skip + (−7% isolated), 3.5-1c pingpong barrier (deadlock risk, already 83% peak), + 4A rope-fp32 (×1.01 — interleave-bound not fp64-bound), 4C FP8-attention + (probe 34.96 dB, gate unreachable — attention is far more e4m3-sensitive than + GEMMs; no kernel written). +- **Key finding:** the FP8 lever (GEMM + attention) is closed at the locked + ≥49 dB gate for this distilled one-step model — its e4m3 error compounds + through the streaming KV cache (4B) and softmax amplifies K errors (4C). + Remaining denoise headroom is structural/lossless (5A zero-copy windowing, + feasibility-confirmed; deferred rope-interleave); the E2E ceiling is now + decode-tail bound (~32% of wall). +- **Recommended-set delta (pending 2nd confirming entry):** + `FLASHVSR_ATTN_BACKEND=triton2` + `FLASHVSR_FUSED_CSR=1`. + +#### 2026-07-08 · Phase 5 (P1–P4) · lossless structural pack — session log + +Base 46.383 FPS / 118.84 ms (triton2+FUSED_CSR). All items bit-exact +(max|diff|=0 gates incl. multi-chunk ring rotation + ON-repeat stability). +Runs: `profiling/runs/phase5/`. + +- **P1 coalesced RoPE kernel** (`diffsynth/models/triton_rope.py`): the fused + torch.compile RoPE read `freqs.real/.imag` as stride-2 fp64 views (~484 GB/s, + 143 µs/call). Hand Triton kernel moves bf16 pairs as packed u32 words and + reads the complex128 storage directly: 128→36 µs (×3.6; ×3.9 at the chunk0 + shape). Bit-exactness required matching torch's fp64→fp32→bf16 double-round + cast chain (c10::BFloat16 is float-constructed); with it, triton==fused== + eager bitwise at both shapes. E2E +3.01%, −5.1 ms/chunk (isolated predicted + −5.5 ✓). +- **P2 incremental pooled-K**: probe first — torch.mean(dim=1) is bitwise + independent of batch row count ((264,128,C) == 4×(66,128,C) slices) → GO. + Pool each kv window once (when appended), carry pooled rows in a rolling + buffer mirroring the cat/trim sequence; self-heals on any length mismatch. + +0.16%, −0.6 ms (part of the 1.6 ms reduce was already latency-hidden). +- **P3 zero-copy V (5A-v1)**: V never feeds mask pooling, so V windowing is + pure layout work. `_VRasterArena` keeps V as raster (F,h,w,C) frames + (contiguous appends, no permute); the v2 kernel gains RASTER_V (rank-4 TMA + box [2,8,8,64] window gather, probe bit-exact incl. ring f0 offsets) and + RASTER_OUT (epilogue stores rows straight to raster token order → win_rev + eliminated). K/Q stay windowed (their pooled means feed the exact mask; + raster pooling would change reduction order — 2B-3 lesson). Fallback: on + any zc failure the raster live view is partitioned on the fly and the + standard ladder runs. +0.86%, −0.84 ms (of the −2.7 estimate: the removed + copies partially overlapped other work, and the scattered raster store is + slightly slower than the contiguous window store). +- **P4a decoder overlap re-eval**: on the shorter denoise chunk the overlap + win grew to **+2.01%** (49.237 FPS, tail 543→124 ms, peak −0.46 GiB) — + meets the pre-registered ≥2% promote gate. Promoted to the production + recommended set; benchmarking note stands (with the flag ON the `[chunks]` + metric measures denoise+decode, so per-chunk attribution runs keep it OFF). +- **P4b cudnn.benchmark**: −0.1% and output bit-identical → cuDNN's heuristic + algo choices were already optimal for the decoder's static shapes. Dropped. + +**Session result @768:** 46.383 → **48.268** FPS denoise-set (+4.1%), → +**49.237** with overlap (+6.2%); steady 118.84 → **112.24 ms**. Peak +15.62→15.64 GiB (+0.02 pooled-K buffers; overlap variant 15.18). +**@1536:** 12.669 → **13.194** (+4.1%), → **13.435** with overlap. +**Cumulative vs Step-0 (38.585 FPS):** +25.1% recommended / **+27.6%** with +overlap; @1536 vs campaign start (11.01): **+22.0%** with overlap. +**Recommended set (updated):** 2A flags + `FLASHVSR_ATTN_BACKEND=triton2` + +`FLASHVSR_FUSED_CSR=1` + `FLASHVSR_ROPE_KERNEL=triton` + +`FLASHVSR_POOLED_K_CACHE=1` + `FLASHVSR_ATTN_ZEROCOPY=1` +(+ `FLASHVSR_DECODER_OVERLAP=1` for production throughput). +Interpretation: the denoise side is now deep in diminishing-returns territory +(attention at 95% of ideal-sparse, GEMMs at 82–85% SOL, FP8 closed by quality, +rope at BW): P1 was the last mid-size lossless brick and it delivered exactly +its isolated prediction; P2/P3 confirmed that the remaining copy/pool costs +were already partially hidden. The E2E ceiling is decode-bound: the decode +tail is untouched ~540 ms serialized (or ~34% co-executed under overlap) — +further FPS needs decoder-side work (Phase 4D compile-fusion under the PSNR +gate) or accepting the ~49–50 FPS plateau @768. + +### Phase 6 — Lossless decoder and LQ kernels (2026-07-09) + +Fresh Phase-5 production baseline (three-run median, 768x1408 F=81) was +**49.229 FPS**. Every candidate was toggled independently on the same stack, +then the cumulative stack was compared bit-for-bit at F=29 and F=89. + +- Pointer-state rotation removes 231 recurrent D2D copies (10.54 GiB traffic) + without changing decoder state order. +- Direct output placement removes the serialized final chunk concatenation and + skips materializing the three causal warm-up RGB frames. +- MemBlock Triton pointwise kernels preserve the eager BF16 materialization + after bias and residual addition; special-value and E2E tests are bit-exact. +- The nearest-neighbor kernel reads each channels-last input pixel once and + writes four outputs; isolated speedup is 7.7–8.6x on production shapes. +- The LQ packer changes only `unfold/permute/reshape` materialization. The + existing `torch.addmm`, output layout, normalization, and arithmetic remain + unchanged; isolated packing speedup is 1.85–2.30x. +- The recurrent concat kernel packs `[current, past]` in one channels-last + pass and is 1.7–1.8x faster in isolation. + +**Cumulative result:** 49.229 → **53.423 FPS** (+8.52%) at 768x1408; peak +allocated memory 15.18 → 15.44 GiB. One phase-closure run at 1536x2560 reached +**14.684 FPS** versus the previous 13.435 (+9.30%), with 47.76 GiB allocated. +The post-review strict rerun reached 53.484 FPS at 768x1408 with every expected +route count satisfied and no recorded fallback. +`test_phase5_lossless.py phase6` reports `max|diff|=0` at 768x1408 F=29/F=89 +and 1536x2560 F=29/F=41. The `F=29` case confirms direct output matches the +existing 21-frame cat behavior for accepted `8n+5` inputs. The final profile +reduced decoder kernel work 518.8 → 393.7 ms +and decoder launches 3860 → 2969. + +Rejected experiments were removed from production code: TGrow packing was +bit-exact but reduced FPS and added ~0.5 GiB; direct causal LQ packing was +bit-exact but 0.37% slower; standalone Triton ReLU was slower, while cuDNN's +fused Conv+ReLU changed BF16 output values. + +### Phase 6b — Decoder conv-body campaign (2026-07-09) + +Quality budget for this phase (user-approved): **>= 49 dB E2E PSNR** against +the current stack; bitwise preferred where free. Fresh NCU on the four cuDNN +decoder conv families showed 75–93% SM/tensor SOL for the MemBlock/transition +convs (no custom-GEMM headroom) and 42% SOL only for the narrow-N final RGB +conv. The analyzer now folds `decode0..N` NVTX ranges into a `decode` phase +and defaults to the true steady window `chunks 3..7`. + +- **6b-P1 `FLASHVSR_TCDECODER_TGROW_UP` (keep, lossless):** the three + `Upsample(2x) -> TGrow(1x1)` pairs are algebraically reordered: the bias-free + 1x1 conv runs at LOW resolution (4x fewer FLOPs), then one Triton kernel + unpacks the temporal channel groups and nearest-upsamples straight into + contiguous channels-last frames (removing the high-res TGrow packing + copies). Isolated 2.2x/5.2x/5.4x on the three sites; measured bit-identical + E2E at 768x1408 (F=25/29/89) and 1536x2560 (F=25/29/41) — cuDNN picked the + same reduction order at both spatial sizes. Production strict 3-run medians: + 53.40 → **54.61 FPS** (+2.27%). +- **6b-P2 `FLASHVSR_TCDECODER_CUDNN_FUSED` (keep, quality-gated):** MemBlock + chains run as cuDNN runtime-fused `conv+bias+ReLU`, `conv+bias+ReLU`, + `conv+bias+residual+ReLU`, and the standalone Conv->ReLU pairs (latent conv, + IdentityConv deepening, full-res tail) fuse via + `aten.cudnn_convolution_relu`. Replaces the separate exact epilogue kernels + (~31 ms/call) with zero extra launches. NOT bit-exact: the fused engine + skips the intermediate BF16 materializations (isolated ~70 dB, ~10–14% of + values ±1 ULP; E2E **55.4–55.8 dB** vs the lossless stack at both + resolutions). Production strict 3-run medians: 54.58 → **55.91 FPS** + (+2.44%). +- **6b-P3 Triton RGB tail conv (drop):** `tl.dot` N=16-padded kernel reached + 65 dB but 0.351 ms vs cuDNN 0.335 ms on (1,128,1408,768): the formulation + re-reads activations 9x through L2 (~2.5 GB/frame) and is L2-bandwidth + bound; beating the 42%-SOL sm80 kernel needs an smem-halo design (Gluon), + which fails the >=15% isolated bar for this phase. + +**Cumulative:** 53.423 → **55.91 FPS** (+4.66%) at 768x1408 F=81 (3-run +medians). 1536x2560 F=41 single-run spot-checks, all three configs measured +back-to-back in one same-process-free batch (`runs/phase6b_1536_isolated/`) +on the same Phase-6 base: baseline **14.36** → +TGROW_UP **14.69** +(+2.30%, 46.0 → 46.4 GiB) → +both **15.07 FPS** (+4.94% vs this fresh 1536 +baseline, 46.4 GiB unchanged). Note this fresh 14.36 FPS/46.0 GiB baseline +run differs slightly from the 14.684 FPS / 47.76 GiB number logged at the +original Phase-6 closure (different session/allocator state); all Phase-6b +deltas above are same-batch, same-session comparisons and are the numbers to +trust for this phase. Decoder-stream kernel work 392.0 → **299.9 ms** and +decoder launches 2969 → **2089** (nsys `phase6b_fused_current` @768, both +flags on). With `TCDECODER_CUDNN_FUSED=0` the stack remains strictly +`max|diff|=0` at **54.61 FPS** @768 / **14.69 FPS** @1536. Campaign total +@768 vs Step 0: 38.585 → 55.91 = **+44.9%** (lossless-only, i.e. +`CUDNN_FUSED=0`: 38.585 → 54.61 = **+41.5%**). + +Next candidates (measured, not yet implemented): DiT residual/LayerNorm/AdaLN +row fusion (~15.3 ms/chunk budget on the denoise stream, expected +2–4%), +MemBlock concat elimination via split-K dual `cudnn_convolution_add_relu` +(~20.6 ms/call concat kernel), true circular K/V arenas (+0.3–0.8%). + +### Phase 7 — 60 FPS push (2026-07-09) + +Target: 60 FPS at 768x1408 F=81, three-run untraced median, with the +user-approved >=49 dB E2E PSNR gate. Fresh Phase-6b baseline was **56.085 FPS** +(55.920 / 56.085 / 56.098); an `nvidia-smi -lgc 1980` A/B was neutral-to-negative +(56.089 / 56.028 / 55.854, median 56.028), so the default clock policy remains. + +- **P7-A `FLASHVSR_DIT_ROW_FUSION` (keep, quality-gated):** a Triton row kernel + fuses affine-free LayerNorm -> AdaLN for `norm1`/`norm2`; a second kernel + performs the broadcast residual gate. Isolated `8448x1536`: LN+AdaLN + 0.242 -> 0.026 ms, gate 0.040 -> 0.024 ms. E2E F=25/F=89: 49.78 / 49.88 dB; + strict 3-run F=81: 56.085 -> **56.398 FPS** (+0.56%). The large isolated gain + is largely hidden by decoder co-execution and warm chunks. +- **P7-B `FLASHVSR_MASKGEN_THRESHOLD_CACHE` (keep, quality-gated):** the first + two per-block geometries retain exact `kthvalue`; later steady chunks reuse + the previous threshold. F=89 used 60 exact and 210 cached threshold routes, + with 51.72 dB against the prior mask path. It composes with P7-A at 49.60 dB. +- **P7-C `FLASHVSR_TCDECODER_SPLITK_CONV` (keep, quality-gated):** split the + first MemBlock 3x3 weights into current/past halves, pre-cache channels-last + views, and use `cudnn_convolution_add_relu`; no recurrent concat is + materialized. Isolated three decoder shapes gained 3–10% at ~68.9 dB; F=89 + E2E was 58.13 dB. A+B+C F=29/F=89 was 49.81 / **49.59 dB**. +- **Cumulative A+B+C:** **57.256 FPS** (57.256 / 57.213 / 57.358) versus the + fresh 56.085 baseline, **+2.09%**. Steady chunks were ~135 ms; peak 15.59 GiB. + All expected routes passed under `FLASHVSR_REQUIRE_FASTPATHS=1` with no + fallback. +- **P7-D `FLASHVSR_CONV3D_EINSUM` (drop):** direct strided contraction made + isolated conv2 8.39 -> 7.91 ms and passed combined quality at 49.56 dB, but + the F=81 A+B+C+D median regressed to **56.741 FPS**. Decoder/DiT co-execution + erased the local saving. +- **P7-E `KV_RINGBUF_SPARE=8` (drop):** removed all 60 F=81 arena compactions + and stayed bit-exact at F=89, but its one-run 57.518 FPS probe did not repeat: + A+B+C+E three-run median was **57.200 FPS** with peak memory 24.3 GiB + (+8.7 GiB). Default spare remains 2. + +**Decision:** promote A+B+C to the GH200 preset. The 60 FPS target was not +reached under the >=49 dB gate: 57.26 FPS leaves 4.8% E2E headroom. Remaining +material levers are Track-2 attention softmax/WGMMA overlap (high integration +risk; prior named-barrier attempt was unstable) or an explicitly approved +quality/algorithm trade in mask density. + +### Entry template (copy-paste per attempt) + +```markdown +#### · · + +- Commit / patch: +- Files changed: +- Flag: FLASHVSR_<...> (default OFF) +- Env vars used (full set): +- Exact benchmark command: +- Resolution / frames: 768x1408 / F=81 (spot-check: 1536x2560 y/n) +- Warmup / steady settings: warmup=1, steady=chunks 2..6 (from [chunks] line) +- FPS before → after (Δ): (3-run medians) +- Steady chunk before → after (Δ): +- Peak mem before → after: +- Correctness: (max|diff| == 0 | PSNR = XX.X dB | mask-equality | n/a) +- Output difference (if applicable): +- Nsight report path (if generated): +- Decision: keep-enabled | keep-behind-flag | revert | investigate | postpone +- Interpretation (2–3 sentences, mandatory): did it match the roadmap + estimate? why / why not? what follows? +``` + +--- + +## Phase 2A closure summary (2026-07-08) + +- All of 2A-1..2A-5 landed with flag + log entry + interpretation + + correctness result; 2A-6 evaluated and postponed via its go/no-go gate. +- Kept enabled (recommended set): `FUSE_ROPE`, `KV_RINGBUF`, + `ATTN_STRIDED_IO`, `MASKGEN_LEAN`, `LQPROJ_LEAN`. Kept behind flag: + `CACHE_ROPE_FREQS` (FPS-neutral @768, lossless, graph-capture prereq). + Rejected during development: LQPROJ sub-item (b) (non-contiguous GEMM + output → 49.9 dB, fails the lossless gate). +- Combined-stack parity: `test_phase2a_lossless.py ALL` → max|diff| == 0 vs + all-flags-OFF (no flag interactions). +- Result @768x1408: 38.585 → 41.580 FPS (+7.8%), steady chunk 156.28 → + 138.46 ms (−17.8 ms); roadmap ceiling for 2A was ~20–26 ms — the gap is the + traced-CPU rope_freqs share (rides under GPU work untraced) and the + deliberately-skipped deep fusions (2B-3 mask topk, 2B-4 im2col). +- @1536x2560 spot-check recorded below; peak-mem cost of the arena documented. +- Recommended next (Phase 2B): decoder overlap (2B-1) first — decode is still + a fully serialized 17–23% of E2E; then FP8 GEMM infra (2B-2, default-OFF), + fused mask top-k (2B-3), fused im2col (2B-4, still eligible since 2A-5 + recovered <2 ms). + +## Cumulative stack + + + +| Step | Enabled Optimizations | FPS | Steady Chunk Time | Peak Memory | Delta vs Phase-2 Baseline | Notes | +|------|----------------------|-----|-------------------|-------------|---------------------------|-------| +| 0 | full-knobs baseline (gemm+NHWC+fuse_norm+triton+TMA+caches) | 38.585 | 156.28 ms | 12.6 GiB | — | Step-0 fresh baseline (2026-07-08, df94d94) | +| 1 | baseline + FUSE_ROPE | 39.023 | 153.62 ms | 12.6 GiB | +1.14% FPS / −2.66 ms | 2A-1b, lossless | +| 2 | + KV_RINGBUF | 39.429 | 150.76 ms | 15.5 GiB | +2.19% FPS / −5.52 ms | 2A-2, lossless; +2.9 GiB arena slack | +| 3 | + ATTN_STRIDED_IO | 41.099 | 141.14 ms | 15.5 GiB | +6.52% FPS / −15.14 ms | 2A-3, lossless | +| 4 | + MASKGEN_LEAN | 41.477 | 139.00 ms | 15.5 GiB | +7.50% FPS / −17.28 ms | 2A-4, exact mask | +| 5 | + LQPROJ_LEAN | 41.580 | 138.46 ms | 15.62 GiB | +7.76% FPS / −17.82 ms | 2A-5, lossless | +| 6 | + DECODER_OVERLAP (kept behind flag, not in default set) | 42.373 | 190.58 ms (denoise+decode) | 15.16 GiB | +9.82% FPS vs Step 0 / +1.71% vs Step 5 | 2B-1, lossless; steady-chunk metric absorbs decode when ON — later per-chunk benchmarking stays flag-OFF; decode tail 561→125 ms | +| 7 | 2A set + ATTN_BACKEND=**triton2** (kept behind flag pending confirmation entry) | 45.466 | 122.06 ms | 15.62 GiB | +17.83% FPS vs Step 0 / +9.05% vs the 2A set | Phase 3 attention v2; PSNR 50.03 dB vs sparse (same class as `triton`'s ~49.7); composes with all 2A flags | +| 8 | + FFMA-softmax (3.5-1a, in triton2) + **FUSED_CSR** | 46.383 | 118.84 ms | 15.62 GiB | +20.2% FPS vs Step 0 / +11.68% vs the 2A set | Phase 3.5 exact-math pack; both lossless-class (PSNR 50.08 / bit-eq CSR); @1536 +10.5% | +| 9 | + ROPE_KERNEL=triton + POOLED_K_CACHE + ATTN_ZEROCOPY | 48.268 | 112.24 ms | 15.64 GiB | **+25.1%** FPS vs Step 0 | Phase 5 P1-P3, all bit-exact (max\|diff\|=0) | +| 10 | + DECODER_OVERLAP (production set) | **49.237** | 164.88 ms (denoise+decode) | 15.18 GiB | **+27.6%** FPS vs Step 0 | P4a: overlap re-promoted at +2.01% on the shorter chunk; keep OFF for per-chunk attribution runs | +| 11 | + Phase-6 pointer/direct-output/pointwise/upsample/LQ-packer/concat | **53.423** | overlap metric | 15.44 GiB | **+38.5%** FPS vs Step 0 / +8.52% vs fresh Phase-5 | All Phase-6 paths bit-identical vs Phase-5 production | +| 12 | + `TCDECODER_TGROW_UP` (Phase 6b, lossless) | **54.61** | overlap metric | 15.6 GiB | **+41.5%** FPS vs Step 0 | bit-identical incl. 1536 F=25/29/41 | +| 13 | + `TCDECODER_CUDNN_FUSED` (Phase 6b, quality-gated) | **55.91** | overlap metric | 15.6 GiB | **+44.9%** FPS vs Step 0 | E2E 55.4 dB PSNR vs Step 12 (gate 49 dB) | +| 14 | + P7 `DIT_ROW_FUSION` + `MASKGEN_THRESHOLD_CACHE` + `TCDECODER_SPLITK_CONV` | **57.256** | overlap metric | 15.59 GiB | **+48.4%** FPS vs Step 0 / +2.41% vs Step 13 | Combined F=29/F=89 PSNR 49.81/49.59 dB vs Step 13 (gate 49 dB) | + +--- + +## Phase closure spot-checks (@1536x2560) + +| Date | Phase closed | Enabled set | FPS @1536 | Steady chunk @1536 | Peak mem @1536 | Notes | +|------|--------------|-------------|-----------|--------------------|----------------|-------| +| 2026-07-08 | 2A | full-knobs + FUSE_ROPE + KV_RINGBUF + ATTN_STRIDED_IO + MASKGEN_LEAN + LQPROJ_LEAN | 11.489 | 501.92 ms | 48.39 GiB | Phase-1 ref: 11.01 / 531.5 ms / 37.4 GiB → +4.35% FPS, −29.6 ms; +11 GiB = arena spare slots at this res (`FLASHVSR_KV_RINGBUF_SPARE` trades it back). Gain smaller than @768 (+7.8%) as attention/decode share grows with res — consistent with ANALYSIS §0. | +| 2026-07-08 | 3 | 2A set + ATTN_BACKEND=triton2 (OFF ref same session: 11.472 / 503.22 ms / 48.39 GiB) | 12.436 | 449.24 ms | 48.39 GiB | Phase-3 attention v2 spot-check: +8.4% FPS, −53.98 ms/chunk, peak unchanged — the kernel win holds at scale (attention share scale-invariant, ANALYSIS §0). | +| 2026-07-08 | 3.5 | 2A set + triton2 + FUSED_CSR (OFF ref same session: 11.461 / 503.38 ms) | 12.669 | 438.20 ms | 48.39 GiB | Phase-3.5 spot-check: +10.5% FPS, −65.2 ms/chunk vs OFF — FFMA-softmax + fused CSR hold/grow at scale. | +| 2026-07-08 | 5 | + ROPE_KERNEL + POOLED_K_CACHE + ATTN_ZEROCOPY | 13.194 | 414.32 ms | 48.49 GiB | Phase-5 P1-P3 spot: +4.1% vs Phase-3.5 set. With +DECODER_OVERLAP: **13.435 FPS** / 46.82 GiB — campaign total @1536: 11.01 → 13.44 (**+22.0%**). | +| 2026-07-09 | 6 | Phase-5 production + all Phase-6 lossless paths | **14.684** | 527.7 ms median (decode overlap included) | 47.76 GiB | Single phase-closure run, strict fast paths; +9.30% vs Phase-5 production. F=41 parity max\|diff\|=0. | +| 2026-07-09 | 6b | Same-batch baseline → +`TGROW_UP` → +`CUDNN_FUSED` | 14.36 → 14.69 → **15.07** | — | 46.0 → 46.4 → 46.4 GiB | Three single strict runs, same session/batch; TGROW_UP bit-identical @1536 (+2.30%), CUDNN_FUSED 55.7–55.8 dB E2E PSNR (+2.59% more). Combined +4.94% vs this fresh baseline. | +| 2026-07-09 | 7 | Phase-6b + `DIT_ROW_FUSION` + `MASKGEN_THRESHOLD_CACHE` + `TCDECODER_SPLITK_CONV` | **15.415** | 641.1 ms (chunk 2 single spot) | 46.46 GiB | Strict F=41 spot, +2.3% vs the 15.07 Phase-6b reference. Combined F=29 quality gate: 49.95 dB vs the Phase-6b production stack. | diff --git a/examples/WanVSR/profiling/PHASE_ROADMAP.md b/examples/WanVSR/profiling/PHASE_ROADMAP.md new file mode 100644 index 00000000..a3f8a2d2 --- /dev/null +++ b/examples/WanVSR/profiling/PHASE_ROADMAP.md @@ -0,0 +1,501 @@ +# FlashVSR Phase-2+ Optimization Roadmap + +Derived from the Phase-1 profiling campaign in [`ANALYSIS.md`](./ANALYSIS.md) +(nsys + ncu + ceiling microbenches, GH200 / sm_90, 2026-07-08). + +This document converts the profiling findings into a staged engineering plan. +**No optimization in this document is implemented yet.** All gain figures are +*estimates derived from profiling attribution and microbenchmarks* — they are +ceilings, not promises, and every one of them must be re-validated by a fresh +benchmark entry in [`PHASE_BENCH_LOG.md`](./PHASE_BENCH_LOG.md) when the +optimization lands. + +--- + +## 0. Baseline and Rules + +### 0.1 Phase-2 baseline (from the Phase-1 campaign, untraced single runs) + +| Metric | Value | Source | +|---|---|---| +| E2E FPS @768x1408, F=81 | **38.55** | `run_pipe_target.py`, all knobs ON | +| Steady chunk time @768 | **156.24 ms** (8 new frames/chunk) | `[chunks]` line, chunks 2–6 avg | +| Peak GPU memory @768 | **12.6 GiB** | `torch.cuda.max_memory_allocated` | +| E2E FPS @1024x1920 / @1536x2560 | 21.77 / 11.01 | resolution shift runs | +| Steady chunk @1024 / @1536 | 274.1 ms / 531.5 ms | | +| GPU idle within steady chunk | 3.1% @768, ~0% @1024+ | not launch-bound | +| Decode share of E2E | 17–23% (serialized tail) | all resolutions | + +Before the first Phase-2A change lands, this baseline MUST be re-measured +fresh (3 runs, median) and recorded as **Step 0** in `PHASE_BENCH_LOG.md`. +All subsequent deltas are computed against that Step-0 entry, not against the +numbers above. + +### 0.2 Environment + +- NVIDIA GH200 480GB (sm_90, 132 SM, 96 GB HBM3), driver 595.58.03 +- CUDA 13.2 · torch 2.12.0a0 (NGC 26.05) · triton 3.7.0 · cuDNN 9.22 +- Python: `/root/FlashVSR/venv/bin/python` (system python lacks imageio) +- GPU clocks locked: `nvidia-smi -lgc 1980,1980` (release: `nvidia-smi -rgc`) +- **Power constraint:** platform enforces a 700 W cap (900 W requested and + denied). Under sustained load the GPU draws 649–687 W and DVFS sags clocks + to 1635–1815 MHz (84–92% of max). Consequence: *removing work* compounds + (saves power → higher sustained clocks); *adding raw math throughput* + partially self-defeats. This is why the roadmap is efficiency-first. + +### 0.3 Baseline environment variables (the "full knobs" config) + +```bash +FLASHVSR_CONV3D_BACKEND=gemm +FLASHVSR_TCDECODER_CHANNELS_LAST=1 +FLASHVSR_FUSE_NORM=1 +FLASHVSR_ATTN_BACKEND=triton # implies FLASHVSR_ATTN_TMA=1 (default) +FLASHVSR_CACHE_MOD=1 +FLASHVSR_CACHE_MASK_BIAS=1 +``` + +### 0.4 Benchmark command template (PRIMARY, untraced) + +Run from `examples/WanVSR/`: + +```bash +FLASHVSR_CONV3D_BACKEND=gemm \ +FLASHVSR_TCDECODER_CHANNELS_LAST=1 \ +FLASHVSR_FUSE_NORM=1 \ +FLASHVSR_ATTN_BACKEND=triton \ +FLASHVSR_CACHE_MOD=1 \ +FLASHVSR_CACHE_MASK_BIAS=1 \ +FLASHVSR_PROF_STEADY=off \ +/root/FlashVSR/venv/bin/python profiling/run_pipe_target.py +``` + +Resolution / length variants via `FLASHVSR_PROF_W`, `FLASHVSR_PROF_H` +(multiples of 128), `FLASHVSR_PROF_FRAMES` (8n+1; default 85 → F=81, +8 chunks). The script prints: + +- `[result] {... fps, wall_s, peak_gib, steady_chunk_ms ...}` — the numbers + that go into the bench log, +- `[chunks] per-chunk ms: [...]` — steady chunk = mean of chunks 2..6. + +### 0.5 Mandatory discipline (non-negotiable rules) + +1. **Untraced runs are the only source of truth for FPS / chunk time / peak + memory.** nsys/ncu runs distort wall time (stop-flush lands inside the + timed window; CPU-side tracing overhead inflates gaps). +2. **Nsight Systems / Nsight Compute are used exclusively for attribution** + (where did the time go, why is a kernel slow) — never for headline numbers. +3. **One optimization at a time.** Every optimization is benchmarked + independently (flag ON vs flag OFF on the same commit) before the next one + is started. +4. **Every optimization is behind a `FLASHVSR_*` environment flag, default + OFF**, or (if a flag is genuinely impossible) trivially revertible in a + single commit. This preserves the PR philosophy: default output is + bit-for-bit the current behaviour. +5. **No performance claim without a fresh `PHASE_BENCH_LOG.md` entry.** + Numbers quoted from memory, from ANALYSIS.md, or from a stale run do not + count. +6. **Do not proceed to the next optimization until the current one has a log + entry AND a short written interpretation** (2–3 sentences: did it match the + estimate, why/why not, decision). +7. Every change lands with a correctness check appropriate to its class: + *lossless* claims require `max|diff| == 0` (pattern: + `test_cache_lossless.py`); *numerically-neutral* claims require PSNR vs + flag-OFF ≥ 49 dB on the standard clip (pattern: `test_fuse_norm.py`). +8. @768x1408 is mandatory for every change. @1536x2560 spot-check is required + only at phase closure (bottleneck structure is scale-invariant per + ANALYSIS §0, so per-change high-res runs are wasted GPU time). + +--- + +## 1. Phase 2A — Low-effort / high-confidence cleanup + +Goal: bank the safe wins first. Everything here is (a) small surface area, +(b) exactly-preserving or trivially PSNR-gated, (c) independently flag-gated. +Estimated combined ceiling from ANALYSIS: **~20–26 ms/chunk** +(156 → ~130–136 ms steady, i.e. roughly +12–17% denoise throughput) — to be +verified item by item. + +**Explicitly excluded from 2A:** attention kernel rewrite (Phase 3), FP8 +anything (Phase 2B infra + Phase 4 gate), decoder overlap (Phase 2B), +any change to the sparse mask semantics (Phase 4). + +### 2A-1 RoPE frequency caching + RoPE apply cleanup + +| | | +|---|---| +| Profiling says | `rope` phase = 10.1 ms/chunk GPU (6.8%); additionally `rope_freqs` construction showed ~44 ms/chunk of CPU-side wall in traced runs and is step-invariant (depends only on `(cur_process_idx, f, h, w)`) | +| Why low effort | freqs caching is a dict keyed on `(idx, f, h, w)` around an existing pure function; apply-fusion reuses the existing `_maybe_compile` pattern from FUSE_NORM | +| Estimated gain | −8 ms/chunk GPU (estimate, unverified) + CPU-side headroom that currently rides under GPU work | +| Confidence | high | +| Risk | low (freqs cache is bit-identical; fused apply is elementwise-only) | +| Files/functions | `diffsynth/pipelines/flashvsr_tiny.py::model_fn_wan_video` (the `rope_freqs` block), `diffsynth/models/wan_video_dit.py::rope_apply` | +| Flags | `FLASHVSR_CACHE_ROPE_FREQS` (lossless), `FLASHVSR_FUSE_ROPE` (fusion) — two separate flags, two separate log entries | +| Benchmark | primary @768; log both flags individually and combined | +| Correctness | freqs cache: `max\|diff\|==0` vs OFF. Fused apply: PSNR ≥ 49 dB vs OFF | +| Default | OFF until log entry; lossless part is a candidate for default-ON later | + +### 2A-2 KV cache ring buffer (remove `kv_cat`) + +| | | +|---|---| +| Profiling says | `kv_cat` = 3.4 ms/chunk (2.3%), `CatArrayBatchedCopy` at 3235 GB/s = 81% HBM3 — at BW limit, so the only fix is not copying; sliding-window trim already drops the oldest window each chunk | +| Why low effort | replace `torch.cat([pre_cache_k, k_w])` + trim-slice with a preallocated `(kv_len+1)`-slot buffer and rolling writes; contained in one forward function + cache init | +| Estimated gain | −3 ms/chunk (estimate, unverified) | +| Confidence | high | +| Risk | low-medium (indexing/rotation bugs would corrupt temporal context — caught by bit-diff) | +| Files/functions | `diffsynth/models/wan_video_dit.py::SelfAttention.forward` (kv concat + `cache_trim`), `diffsynth/pipelines/flashvsr_tiny.py::__call__` (`pre_cache_k/v` lifecycle) | +| Flag | `FLASHVSR_KV_RINGBUF` | +| Benchmark | primary @768 | +| Correctness | `max\|diff\|==0` vs OFF over a full multi-chunk clip (exercises rotation) | +| Default | OFF until proven, then candidate for default-ON (lossless) | + +### 2A-3 Attention-path transpose / contiguous cleanup + +| | | +|---|---| +| Profiling says | `attn_core` − kernel = 11.6 ms/chunk (7.7% of GPU): three `.contiguous()` transposes `(S,n,d)→(n,S,d)` for q/k/v plus the output transpose back, all introduced by the triton backend glue; ncu shows them as SM-bound strided copies at only 980 GB/s | +| Why low-medium effort | the Triton kernel's addressing already goes through strides; passing real (non-contiguous) strides for q/k/v/out removes the copies without touching kernel math. Glue-side change + kernel signature audit | +| Estimated gain | −11.6 ms/chunk ceiling (estimate, unverified; some reorder cost may remain) | +| Confidence | high | +| Risk | medium (wrong stride handling = garbage attention — caught immediately by cos/PSNR checks) | +| Files/functions | `diffsynth/models/wan_video_dit.py::flash_attention` (triton branch glue), `diffsynth/models/triton_block_sparse_attn.py` (`triton_block_sparse_attention` wrapper + `_bsfa*_kernel` stride params; note TMA descriptors constrain layouts — if TMA requires contiguity, fix the TMA=0 path first and measure) | +| Flag | `FLASHVSR_ATTN_STRIDED_IO` | +| Benchmark | primary @768; also compare `FLASHVSR_ATTN_TMA=0/1` interaction | +| Correctness | kernel-level cosine ≥ 0.9999 vs current triton path; E2E PSNR ≥ 49 dB vs sparse backend (same gate the PR used) | +| Default | OFF until log entry | + +### 2A-4 mask_gen allocation / sync cleanup (NOT the fused kernel) + +| | | +|---|---| +| Profiling says | `mask_gen` = 7.4 ms/chunk (4.9%) through a chain of tiny kernels; separately, the *sparse* backend allocates `cu_seqlens_q/k` + `head_mask_type` via `torch.tensor(..., device=...)` **per call** — a hidden H2D sync that explains its 8.2% idle | +| Why low effort | buffer reuse and hoisting constant tensors out of the hot path; no algorithm change (the fused top-k kernel is Phase 2B-3) | +| Estimated gain | −1–2 ms/chunk @768 on the triton path (estimate); larger effect on the sparse fallback path (idle reduction) | +| Confidence | medium-high | +| Risk | low | +| Files/functions | `diffsynth/models/wan_video_dit.py::generate_draft_block_mask` (intermediate allocs), `flash_attention` sparse branch (persistent cu_seqlens/head_mask_type keyed on shape/device) | +| Flag | `FLASHVSR_MASKGEN_LEAN` | +| Benchmark | primary @768 (triton) + one sparse-backend run to capture the sync win | +| Correctness | `max\|diff\|==0` (allocation/hoisting only) | +| Default | OFF until log entry | + +### 2A-5 LQ projector allocation / layout cleanup + +| | | +|---|---| +| Profiling says | lq_conv1+conv2 = 11.7 ms/chunk (7.8%); path split: pad+cat 8%, im2col copy 29%, addmm 65%. Also per-clip `cache_x = x[..., -CACHE_T:].clone()` copies in `stream_forward` | +| Why low effort | this item is only the cheap part: avoid the separate pad+cat materialization (fold cache frames into the im2col source view), reuse im2col/patch buffers across clips, drop redundant clones. The *fused im2col-GEMM kernel* is Phase 2B-4 | +| Estimated gain | −1–2 ms/chunk (estimate, unverified) | +| Confidence | medium-high | +| Risk | low (streaming-cache semantics must be preserved — bit-diff catches drift across chunk boundaries) | +| Files/functions | `examples/WanVSR/utils/utils.py::CausalConv3d.forward`, `_conv3d_gemm` / `_im2col_gemm_rows`, `Causal_LQ4x_Proj.stream_forward` | +| Flag | `FLASHVSR_LQPROJ_LEAN` | +| Benchmark | primary @768 | +| Correctness | `max\|diff\|==0` vs OFF across a full clip (streaming cache exercised) | +| Default | OFF until log entry | + +### 2A-6 (Optional experiment) CUDA Graphs on the steady chunk + +| | | +|---|---| +| Profiling says | GPU idle within steady chunk = 3.1% @768 (≤4.8 ms/chunk ceiling), ~0% @1024+. This is the *smallest* item in 2A and may not pay for its complexity | +| Why gated as experiment | graph capture requires a sync-free, allocation-stable, shape-static chunk body. Today the chunk body is *probably not* capture-safe (mask topk sizes, cache rotation, python-side branching). Go/no-go gate: after 2A-1..2A-5 land, re-measure idle; only attempt capture if idle ≥ 2% AND a capture dry-run shows no illegal ops | +| Estimated gain | ≤ −4.8 ms/chunk @768, ~0 at higher resolutions (estimate) | +| Confidence | medium | +| Risk | medium (silent staleness bugs if any buffer is re-bound between replays) | +| Files/functions | `diffsynth/pipelines/flashvsr_tiny.py::__call__` steady-chunk body | +| Flag | `FLASHVSR_CUDA_GRAPHS` | +| Benchmark | primary @768 (where the idle exists); confirm no regression @1536 | +| Correctness | `max\|diff\|==0` vs OFF | +| Default | OFF; likely stays OFF unless the win is clean | + +### Phase 2A exit criteria + +- Each of 2A-1..2A-5 (2A-6 optional) has: flag, log entry, interpretation, + correctness result. +- Cumulative log row updated; @1536 spot-check run recorded. +- Expected (to verify): steady chunk ~130–136 ms, E2E ≈ 43–46 FPS @768. + +--- + +## 2. Phase 2B — Medium-effort structural wins + +These need more design care than 2A cleanup: streams, new numerics paths, or +custom kernels — but all have strong profiling evidence. + +### 2B-1 Decoder overlap on a side CUDA stream ⟵ highest E2E leverage in 2B + +| | | +|---|---| +| Expected impact | decode is 343–430 ms of a ~2.0 s run @768 (17–21% E2E) and 22–23% of (chunks+decode) at 1024/1536. Fully hidden ⇒ **+21–29% FPS E2E** (estimate, unverified) — likely the best gain/effort in the whole roadmap | +| Why not 2A | touches scheduling semantics, not just code: per-chunk streaming decode on a second stream, event-based handoff of `cur_latents`, allocator behaviour across streams, and TCDecoder's stateful mem-blocks (`TAEHV.mem`) must only ever be touched by the decode stream. The tiny pipeline currently decodes once at the end; the long pipeline (`flashvsr_tiny_long.py`) already decodes per chunk — use it as the semantic reference | +| Design sketch | after each chunk's latent update, `record_event`; decode stream waits on the event, decodes the chunk's latents with the matching LQ cond slice, appends to output. Denoise stream never waits on decode except at the very end | +| Correctness validation | output must be `max\|diff\|==0` vs sequential decode (same TCDecoder state transitions in the same order); explicit test with short AND long clips; race detection via 3 repeated runs (identical outputs) | +| Quality/numerical risk | none if bit-identical is enforced; the risk is concurrency bugs, not math | +| Benchmark plan | primary @768 AND @1536 (decode share is larger at high res); also record peak memory (decode buffers now coexist with denoise) | +| Rollback | `FLASHVSR_DECODE_OVERLAP=0` restores the end-of-loop decode verbatim | +| Files/functions | `diffsynth/pipelines/flashvsr_tiny.py::__call__` (chunk loop + decode call), `examples/WanVSR/utils/TCDecoder.py` (`decode_video`, `apply_model_with_memblocks`, `clean_mem`) | + +### 2B-2 FP8 GEMM infrastructure (`torch._scaled_mm`) + +| | | +|---|---| +| Expected impact | GEMMs ≈ 33 ms/chunk; microbench measured ×1.55–1.72 vs bf16 at exact shapes (qkv/o ×1.59, ffn1 ×1.55, ffn2 ×1.72, lq_linear ×1.55) ⇒ ~−13 ms/chunk ceiling (estimate). bf16 GEMMs are already at 82–85% SOL — FP8 is the only remaining GEMM lever | +| Why not 2A | new numerics path: weight/activation casting strategy, scale management, and a quality gate are required. **Implementation lands in 2B, but the enable decision is governed by the Phase-4 quality protocol** — this item ships default-OFF regardless of speed | +| Scope | qkv/o, ffn1/ffn2, LQ per-layer linears first (per-tensor scales, e4m3 weights pre-cast once); cross-attn q/o optional second step | +| Correctness validation | Phase-4 protocol: PSNR/SSIM vs flag-OFF on examples 0–3, per-layer activation error audit if PSNR < 49 dB | +| Quality/numerical risk | medium — distilled one-step model may be sensitive; unknown until measured | +| Benchmark plan | primary @768 with flag ON vs OFF; per-shape kernel check via one ncu run (confirm FP8 kernels actually selected) | +| Rollback | `FLASHVSR_FP8_GEMM=0` | +| Files/functions | `diffsynth/models/wan_video_dit.py` (Linear call sites: SelfAttention q/k/v/o, DiTBlock ffn, CrossAttention q/o), `examples/WanVSR/utils/utils.py::Causal_LQ4x_Proj.linear_layers` | + +### 2B-3 Fused / semi-fused mask generation (top-k path) + +| | | +|---|---| +| Expected impact | mask_gen = 7.4 ms/chunk (4.9%); the chain is `mean-pool → einsum → +bias → softmax → topk(gatherTopK @ SM 5.2%, 0.1 waves) → 7–9 radix mini-kernels → compare`. A single fused kernel (or per-head threshold-selection kernel) should land near ~1 ms/chunk ⇒ −5 ms (estimate) | +| Why not 2A | requires a custom Triton kernel (fused softmax+select) or a rewritten selection algorithm; more design/testing than buffer hygiene | +| Correctness validation | the produced boolean block mask must be **identical** to the reference implementation on real inputs (not just statistically similar) — direct tensor equality across a full clip; then E2E bit-diff | +| Quality/numerical risk | low if mask equality is enforced; any tie-breaking difference in top-k must be resolved to match torch semantics or shown to be E2E-neutral (then it moves to Phase 4) | +| Benchmark plan | primary @768; mask_gen shrinks relatively at higher res (2.1% @1536) so no high-res requirement | +| Rollback | `FLASHVSR_FUSED_MASKGEN=0` | +| Files/functions | `diffsynth/models/wan_video_dit.py::generate_draft_block_mask` (+ new kernel module) | + +### 2B-4 Fused im2col-GEMM for the LQ projector (only if 2A-5 is not enough) + +| | | +|---|---| +| Expected impact | im2col copy = 29% of the conv path ⇒ −4 ms/chunk ceiling (estimate). Gate: skip this item if 2A-5 already gets ≥2 ms and the projector drops below ~6% of GPU | +| Why not 2A | a real fused kernel (CUTLASS 3.x conv3d or Triton implicit-GEMM with the causal window) vs cuDNN is mandatory — cuDNN 9.22 direct conv3d measured **18× slower** (152 ms vs 8.4 ms) at this shape, so there is no library shortcut | +| Correctness validation | bit-diff vs the current gemm path (`test_conv3d_gemm_parity.py` pattern: single-call + streaming-cache) | +| Quality/numerical risk | low (same math, same accumulation dtype required — fp32 accumulate) | +| Benchmark plan | primary @768 + isolated kernel bench at conv1/conv2 shapes; @1536 spot-check (projector share grows slightly with res) | +| Rollback | `FLASHVSR_CONV3D_BACKEND=gemm` (existing path untouched); new path = `FLASHVSR_CONV3D_BACKEND=fused` | +| Files/functions | `examples/WanVSR/utils/utils.py::_conv3d_gemm` (+ new kernel) | + +### Phase 2B exit criteria + +- 2B-1 decode overlap proven bit-identical and logged @768 + @1536. +- 2B-2 FP8 infra merged default-OFF with a Phase-4 gate ticket. +- 2B-3 mask equality proven; 2B-4 done or explicitly skipped with rationale. +- Cumulative table updated with the 2A+2B stack. + +--- + +## 3. Phase 3 — Major attention kernel work (`_bsfa` v2) + +The single largest target, deliberately NOT first: highest effort, highest +risk, and its payoff stacks with (rather than blocks) everything above. + +### 3.1 Evidence recap (why the kernel must be rebuilt, not tuned) + +From ncu (`--set full`, steady chunk, @768): + +``` +_bsfa_tma_kernel: 2.03 ms · SM SOL 40.2% · tensor pipe 40.2% (active 46.4%) +DRAM 134 GB/s (3.3%) · L2 hit 92.5% · achieved WGMMA ≈ 393 TFLOP/s +occupancy 12.5% = 1 CTA/SM (178 reg/thread + 229 KB smem/block) · 8 warps/SM +stalls/issue: barrier 2.39 · wait 1.00 · short_sb 0.55 (total 6.37) +scheduler: 68.6% of cycles have NO eligible warp +TMA=0 variant: 2.20 ms · 235 reg · 164 KB smem · long_sb 0.87 (TMA removed it) +``` + +Reference points at the exact shape (q 8448 × kv 25344, h12, d128, density 0.606): + +| Kernel | time | meaning | +|---|---|---| +| cuDNN dense SDPA (full attention) | 1.86 ms | computes 1.65× the FLOPs, still faster | +| ideal sparse = dense × 0.606 | **1.13 ms** | efficiency ceiling | +| `_bsfa_tma` (current) | 2.03 ms | **56% of ideal** | +| FlexAttention + BlockMask | 1.92 ms | torch-native is not the answer | + +Interpretation: a single-CTA-per-SM Triton pipeline where all 8 warps +synchronize at every stage barrier; with nothing else resident, the tensor +pipe idles ~60% of the time. This is precisely the failure mode +warp-specialization (producer/consumer) and/or 2-CTA residency solve. + +### 3.2 Candidate directions (in evaluation order) + +1. **Triton warp-specialized rewrite** — keep the existing mask format and + glue; restructure into producer (TMA loads) / consumer (WGMMA) warp groups, + tune `num_stages`/tile sizes so smem ≤ ~114 KB or regs ≤ ~96 to admit a + second CTA. Lowest integration cost; Triton 3.7 WS maturity is the risk. +2. **CUTLASS FMHA-based implementation** — start from CUTLASS 3.x Hopper FMHA + (WS pingpong pipeline, proven ≥60–75% tensor util) and add 2D block-mask + skipping. Highest ceiling, highest integration cost (C++ extension build). +3. **FA3-style hand-rolled pipeline ideas** applied to whichever base wins: + pingpong warpgroups, softmax/WGMMA overlap, TMA stores. +4. **Exact sparse mask preservation is mandatory** in all directions — the + trained sparse pattern is quality-bearing; the mask semantics of + `block_sparse_attn` must be reproduced block-for-block. +5. (Later, optional) **FP8 attention** (QK^T and/or PV in e4m3) — belongs to + Phase 4; est. additional −15–25 ms/chunk, quality unknown. + +### 3.3 Targets, tests, fallback + +| | | +|---|---| +| Expected gain | 2.03 → ~1.13–1.3 ms/call ⇒ **−21 to −26 ms/chunk @768** (estimate); relative share grows at higher res (attention ≈ 47% everywhere) | +| Development risk | high (kernel correctness across mask densities/shapes; Hopper pipeline subtleties; possible Triton compiler limitations) | +| Acceptance metrics (ncu, mandatory) | tensor pipe active ≥ 60% elapsed · barrier stall < 1.0/issue · ≥2 CTA/SM or WS with ≥16 warps/SM · duration ≤ 1.3 ms at the reference shape | +| Correctness tests | kernel-level cosine ≥ 0.9999 vs `block_sparse_attn` on randomized real-shape inputs (multiple densities incl. degenerate all-true/all-false rows); E2E PSNR ≥ 49 dB vs sparse backend; multi-chunk streaming bit-stability of the KV/cache interface | +| Fallback path | runtime chain `v2 → _bsfa_tma → block_sparse_attn` behind `FLASHVSR_ATTN_BACKEND=triton2` (sm_90-guarded, silent fallback on any error — same pattern as the existing backends) | +| Benchmark | primary @768 + @1536; ncu acceptance run archived in the log entry | +| Separate PR? | **Yes.** Self-contained, reviewable, revertible; roadmap phases 2A/2B must not wait on it | + +--- + +## 4. Phase 4 — Quality-risk / precision-risk optimizations + +Anything that can change output pixels beyond bit-identical or the +established ≥49 dB gate lives here, is **default-OFF permanently** until it +passes the protocol, and is documented with its measured quality delta. + +### 4.1 Candidates + +1. **FP8 GEMMs enable decision** (infra from 2B-2) +2. **FP8 attention** (QK^T/PV e4m3, from Phase 3 base) +3. **Aggressive elementwise fusion that changes operation order** (e.g. + folding RMSNorm/modulate chains across dtype boundaries beyond what + FUSE_NORM does today) +4. **Any approximation of the attention mask / sparse layout** (topk_ratio + reduction, coarser block masks, approximate selection) +5. **Any cache-behaviour change that could affect temporal consistency** + (KV window policy, decoder mem-block reuse across chunks) + +### 4.2 Quality protocol (applies to every Phase-4 item) + +- **Reference set:** examples 0–3 at 768x1408 (and one clip @1536), generated + fresh with the flag OFF at the current commit (do not reuse stale outputs). +- **Metrics:** per-clip PSNR and SSIM vs reference; `max|diff|`; report the + worst clip, not the average. Existing harness patterns: + `test_fuse_norm.py` (PSNR), `test_cache_lossless.py` (bit-diff). +- **Visual check:** side-by-side of the worst-PSNR clip (temporal flicker is + the failure mode PSNR misses — scrub frame-by-frame around scene motion). +- **Acceptance tiers:** ≥49 dB → eligible for "recommended config" listing; + 45–49 dB → ships flag-gated with documented delta; <45 dB → revert or + redesign. +- **Comparison rule:** quality is always measured against flag-OFF at the + same commit (isolates the numeric change from unrelated drift). +- **Default:** OFF. A Phase-4 flag may only become part of the recommended + config line after two independent bench+quality entries agree. + +--- + +## 5. Benchmark and logging system + +All results — successes, failures, reverts — go to +[`PHASE_BENCH_LOG.md`](./PHASE_BENCH_LOG.md), chronologically. Failures are +as valuable as wins; do not delete entries, mark them `revert`. + +### 5.1 Required fields per entry + +Date/time · phase · optimization name · commit hash (or patch name) · files +changed · env vars used · exact benchmark command · resolution · frames · +warmup/steady settings · FPS before/after/Δ · steady chunk before/after/Δ · +peak mem before/after · correctness result · output-difference result (if +applicable) · Nsight report path (if generated) · decision +(`keep-enabled` / `keep-behind-flag` / `revert` / `investigate` / `postpone`) +· **interpretation (2–3 sentences, mandatory)**. + +### 5.2 Tables + +Per-change table: + +| Date | Phase | Optimization | Flag | FPS Before | FPS After | Delta | Steady Chunk Before | Steady Chunk After | Peak Mem Before | Peak Mem After | Correctness | Decision | +|------|-------|--------------|------|------------|-----------|-------|---------------------|--------------------|-----------------|----------------|-------------|----------| + +Cumulative stack table (updated whenever a flag joins the recommended set): + +| Step | Enabled Optimizations | FPS | Steady Chunk Time | Peak Memory | Delta vs Phase-2 Baseline | Notes | +|------|----------------------|-----|-------------------|-------------|---------------------------|-------| + +### 5.3 The gate + +> **Do not continue to the next optimization until the current optimization +> has a benchmark entry and a short written interpretation.** + +No exceptions — including "obvious" wins and including reverts. + +--- + +## 6. Benchmark commands + +Infrastructure lives in `examples/WanVSR/profiling/`: +`run_pipe_target.py` · `nsys_run.sh` · `ncu_run.sh` · `ncu_batch.sh` · +`analyze_gaps.py` · `ncu_extract.py` · `bench_ceilings.py`. + +### 6.1 Primary (untraced — the only FPS source of truth) + +```bash +cd examples/WanVSR +FLASHVSR_CONV3D_BACKEND=gemm \ +FLASHVSR_TCDECODER_CHANNELS_LAST=1 \ +FLASHVSR_FUSE_NORM=1 \ +FLASHVSR_ATTN_BACKEND=triton \ +FLASHVSR_CACHE_MOD=1 \ +FLASHVSR_CACHE_MASK_BIAS=1 \ +FLASHVSR_PROF_STEADY=off \ +/root/FlashVSR/venv/bin/python profiling/run_pipe_target.py +``` + +Add the flag under test (`FLASHVSR_=1`) for the "after" run; run +before/after back-to-back in the same session; 3 runs, log the median. +High-res spot check: prepend `FLASHVSR_PROF_W=1536 FLASHVSR_PROF_H=2560`. + +### 6.2 Optional attribution runs (never for headline FPS) + +Traced wall time is distorted (capture stop-flush lands inside the timer; +CPU tracing overhead inflates gaps) — use these only to explain a result: + +```bash +# timeline attribution (minimal trace + GPU metrics): +FLASHVSR_NVTX=1 FLASHVSR_PROF_STEADY=0:-1 \ + ./profiling/nsys_run.sh --gpu-metrics-devices=0 --gpu-metrics-frequency=20000 +/root/FlashVSR/venv/bin/python profiling/analyze_gaps.py profiling/reports/ \ + --ref-wall-per-chunk + +# kernel deep-dive: +FLASHVSR_NVTX=1 FLASHVSR_PROF_WARMUP=0 \ + ./profiling/ncu_run.sh --target-processes application-only \ + --set full -k "regex:" --launch-skip 6 --launch-count 4 +python3 profiling/ncu_extract.py profiling/reports/ncu/.ncu-rep --csv .csv +``` + +### 6.3 Known workarounds & operational safety (keep these) + +- **ncu child-injection deadlock fix** (already baked into `ncu_run.sh`): + `TRITON_LIBCUDA_PATH=/usr/lib/aarch64-linux-gnu` and + `--target-processes application-only`. Root cause: triton's + `libcuda_dirs()` spawns `ldconfig`; ncu's child injection intermittently + deadlocks that handshake (main python stuck in `subprocess.communicate`). +- Launch long profiling jobs detached: `setsid nohup > log 2>&1 &` — + interactive aborts kill the whole process group otherwise. +- Never write `pkill -f ` / `pgrep -f ` where `` + appears verbatim in your own command line (self-kill); use the + `[b]racket` trick: `pgrep -f "run_pipe_[t]arget"`. +- Keep clocks locked during a measurement campaign (`nvidia-smi -lgc + 1980,1980`); note that the 700 W platform cap still sags clocks under load + — never compare runs taken minutes apart without checking `dmon` logs if a + result looks anomalous. + +--- + +## 7. Sequencing summary + +``` +Step 0 Fresh 3-run baseline → PHASE_BENCH_LOG.md +Phase 2A 1 RoPE cache/fusion → 2 KV ring buffer → 3 attn strided IO + → 4 mask_gen lean → 5 LQ proj lean → (6 CUDA Graphs go/no-go) +Phase 2B 1 decoder overlap → 2 FP8 GEMM infra (OFF) → 3 fused mask topk + → 4 fused im2col (conditional) +Phase 3 attention kernel v2 (separate PR, parallel-track allowed once 2A done) +Phase 4 quality-gated enables: FP8 GEMM, FP8 attention, fusion reorders, + sparsity experiments +``` + +Rationale for the order: 2A banks low-risk efficiency first (which, under the +700 W cap, also buys back clock headroom), 2B adds the structural wins with +contained blast radius, Phase 3 is the big rock developed against an already +faster baseline, and Phase 4 only ever trades quality knowingly, never by +accident. diff --git a/examples/WanVSR/profiling/analyze_gaps.py b/examples/WanVSR/profiling/analyze_gaps.py new file mode 100644 index 00000000..cc9cea69 --- /dev/null +++ b/examples/WanVSR/profiling/analyze_gaps.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Deep nsys report analyzer for FlashVSR profiling campaign. + +Consumes a report dir produced by nsys_run.sh (profile.nsys-rep) and emits: + - analysis.md : human-readable findings + - kernels.csv : top kernels within the steady window (time, count, grid, category) + - phases.csv : GPU time attributed to each NVTX leaf phase + - gaps.csv : every GPU idle gap >= threshold with NVTX phase attribution + - summary.json : machine-readable rollup (used later by ANALYSIS.md synthesis) + +Method notes +------------ +* GPU busy time = union of kernel+memcpy+memset intervals on the device + (all streams merged), computed inside each steady NVTX chunk window. +* Phase attribution maps each GPU activity to the NVTX range that was open on + the launching CPU thread at launch-API time (correlationId join), taking the + innermost containing range whose name is in the known phase set. +* Tracing inflates CPU time, so traced idle% is an UPPER bound. If + --ref-wall-per-chunk (untraced seconds/chunk) is given we also report a + corrected idle% = 1 - busy_per_chunk / ref_wall_per_chunk. +* GPU metrics (if sampled) are averaged per phase using sample timestamps + falling inside phase-attributed GPU activity intervals. +""" +import argparse +import bisect +import csv +import json +import os +import re +import sqlite3 +import subprocess +import sys +from collections import defaultdict + +# True steady window for overlap mode: chunk2 still carries the decode-0/1 +# backlog, so the default steady set is chunks 3-7 (override via +# FLASHVSR_ANALYZE_STEADY="chunk2,chunk3,..." for old-style windows). +STEADY_CHUNKS = [ + c.strip() for c in os.environ.get( + "FLASHVSR_ANALYZE_STEADY", + "chunk3,chunk4,chunk5,chunk6,chunk7").split(",") if c.strip() +] + +LEAF_PHASES = [ + "lq_proj0", "lq_proj", "lq_conv1", "lq_conv2", "lq_linears", + "patchify", "rope_freqs", "mod1", "qkv_norm", "rope", "win_part", + "kv_cat", "reorder", "mask_gen", "attn_core", "cache_trim", "win_rev", + "o_proj", "gate1", "xattn", "ffn", "head", "unpatchify", + "decode", "color_fix", +] +# parent ranges we also record for rollups +PARENT_PHASES = ["dit_forward", "self_attn"] +CHUNK_RE = re.compile(r"^chunk(\d+)$") +# Overlap mode emits per-chunk decode ranges (decode0..decodeN); fold them all +# into the single "decode" leaf phase so decoder-stream kernels stop showing +# up as (unattributed). +DECODE_RE = re.compile(r"^decode(\d+)$") + +KEY_METRICS = [ + "SMs Active [Throughput %]", + "SM Issue [Throughput %]", + "Tensor Active [Throughput %]", + "DRAM Read Bandwidth [Throughput %]", + "DRAM Write Bandwidth [Throughput %]", +] + + +def categorize(name): + n = name.lower() + if any(k in n for k in ["_bsfa", "bsfa_", "block_sparse", "fmha", "flash", "attention"]): + return "attention" + if "softmax" in n: + return "softmax/mask" + if ("cudnn" in n or "implicit_gemm" in n or "xmma_fprop" in n + or ("conv" in n and "gemm" not in n)): + return "conv-cudnn(decoder)" + if any(k in n for k in ["nvjet", "gemm", "cutlass", "cublas", "addmm", "matmul", "wgmma"]): + return "gemm" + if any(k in n for k in ["nchwtonhwc", "nhwctonchw", "tonhwc", "tonchw"]): + return "layout" + if any(k in n for k in ["topk", "radix", "sort", "scan", "bitonic"]): + return "topk/sort(mask)" + if any(k in n for k in ["upsample", "interpolate", "resize", "grid_sample"]): + return "resample(decoder)" + if any(k in n for k in ["cat", "copy", "pad", "transpose", "permute", "contiguous", + "gather", "scatter", "index", "slice", "unfold", "im2col", "col2im"]): + return "copy/cat/im2col" + if any(k in n for k in ["norm", "rms", "silu", "gelu", "relu", "sigmoid", "elementwise", + "vectorized", "reduce", "mul", "add", "div", "triton_", "mean", "pow"]): + return "elementwise/norm" + if "memcpy" in n or "memset" in n: + return "memcpy/memset" + return "other" + + +def tid_of(global_tid): + return global_tid # keep raw; NVTX and RUNTIME use same encoding + + +class Analyzer: + def __init__(self, report_dir, ref_wall_per_chunk=None, gap_threshold_us=2.0): + self.dir = report_dir + self.ref_wall_per_chunk = ref_wall_per_chunk + self.gap_thr_ns = int(gap_threshold_us * 1000) + self.db = self._open_db() + self.strings = self._load_strings() + + # ---------- loading ---------- + def _open_db(self): + rep = os.path.join(self.dir, "profile.nsys-rep") + db = os.path.join(self.dir, "profile.sqlite") + if not os.path.exists(db): + subprocess.run(["nsys", "export", "-t", "sqlite", "--force-overwrite=true", + "-o", db, rep], check=True, capture_output=True) + return sqlite3.connect(db) + + def _load_strings(self): + return dict(self.db.execute("SELECT id, value FROM StringIds")) + + def _table_exists(self, name): + return self.db.execute( + "SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?", + (name,)).fetchone()[0] > 0 + + def load(self): + # GPU activities + self.kernels = self.db.execute( + "SELECT start, end, streamId, correlationId, shortName, demangledName, " + "gridX, gridY, gridZ, blockX, blockY, blockZ, registersPerThread, " + "staticSharedMemory, dynamicSharedMemory " + "FROM CUPTI_ACTIVITY_KIND_KERNEL ORDER BY start").fetchall() + self.memcpys = self.db.execute( + "SELECT start, end, streamId, correlationId, bytes, copyKind " + "FROM CUPTI_ACTIVITY_KIND_MEMCPY ORDER BY start").fetchall() \ + if self._table_exists("CUPTI_ACTIVITY_KIND_MEMCPY") else [] + self.memsets = self.db.execute( + "SELECT start, end, streamId, correlationId FROM CUPTI_ACTIVITY_KIND_MEMSET " + "ORDER BY start").fetchall() \ + if self._table_exists("CUPTI_ACTIVITY_KIND_MEMSET") else [] + + # Runtime API rows for correlation (launch thread + time) and sync analysis + self.rt_by_corr = {} + self.sync_apis = [] + self.launch_apis = [] + for start, end, gtid, corr, name_id in self.db.execute( + "SELECT start, end, globalTid, correlationId, nameId " + "FROM CUPTI_ACTIVITY_KIND_RUNTIME"): + name = self.strings.get(name_id, "") + self.rt_by_corr[corr] = (start, end, gtid, name) + if "Synchronize" in name or "cudaEventSynchronize" in name: + self.sync_apis.append((start, end, gtid, name)) + if name.startswith("cudaLaunchKernel") or name.startswith("cuLaunchKernel"): + self.launch_apis.append((start, end, gtid)) + self.launch_apis.sort() + + # NVTX ranges (eventType 59 = PushPop) + self.nvtx = [] + for start, end, text, text_id, gtid in self.db.execute( + "SELECT start, end, text, textId, globalTid FROM NVTX_EVENTS " + "WHERE eventType = 59 AND end IS NOT NULL"): + name = text if text else self.strings.get(text_id, "") + if DECODE_RE.match(name or ""): + name = "decode" + self.nvtx.append((start, end, name, gtid)) + self.nvtx.sort() + + # per-tid sorted ranges for innermost lookup + self.nvtx_by_tid = defaultdict(list) + for start, end, name, gtid in self.nvtx: + self.nvtx_by_tid[gtid].append((start, end, name)) + self.nvtx_starts = {t: [r[0] for r in v] for t, v in self.nvtx_by_tid.items()} + + # chunk windows + self.chunks = {} + for start, end, name, gtid in self.nvtx: + m = CHUNK_RE.match(name or "") + if m: + self.chunks[name] = (start, end) + self.top_ranges = {name: (s, e) for s, e, name, _ in self.nvtx + if name in ("decode", "color_fix")} + + # GPU metrics + self.metrics = {} + if self._table_exists("GPU_METRICS") and self._table_exists("TARGET_INFO_GPU_METRICS"): + id2name = dict(self.db.execute( + "SELECT metricId, metricName FROM TARGET_INFO_GPU_METRICS")) + rows = self.db.execute( + "SELECT timestamp, metricId, value FROM GPU_METRICS").fetchall() + per = defaultdict(list) + for ts, mid, val in rows: + nm = id2name.get(mid) + if nm in KEY_METRICS: + per[nm].append((ts, val)) + self.metrics = {k: sorted(v) for k, v in per.items()} + + # ---------- helpers ---------- + def innermost_phase(self, gtid, t, allowed): + starts = self.nvtx_starts.get(gtid) + if not starts: + return None + i = bisect.bisect_right(starts, t) - 1 + ranges = self.nvtx_by_tid[gtid] + best = None + # walk left; the first containing range is the innermost, but keep + # walking to find the innermost whose name is in `allowed`. + steps = 0 + while i >= 0 and steps < 4096: + s, e, name = ranges[i] + if s <= t <= e: + if name in allowed: + return name + if best is None: + best = name + i -= 1 + steps += 1 + return None + + @staticmethod + def union_busy(intervals, w0, w1): + """Total covered time of intervals clipped to [w0,w1].""" + busy = 0 + cur_s = cur_e = None + for s, e in intervals: + if e <= w0 or s >= w1: + continue + s, e = max(s, w0), min(e, w1) + if cur_s is None: + cur_s, cur_e = s, e + elif s <= cur_e: + cur_e = max(cur_e, e) + else: + busy += cur_e - cur_s + cur_s, cur_e = s, e + if cur_s is not None: + busy += cur_e - cur_s + return busy + + @staticmethod + def gaps_in(intervals, w0, w1, thr): + """Idle gaps >= thr within [w0,w1] given sorted activity intervals.""" + gaps = [] + cursor = w0 + for s, e in intervals: + if e <= w0 or s >= w1: + continue + s, e = max(s, w0), min(e, w1) + if s > cursor and s - cursor >= thr: + gaps.append((cursor, s)) + cursor = max(cursor, e) + if w1 > cursor and w1 - cursor >= thr: + gaps.append((cursor, w1)) + return gaps + + # ---------- analyses ---------- + def run(self): + self.load() + out = {} + acts = [(k[0], k[1]) for k in self.kernels] + \ + [(m[0], m[1]) for m in self.memcpys] + \ + [(m[0], m[1]) for m in self.memsets] + acts.sort() + self.acts = acts + + steady = [(c, self.chunks[c]) for c in STEADY_CHUNKS if c in self.chunks] + if not steady: # fall back to whatever chunks exist + steady = sorted(self.chunks.items())[2:7] + out["steady_chunks"] = {c: {"span_ms": (w[1] - w[0]) / 1e6} for c, w in steady} + + # --- per-chunk busy/idle & launches --- + chunk_rows = [] + for cname, (w0, w1) in steady: + span = w1 - w0 + busy = self.union_busy(acts, w0, w1) + nk = sum(1 for k in self.kernels if k[0] >= w0 and k[1] <= w1) + launch_cpu = sum(min(e, w1) - max(s, w0) for s, e, _ in self.launch_apis + if e > w0 and s < w1) + row = dict(chunk=cname, span_ms=span / 1e6, gpu_busy_ms=busy / 1e6, + idle_pct=100 * (1 - busy / span), kernels=nk, + launch_cpu_ms=launch_cpu / 1e6) + if self.ref_wall_per_chunk: + row["idle_pct_corrected"] = 100 * (1 - (busy / 1e9) / self.ref_wall_per_chunk) + chunk_rows.append(row) + out["per_chunk"] = chunk_rows + + # --- phase attribution --- + allowed = set(LEAF_PHASES) + phase_gpu = defaultdict(float) + phase_cnt = defaultdict(int) + phase_intervals = defaultdict(list) + kern_agg = defaultdict(lambda: [0.0, 0, None]) # name -> [ns, count, meta] + steady_windows = [w for _, w in steady] + + def in_steady(s, e): + return any(s >= w0 and e <= w1 for w0, w1 in steady_windows) + + for k in self.kernels: + s, e, stream, corr, short_id, dem_id = k[0], k[1], k[2], k[3], k[4], k[5] + if not in_steady(s, e): + continue + name = self.strings.get(short_id) or self.strings.get(dem_id) or "?" + agg = kern_agg[name] + agg[0] += (e - s) + agg[1] += 1 + if agg[2] is None: + agg[2] = (k[6], k[7], k[8], k[9], k[10], k[11], k[12], k[13], k[14]) + rt = self.rt_by_corr.get(corr) + phase = None + if rt: + phase = self.innermost_phase(rt[2], rt[0], allowed) + phase = phase or "(unattributed)" + phase_gpu[phase] += (e - s) + phase_cnt[phase] += 1 + phase_intervals[phase].append((s, e)) + for m in self.memcpys: + s, e, corr = m[0], m[1], m[3] + if not in_steady(s, e): + continue + rt = self.rt_by_corr.get(corr) + phase = self.innermost_phase(rt[2], rt[0], allowed) if rt else None + phase_gpu[(phase or "(unattributed)") + "|memcpy"] += (e - s) + + total_phase = sum(phase_gpu.values()) + out["phases"] = {p: dict(gpu_ms=v / 1e6, pct=100 * v / total_phase, + count=phase_cnt.get(p, 0)) + for p, v in sorted(phase_gpu.items(), key=lambda x: -x[1])} + + # --- kernel table --- + ktable = sorted(((v[0], v[1], n, v[2]) for n, v in kern_agg.items()), + reverse=True) + out["kernels_total_ms"] = sum(v[0] for v in kern_agg.values()) / 1e6 + self.ktable = ktable + + # --- gap analysis --- + gap_rows = [] + gap_total = 0 + for cname, (w0, w1) in steady: + for gs, ge in self.gaps_in(acts, w0, w1, self.gap_thr_ns): + dur = ge - gs + gap_total += dur + mid = (gs + ge) // 2 + # phase on the main thread at gap time (any tid owning chunk ranges) + phase = None + for gtid in self.nvtx_by_tid: + p = self.innermost_phase(gtid, mid, set(LEAF_PHASES + PARENT_PHASES)) + if p: + phase = p + break + # overlapping sync API? + sync = any(s <= ge and e >= gs for s, e, _, _ in self.sync_apis) + # next kernel after gap + idx = bisect.bisect_left(self.acts, (ge, ge)) + nxt = None + for k in self.kernels: + if k[0] >= ge: + nxt = self.strings.get(k[4]) or "?" + break + gap_rows.append(dict(chunk=cname, start_us=(gs - w0) / 1e3, + dur_us=dur / 1e3, phase=phase or "?", + sync=sync, next_kernel=(nxt or "?")[:80])) + gap_rows.sort(key=lambda r: -r["dur_us"]) + out["gaps"] = dict(total_ms=gap_total / 1e6, count=len(gap_rows), + threshold_us=self.gap_thr_ns / 1e3) + # gap rollup per phase + by_phase = defaultdict(float) + for r in gap_rows: + by_phase[r["phase"]] += r["dur_us"] + out["gaps"]["by_phase_us"] = dict(sorted(by_phase.items(), key=lambda x: -x[1])) + self.gap_rows = gap_rows + + # --- memcpy inventory (steady) --- + cp = defaultdict(lambda: [0, 0.0, 0]) # kind -> [count, ms, bytes] + for m in self.memcpys: + s, e, kind, nbytes = m[0], m[1], m[5], m[4] + if not in_steady(s, e): + continue + c = cp[kind] + c[0] += 1 + c[1] += (e - s) / 1e6 + c[2] += nbytes or 0 + out["memcpy_steady"] = {str(k): dict(count=v[0], ms=round(v[1], 3), + MiB=round(v[2] / 2**20, 1)) + for k, v in cp.items()} + + # --- wall budget over whole capture (if chunk0 exists) --- + if "chunk0" in self.chunks: + all_chunks = sorted(self.chunks.items(), key=lambda x: x[1][0]) + t0 = all_chunks[0][1][0] + t1 = max(e for _, (s, e) in all_chunks) + budget = {c: (e - s) / 1e6 for c, (s, e) in all_chunks} + for nm, (s, e) in self.top_ranges.items(): + budget[nm] = (e - s) / 1e6 + t1 = max(t1, e) + covered = sum(budget.values()) + budget["(uncovered python/IO)"] = (t1 - t0) / 1e6 - covered + out["wall_budget_ms"] = budget + + # --- GPU metrics per phase --- + if self.metrics: + out["gpu_metrics_phase_avg"] = {} + interesting = ["attn_core", "ffn", "qkv_norm", "o_proj", "xattn", + "mask_gen", "lq_proj", "decode", "kv_cat", "reorder"] + for metric, samples in self.metrics.items(): + ts = [t for t, _ in samples] + vals = [v for _, v in samples] + mrow = {} + # steady-window average + sel = self._avg_in_windows(ts, vals, steady_windows) + mrow["steady_all"] = sel + for ph in interesting: + ivs = phase_intervals.get(ph) + if ivs: + mrow[ph] = self._avg_in_windows(ts, vals, ivs) + if "decode" in self.top_ranges: + mrow["decode_window"] = self._avg_in_windows( + ts, vals, [self.top_ranges["decode"]]) + out["gpu_metrics_phase_avg"][metric] = mrow + + self.out = out + return out + + @staticmethod + def _avg_in_windows(ts, vals, windows): + tot = 0.0 + n = 0 + for w0, w1 in windows: + i0 = bisect.bisect_left(ts, w0) + i1 = bisect.bisect_right(ts, w1) + for j in range(i0, i1): + tot += vals[j] + n += 1 + return round(tot / n, 2) if n else None + + # ---------- outputs ---------- + def write(self): + out = self.out + with open(os.path.join(self.dir, "summary.json"), "w") as f: + json.dump(out, f, indent=2) + + with open(os.path.join(self.dir, "kernels.csv"), "w", newline="") as f: + w = csv.writer(f) + w.writerow(["total_ms", "count", "avg_us", "category", "grid", "block", + "regs", "smem_static", "smem_dyn", "name"]) + for tot, cnt, name, meta in self.ktable[:60]: + grid = f"{meta[0]}x{meta[1]}x{meta[2]}" if meta else "" + blk = f"{meta[3]}x{meta[4]}x{meta[5]}" if meta else "" + w.writerow([round(tot / 1e6, 3), cnt, round(tot / cnt / 1e3, 1), + categorize(name), grid, blk, + meta[6] if meta else "", meta[7] if meta else "", + meta[8] if meta else "", name[:160]]) + + with open(os.path.join(self.dir, "phases.csv"), "w", newline="") as f: + w = csv.writer(f) + w.writerow(["phase", "gpu_ms", "pct", "count"]) + for p, d in out["phases"].items(): + w.writerow([p, round(d["gpu_ms"], 3), round(d["pct"], 2), d["count"]]) + + with open(os.path.join(self.dir, "gaps.csv"), "w", newline="") as f: + w = csv.writer(f) + w.writerow(["chunk", "start_us", "dur_us", "phase", "sync", "next_kernel"]) + for r in self.gap_rows[:400]: + w.writerow([r["chunk"], round(r["start_us"], 1), round(r["dur_us"], 1), + r["phase"], r["sync"], r["next_kernel"]]) + + md = [f"# nsys analysis: {os.path.basename(self.dir)}\n"] + md.append("## Steady chunks (busy/idle)\n") + md.append("| chunk | span ms | GPU busy ms | idle % (traced)" + + (" | idle % (corrected)" if self.ref_wall_per_chunk else "") + + " | kernels | launch CPU ms |") + md.append("|---|---|---|---|---|" + ("---|" if self.ref_wall_per_chunk else "")) + for r in out["per_chunk"]: + row = (f"| {r['chunk']} | {r['span_ms']:.1f} | {r['gpu_busy_ms']:.1f} " + f"| {r['idle_pct']:.1f} ") + if self.ref_wall_per_chunk: + row += f"| {r['idle_pct_corrected']:.1f} " + row += f"| {r['kernels']} | {r['launch_cpu_ms']:.1f} |" + md.append(row) + md.append("\n## GPU time by NVTX phase (steady window)\n") + md.append("| phase | GPU ms | % | launches |") + md.append("|---|---|---|---|") + for p, d in out["phases"].items(): + md.append(f"| {p} | {d['gpu_ms']:.2f} | {d['pct']:.1f} | {d['count']} |") + md.append("\n## Top kernels (steady window)\n") + md.append("| total ms | n | avg us | category | grid | name |") + md.append("|---|---|---|---|---|---|") + for tot, cnt, name, meta in self.ktable[:40]: + grid = f"{meta[0]},{meta[1]},{meta[2]}" if meta else "" + md.append(f"| {tot/1e6:.2f} | {cnt} | {tot/cnt/1e3:.1f} " + f"| {categorize(name)} | {grid} | `{name[:90]}` |") + md.append(f"\n## Gaps (>= {out['gaps']['threshold_us']:.0f} us)\n") + md.append(f"total gap: **{out['gaps']['total_ms']:.2f} ms** over " + f"{out['gaps']['count']} gaps in steady window") + md.append("\nby phase (us): " + json.dumps(out["gaps"]["by_phase_us"])) + md.append("\n### Largest 25 gaps\n") + md.append("| chunk | at us | dur us | phase | sync | next kernel |") + md.append("|---|---|---|---|---|---|") + for r in self.gap_rows[:25]: + md.append(f"| {r['chunk']} | {r['start_us']:.0f} | {r['dur_us']:.1f} " + f"| {r['phase']} | {r['sync']} | `{r['next_kernel'][:60]}` |") + if "memcpy_steady" in out: + md.append("\n## Memcpy (steady)\n```json\n" + + json.dumps(out["memcpy_steady"], indent=2) + "\n```") + if "wall_budget_ms" in out: + md.append("\n## Wall budget (full measured call, traced)\n```json\n" + + json.dumps({k: round(v, 1) for k, v in out["wall_budget_ms"].items()}, + indent=2) + "\n```") + if "gpu_metrics_phase_avg" in out: + md.append("\n## GPU metrics averages\n```json\n" + + json.dumps(out["gpu_metrics_phase_avg"], indent=2) + "\n```") + with open(os.path.join(self.dir, "analysis.md"), "w") as f: + f.write("\n".join(md) + "\n") + + print(f"[analyze] {self.dir}: busy tables written " + f"(kernels {out['kernels_total_ms']:.1f} ms in steady window; " + f"gaps {out['gaps']['total_ms']:.2f} ms)") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("report_dir") + ap.add_argument("--ref-wall-per-chunk", type=float, default=None, + help="untraced wall seconds per steady chunk (for corrected idle%)") + ap.add_argument("--gap-us", type=float, default=2.0) + args = ap.parse_args() + a = Analyzer(args.report_dir, args.ref_wall_per_chunk, args.gap_us) + a.run() + a.write() + + +if __name__ == "__main__": + main() diff --git a/examples/WanVSR/profiling/bench_attn_v2.py b/examples/WanVSR/profiling/bench_attn_v2.py new file mode 100644 index 00000000..af216cb9 --- /dev/null +++ b/examples/WanVSR/profiling/bench_attn_v2.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Phase-3 standalone bench for the block-sparse attention kernel rewrite. + +Times kernel candidates at the exact steady-state reference shape +(q 8448 x kv 25344, h12, d128, block 128, density ~0.606) against: + - v1 `_bsfa_tma_kernel_snd` (the current production kernel, kernel-only) + - cuDNN dense SDPA (ceiling reference) + - `block_sparse_attn` (correctness reference) + +Usage (from examples/WanVSR/): + python profiling/bench_attn_v2.py --capture # one-time: save a REAL mask + python profiling/bench_attn_v2.py # bench all available routes + python profiling/bench_attn_v2.py --routes v1,ws,occ,v2 + +The real mask is captured from a short pipeline run (chunk 2, last DiT block) +so the per-row cnt distribution and spatial locality match production. +""" +import argparse +import math +import os +import sys +import time + +_here = os.path.dirname(os.path.abspath(__file__)) +_wanvsr = os.path.dirname(_here) +_root = os.path.dirname(os.path.dirname(_wanvsr)) +sys.path.insert(0, _wanvsr) +sys.path.insert(0, _root) + +import torch # noqa: E402 + +DEV = "cuda" +DT = torch.bfloat16 +MASK_CACHE = os.path.join(_here, "cache", "attn_mask_768_steady.pt") + +# Reference shape (steady chunk @768x1408). NOTE: the attention consumes the +# UNTRIMMED kv window (pre 3 slots + 1 new = 264 blocks = 33792 tokens) at +# density 0.4545 — FLOP-identical to ANALYSIS's "kv 25344 @ 0.606" framing +# (avg ~120 active kv blocks per q row either way), but the kernel's real +# iteration space is 264 blocks. +NQ, NKV, H, D = 8448, 33792, 12, 128 + + +def bench(fn, it=30, warm=10): + for _ in range(warm): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(it): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / it * 1e3 # ms + + +def cos_max(a, b): + af, bf = a.float().flatten(), b.float().flatten() + cos = torch.nn.functional.cosine_similarity(af, bf, dim=0).item() + return cos, (af - bf).abs().max().item() + + +# --------------------------------------------------------------------------- +# Real-mask capture (monkeypatch, no production code changes) +# --------------------------------------------------------------------------- + +def capture_real_mask(): + os.environ.setdefault("FLASHVSR_CONV3D_BACKEND", "gemm") + os.environ.setdefault("FLASHVSR_TCDECODER_CHANNELS_LAST", "1") + os.environ.setdefault("FLASHVSR_FUSE_NORM", "1") + os.environ.setdefault("FLASHVSR_ATTN_BACKEND", "triton") + os.environ.setdefault("FLASHVSR_CACHE_MOD", "1") + os.environ.setdefault("FLASHVSR_CACHE_MASK_BIAS", "1") + os.environ.setdefault("FLASHVSR_FUSE_ROPE", "1") + os.environ.setdefault("FLASHVSR_KV_RINGBUF", "1") + os.environ.setdefault("FLASHVSR_ATTN_STRIDED_IO", "1") + os.environ.setdefault("FLASHVSR_MASKGEN_LEAN", "1") + os.environ.setdefault("FLASHVSR_LQPROJ_LEAN", "1") + os.environ["FLASHVSR_PROF_FRAMES"] = "41" # chunks 0..2 + os.environ["FLASHVSR_PROF_STEADY"] = "off" + os.environ["FLASHVSR_PROF_WARMUP"] = "0" + + import importlib.util + import diffsynth.models.wan_video_dit as ditmod + + captured = {} + orig = ditmod.generate_draft_block_mask + + def wrapper(*args, **kwargs): + m = orig(*args, **kwargs) + # keep the LAST steady-chunk mask: q side 66 blocks, kv fully grown + if m.shape[-2] * 128 == NQ and m.shape[-1] * 128 == NKV: + captured["mask"] = m.detach().to("cpu", torch.bool).clone() + return m + + ditmod.generate_draft_block_mask = wrapper + try: + spec = importlib.util.spec_from_file_location( + "infer_v1_1_tiny", os.path.join(_wanvsr, "infer_flashvsr_v1.1_tiny.py")) + _infer = importlib.util.module_from_spec(spec) + spec.loader.exec_module(_infer) + pipe = _infer.init_pipeline() + sys.path.insert(0, _here) + from run_pipe_target import build_lq # reuse the LQ cache + LQ, _, _ = build_lq("./inputs/example0.mp4", 768, 1408, 41) + with torch.no_grad(): + pipe(prompt="", negative_prompt="", cfg_scale=1.0, + num_inference_steps=1, seed=0, LQ_video=LQ, num_frames=41, + height=1408, width=768, is_full_block=False, if_buffer=True, + topk_ratio=2.0 * 768 * 1280 / (1408 * 768), kv_ratio=3.0, + local_range=11, color_fix=True) + finally: + ditmod.generate_draft_block_mask = orig + assert "mask" in captured, "no steady-shaped mask seen" + m = captured["mask"] + torch.save(m, MASK_CACHE) + d = m.float().mean().item() + print(f"[capture] saved {MASK_CACHE} shape={tuple(m.shape)} density={d:.4f}") + + +def load_mask(): + if os.path.exists(MASK_CACHE): + m = torch.load(MASK_CACHE, map_location="cpu", weights_only=True) + print(f"[mask] real mask {tuple(m.shape)} density={m.float().mean():.4f}") + return m.to(DEV) + print("[mask] WARNING: no captured mask; synthesizing density-0.4545") + torch.manual_seed(0) + m = torch.rand(1, H, NQ // 128, NKV // 128) < 0.4545 + return m.to(DEV) + + +# --------------------------------------------------------------------------- +# References +# --------------------------------------------------------------------------- + +def ref_block_sparse(q_snd, k_snd, v_snd, mask4): + from block_sparse_attn import block_sparse_attn_func + seqlen, seqlen_kv = q_snd.shape[0], k_snd.shape[0] + cu_q = torch.tensor([0, seqlen], device=DEV, dtype=torch.int32) + cu_k = torch.tensor([0, seqlen_kv], device=DEV, dtype=torch.int32) + hmt = torch.tensor([1] * H, device=DEV, dtype=torch.int32) + + def call(): + return block_sparse_attn_func( + q_snd, k_snd, v_snd, cu_q, cu_k, hmt, None, mask4, + seqlen, seqlen_kv, 0.0, deterministic=False, softmax_scale=None, + is_causal=False, exact_streaming=False, return_attn_probs=False) + return call + + +def ref_dense(q_snd, k_snd, v_snd): + from torch.nn.attention import sdpa_kernel, SDPBackend + qd = q_snd.transpose(0, 1).unsqueeze(0).contiguous() + kd = k_snd.transpose(0, 1).unsqueeze(0).contiguous() + vd = v_snd.transpose(0, 1).unsqueeze(0).contiguous() + + def call(): + with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + return torch.nn.functional.scaled_dot_product_attention(qd, kd, vd) + return call + + +# --------------------------------------------------------------------------- +# v1 kernel-only launcher (mirrors triton_block_sparse_attention_snd's TMA path) +# --------------------------------------------------------------------------- + +def make_v1_kernel_only(q, k, v, bm, BLOCK_M=128, BLOCK_N=128, num_warps=8, + num_stages=3, kernel=None): + import triton + from triton.tools.tensor_descriptor import TensorDescriptor + from diffsynth.models.triton_block_sparse_attn import ( + _make_csr, _bsfa_tma_kernel_snd) + if kernel is None: + kernel = _bsfa_tma_kernel_snd + Nq, Hh, Dd = q.shape + Nkv = k.shape[0] + sm_scale = 1.0 / math.sqrt(Dd) + Nqb = triton.cdiv(Nq, BLOCK_M) + Nkvb = triton.cdiv(Nkv, BLOCK_N) + # expand 128-granular mask rows/cols exactly onto smaller tiles + em = bm + if BLOCK_M != 128: + em = em.repeat_interleave(128 // BLOCK_M, dim=1) + if BLOCK_N != 128: + em = em.repeat_interleave(128 // BLOCK_N, dim=2) + assert em.shape[1] == Nqb and em.shape[2] == Nkvb + idx, cnt = _make_csr(em) + o = torch.empty_like(q) + q2 = q.view(Nq, Hh * Dd) + k2 = k.view(Nkv, Hh * Dd) + v2 = v.view(Nkv, Hh * Dd) + q_desc = TensorDescriptor.from_tensor(q2, [BLOCK_M, Dd]) + k_desc = TensorDescriptor.from_tensor(k2, [BLOCK_N, Dd]) + v_desc = TensorDescriptor.from_tensor(v2, [BLOCK_N, Dd]) + grid = (Nqb, Hh) + + def call(): + kernel[grid]( + q_desc, k_desc, v_desc, o, idx, cnt, sm_scale, + o.stride(1), o.stride(0), o.stride(2), + idx.stride(0), idx.stride(1), idx.stride(2), + cnt.stride(0), cnt.stride(1), + Hh, Nq, Nkv, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=Dd, + num_warps=num_warps, num_stages=num_stages) + return o + return call + + +# --------------------------------------------------------------------------- +# Route WS: one-line tl.range(warp_specialize=True) variant of the v1 kernel +# --------------------------------------------------------------------------- +try: + import triton + import triton.language as tl + + @triton.jit + def _bsfa_tma_kernel_snd_ws( + q_desc, k_desc, v_desc, O, KVIdx, KVCnt, sm_scale, + soh, som, sok, sih, sim, sic, sch, scm, H, N_Q, N_KV, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, + ): + start_m = tl.program_id(0) + off_h = tl.program_id(1) + col0 = off_h * HEAD_DIM + q = q_desc.load([start_m * BLOCK_M, col0]) + qs = (q * sm_scale).to(q.dtype) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + offs_n = tl.arange(0, BLOCK_N) + cnt = tl.load(KVCnt + off_h * sch + start_m * scm) + base = KVIdx + off_h * sih + start_m * sim + for j in tl.range(0, cnt, warp_specialize=True): + kvb = tl.load(base + j * sic) + kk = k_desc.load([kvb * BLOCK_N, col0]) + qk = tl.dot(qs, kk.T) + n = kvb * BLOCK_N + offs_n + qk = tl.where(n[None, :] < N_KV, qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2((qk - m_ij[:, None]) * 1.44269504) + alpha = tl.math.exp2((m_i - m_ij) * 1.44269504) + l_i = l_i * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + vv = v_desc.load([kvb * BLOCK_N, col0]) + acc += tl.dot(p.to(vv.dtype), vv) + m_i = m_ij + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = tl.arange(0, HEAD_DIM) + tl.store(O + off_h * soh + offs_m[:, None] * som + offs_k[None, :] * sok, + acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_Q) +except Exception: # pragma: no cover + _bsfa_tma_kernel_snd_ws = None + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--capture", action="store_true") + ap.add_argument("--routes", default="v1,dense,ws,occ,v2") + ap.add_argument("--iters", type=int, default=30) + ap.add_argument("--check", action="store_true", help="cos vs block_sparse") + args = ap.parse_args() + + if args.capture: + capture_real_mask() + return + + torch.manual_seed(0) + routes = args.routes.split(",") + bm4 = load_mask() # (1, H, Nqb, Nkvb) bool + bm = bm4[0] + q = torch.randn(NQ, H, D, device=DEV, dtype=DT) + k = torch.randn(NKV, H, D, device=DEV, dtype=DT) + v = torch.randn(NKV, H, D, device=DEV, dtype=DT) + density = bm.float().mean().item() + print(f"[shape] q {NQ} kv {NKV} h{H} d{D} density {density:.4f}") + + ref_out = None + if args.check: + ref_out = ref_block_sparse(q, k, v, bm4)() + + results = {} + if "v1" in routes: + call = make_v1_kernel_only(q, k, v, bm) + t = bench(call, args.iters) + results["v1 kernel-only (M128 N128 w8 s3)"] = (t, call()) + if "dense" in routes: + t = bench(ref_dense(q, k, v), args.iters) + results["cuDNN dense (full attention)"] = (t, None) + print(f"[ref ] ideal sparse = dense x {density:.3f} = {t*density:.3f} ms") + if "ws" in routes and _bsfa_tma_kernel_snd_ws is not None: + for ns in (2, 3, 4): + try: + call = make_v1_kernel_only(q, k, v, bm, num_stages=ns, + kernel=_bsfa_tma_kernel_snd_ws) + t = bench(call, args.iters) + results[f"ws one-liner (M128 N128 w8 s{ns})"] = (t, call()) + except Exception as e: + print(f"[ws s{ns}] FAILED: {type(e).__name__}: {str(e)[:200]}") + if "occ" in routes: + for (bmz, bnz, w, ns) in ((128, 64, 8, 3), (128, 64, 8, 4), + (64, 128, 4, 2), (64, 64, 4, 2), + (64, 64, 4, 3), (64, 64, 8, 3)): + try: + call = make_v1_kernel_only(q, k, v, bm, BLOCK_M=bmz, + BLOCK_N=bnz, num_warps=w, + num_stages=ns) + t = bench(call, args.iters) + results[f"occ (M{bmz} N{bnz} w{w} s{ns})"] = (t, call()) + except Exception as e: + print(f"[occ M{bmz}N{bnz}w{w}s{ns}] FAILED: {type(e).__name__}: {str(e)[:160]}") + if "v2" in routes: + try: + from diffsynth.models.triton_block_sparse_attn_v2 import ( + triton_block_sparse_attention_v2, bsfa_v2_kernel_only) + call = bsfa_v2_kernel_only(q, k, v, bm) + t = bench(call, args.iters) + results["v2 gluon WS"] = (t, call()) + except ImportError: + print("[v2] module not present yet") + except Exception as e: + print(f"[v2] FAILED: {type(e).__name__}: {str(e)[:300]}") + + # 2 GEMMs (QK^T, PV) x 2 flops/MAC x D per 128x128 active block + flops = 4 * D * (bm.float().sum().item() * 128 * 128) + print(f"\n{'route':42s} {'ms':>8s} {'TF/s':>7s} cos/max|d| vs sparse") + for name, (t, out) in results.items(): + tf = flops / t / 1e9 + extra = "" + if ref_out is not None and out is not None: + c, mx = cos_max(out, ref_out) + extra = f"cos={c:.6f} max|d|={mx:.4f}" + print(f"{name:42s} {t:8.3f} {tf:7.0f} {extra}") + + +if __name__ == "__main__": + main() diff --git a/examples/WanVSR/profiling/bench_ceilings.py b/examples/WanVSR/profiling/bench_ceilings.py new file mode 100644 index 00000000..b0adc95f --- /dev/null +++ b/examples/WanVSR/profiling/bench_ceilings.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Phase-3 ceiling microbenchmarks @768x1408 real shapes (GH200). + +H2: dense cuDNN SDPA reference at the exact attention shape -> how fast could + an "ideal" kernel be (density-scaled), vs the measured _bsfa 2.03 ms. +H4: bf16 vs fp8 (torch._scaled_mm) GEMM at the DiT shapes -> FP8 ceiling. +H3: im2col+GEMM split for the LQ-projector conv2 + cuDNN 9.22 conv3d check. + +All shapes match the steady-state 768x1408 run: + q tokens 8448 (66 blocks x 128), kv 25344 (3 chunks), 12 heads, d=128, + block-mask density ~0.606; DiT dim 1536, ffn 8960; conv2 2048->3072 + k=(4,3,3) s=(2,1,1) input (1,2048,4,88,48) + cache 2 frames. +""" +import os +import sys +import time + +os.environ.setdefault("FLASHVSR_CONV3D_BACKEND", "gemm") +_here = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.dirname(_here)) + +import torch +import torch.nn.functional as F + +DEV = "cuda" +DT = torch.bfloat16 + + +def bench(fn, it=30, warm=10): + for _ in range(warm): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(it): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / it * 1e3 # ms + + +def sec(title): + print(f"\n=== {title} ===") + + +def h2_attention(): + sec("H2: attention reference (q 8448, kv 25344, h12, d128)") + q = torch.randn(1, 12, 8448, 128, device=DEV, dtype=DT) + k = torch.randn(1, 12, 25344, 128, device=DEV, dtype=DT) + v = torch.randn_like(k) + from torch.nn.attention import sdpa_kernel, SDPBackend + + def dense(): + with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + return F.scaled_dot_product_attention(q, k, v) + + t_dense = bench(dense) + flops_dense = 2 * 2 * 12 * 8448 * 25344 * 128 # qk + pv + print(f"cuDNN dense SDPA : {t_dense:.3f} ms ({flops_dense/t_dense/1e9:.0f} TFLOP/s)") + density = 0.606 + print(f"ideal sparse (dense x {density}) : {t_dense*density:.3f} ms") + print(f"measured _bsfa_tma : 2.03 ms -> efficiency vs ideal-sparse " + f"{t_dense*density/2.03*100:.0f}%") + try: + import flash_attn_interface # noqa: F401 + has_fa3 = True + except Exception: + has_fa3 = False + print(f"flash_attn_3 available: {has_fa3}") + + +def h4_gemm_fp8(): + sec("H4: bf16 vs FP8 GEMM at DiT shapes (M=8448)") + M = 8448 + shapes = [("qkv/o 1536x1536", 1536, 1536), + ("ffn1 1536x8960", 1536, 8960), + ("ffn2 8960x1536", 8960, 1536), + ("lq_lin 3072x1536", 3072, 1536)] + for name, K, N in shapes: + a = torch.randn(M, K, device=DEV, dtype=DT) + b = torch.randn(N, K, device=DEV, dtype=DT) + t_bf16 = bench(lambda: a @ b.t()) + a8 = a.to(torch.float8_e4m3fn) + b8 = b.to(torch.float8_e4m3fn) + sa = torch.ones(M, 1, device=DEV) + sb = torch.ones(1, N, device=DEV) + + def fp8(): + return torch._scaled_mm(a8, b8.t(), scale_a=sa, scale_b=sb, + out_dtype=DT) + try: + t_fp8 = bench(fp8) + fl = 2 * M * K * N + print(f"{name:18s} bf16 {t_bf16*1e3:7.1f} us ({fl/t_bf16/1e9:5.0f} TF/s) | " + f"fp8 {t_fp8*1e3:7.1f} us ({fl/t_fp8/1e9:5.0f} TF/s) | x{t_bf16/t_fp8:.2f}") + except Exception as e: + print(f"{name:18s} bf16 {t_bf16*1e3:7.1f} us | fp8 FAILED: {e}") + + +def h3_conv(): + sec("H3: LQ conv2 im2col+GEMM split (2048->3072, in (1,2048,4,88,48)+cache2)") + from utils.utils import CausalConv3d + conv = CausalConv3d(2048, 3072, (4, 3, 3), stride=(2, 1, 1), + padding=(1, 1, 1)).to(DEV, DT).eval() + x = torch.randn(1, 2048, 4, 88, 48, device=DEV, dtype=DT) + cache = torch.randn(1, 2048, 2, 88, 48, device=DEV, dtype=DT) + + with torch.no_grad(): + t_gemm_path = bench(lambda: conv(x, cache)) + + # split: pad+cat / unfold(im2col) / addmm + w = conv.weight + b = conv.bias + pad = conv._padding + + def prep(): + xx = torch.cat([cache, x], dim=2) + p = list(pad) + p[4] = max(0, p[4] - cache.shape[2]) + return F.pad(xx, p, mode="replicate") + + xp = prep() + + def im2col(): + cols = xp.unfold(2, 4, 2).unfold(3, 3, 1).unfold(4, 3, 1) + return cols.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(-1, 2048 * 36) + + cols = im2col().contiguous() + wmat = w.permute(0, 1, 2, 3, 4).reshape(3072, -1) + + # weight layout for addmm: patches are (Cin,kt,kh,kw) flattened + def gemm_only(): + return torch.addmm(b, cols, wmat.t()) + + t_prep = bench(prep) + t_im2col = bench(lambda: im2col().contiguous()) + t_gemm = bench(gemm_only) + fl = 2 * cols.shape[0] * cols.shape[1] * 3072 + print(f"full gemm path : {t_gemm_path:.3f} ms") + print(f" pad+cat : {t_prep:.3f} ms") + print(f" im2col copy : {t_im2col:.3f} ms") + print(f" addmm only : {t_gemm:.3f} ms ({fl/t_gemm/1e9:.0f} TFLOP/s)") + print(f" im2col share of path: {t_im2col/t_gemm_path*100:.0f}%") + + # cuDNN 9.22 reference (the 'auto' backend path) + conv_ref = torch.nn.Conv3d(2048, 3072, (4, 3, 3), stride=(2, 1, 1)).to(DEV, DT).eval() + xx = prep() + t_cudnn = bench(lambda: conv_ref(xx)) + print(f"cuDNN conv3d : {t_cudnn:.3f} ms (gemm path is x{t_cudnn/t_gemm_path:.1f} faster)") + + +def main(): + torch.manual_seed(0) + print(f"device: {torch.cuda.get_device_name(0)}, torch {torch.__version__}") + h2_attention() + h4_gemm_fp8() + h3_conv() + + +if __name__ == "__main__": + main() diff --git a/examples/WanVSR/profiling/bench_phase7.sh b/examples/WanVSR/profiling/bench_phase7.sh new file mode 100755 index 00000000..b617bcaf --- /dev/null +++ b/examples/WanVSR/profiling/bench_phase7.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +# Phase-7 benchmark driver: production knob set + overrides, N runs, median FPS. +# Usage: bench_phase7.sh LABEL [NRUNS] [KEY=VAL ...] +# Logs to profiling/runs/phase7/