From 36eb19fd614eb2a13b0ecbc7325a3d2e5ce04c2d Mon Sep 17 00:00:00 2001 From: Zhenshan Xie Date: Mon, 9 Nov 2026 14:35:41 -0800 Subject: [PATCH] docs(sana_wm): record why FP8 is not shipped Add a README for the sana_wm family with the FP8 study of the stage-1 DiT. About 80% of the block GEMM and convolution time runs inside the family plugins (GLUMBConvTemp FFN, cross attention, the bf16 gate and projection linears), which call ATen kernels, so TensorRT Quantize/Dequantize does not reach it. Only attn.qkv and the camera q/k/v projections are graph matmuls. Timed in isolation at the manifest config (72,160 rows), FP8 makes them 1.4x to 1.7x faster, which saves about 1 ms per block, 7.3% (B200) and 7.8% (B300) of the block GEMM and convolution time on TensorRT 11.4 (11.1 is the same). The refiner DiT packs every block into one plugin and has no graph GEMM. FP8 is not shipped and no FP8 build was made. NVFP4 has the same limit. The README lists what would change the decision. Documentation only; no code change. Signed-off-by: Zhenshan Xie --- families/sana_wm/README.md | 149 +++++++++++++++++++++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 families/sana_wm/README.md diff --git a/families/sana_wm/README.md b/families/sana_wm/README.md new file mode 100644 index 0000000000..5b579eedf8 --- /dev/null +++ b/families/sana_wm/README.md @@ -0,0 +1,149 @@ +# SANA-WM + +This family builds `Efficient-Large-Model/SANA-WM_bidirectional` for +`world_model_generation`: an image, a prompt and a camera path give a 321-frame +704 x 1280 video. The bundle holds TensorRT plans for the stage-1 text encoder +(Gemma 2), the stage-1 DiT (20 blocks, 60 steps, batch 2 for CFG), the refiner +text encoder and text connector, the refiner DiT (48 LTX-2 blocks, 3 steps) and +the LTX-2 VAE; the native C++ runtime runs the pipeline. All plans are bf16. Most +of the DiT compute runs in family-owned TensorRT plugins (`native_plugins/`) +that call ATen (PyTorch C++) kernels. The family has no quantized build; the +build refuses `quantization`. + +## FP8 quantization study (not shipped) + +FP8 was studied for the stage-1 DiT and is not shipped, and no FP8 build was +made: about 80% of the DiT GEMM and convolution time runs inside the family +plugins, where TensorRT Quantize/Dequantize does not apply. Only `attn.qkv` and +the camera `q/k/v_proj_cam` projections are TensorRT graph matmuls. Making +those four FP8 saves at most about 1 ms per block (7% of the GEMM and +convolution time, a smaller share of the whole DiT call), and it would need a +new calibration path for a 36,080-token input. This section records the +evidence, so that the study does not have to be repeated. + +### Summary and decision + +Stage-1 DiT, one block, manifest config (36,080 tokens, batch 2 = 72,160 rows), +TensorRT 11.4: + +| Item | B200 | B300 | +|---|---|---| +| Graph GEMMs (`attn.qkv`, `q/k/v_proj_cam`), bf16 | 2.98 ms | 3.22 ms | +| Same GEMMs in FP8 (TensorRT Quantize/Dequantize) | 1.93 ms | 2.14 ms | +| GEMMs and convolutions inside plugins (bf16 ATen) | 11.43 ms | 10.74 ms | +| Graph share of the GEMM and convolution time | 21% | 23% | +| FP8 saving per block, share of the GEMM and convolution time | 1.06 ms, 7.3% | 1.09 ms, 7.8% | +| FP8 saving per 60-step video (20 blocks x 60 calls) | at most 1.3 s | at most 1.3 s | + +Decision: do not ship FP8 for SANA-WM. FP8 on the graph part cannot meet the +speed bar that the other FP8 builds meet (the DiT call also has the GDN scan, +attention, camera and normalization plugins, which FP8 does not touch), and a +quality study of a 321-frame video model would cost more than the gain. The +refiner DiT has no graph GEMM at all. + +### Where the DiT compute runs + +Per-token linear weights of one stage-1 block (from the checkpoint header, +`dit/sana_wm_1600m_720p.safetensors`) and where they run on `main`: + +| Weights | Params (M) | Runs in | FP8 Q/DQ possible | +|---|---|---|---| +| `mlp.inverted_conv` (1x1, 2240 to 13440), `mlp.depth_conv` (3x3 depthwise), `mlp.point_conv` (1x1, 6720 to 2240), `mlp.t_conv` (3x1 temporal) | 30.1 + 0.1 + 15.1 + 15.1 | `SanaWmGlumbconvTemp` plugin: weights are plugin fields, `at::conv2d` in `enqueue` | no | +| `cross_attn.q_linear`, `cross_attn.proj` (and `kv_linear` on 300 text tokens) | 5.0 + 5.0 (+ 10.0) | `SanaWmCrossAttention` plugin: `at::matmul`, `at::linear`, `at::scaled_dot_product_attention` | no | +| `attn.output_gate`, `attn.proj`, `attn.out_proj_cam`, `plucker_proj` | 4 x 5.0 | `SanaWmGateProj` bf16 linear plugin: `at::linear` | no | +| `attn.qkv` | 15.1 | TensorRT graph matmul (`_add_linear`; the plugin path is built with `enabled=False`) | yes | +| `attn.q_proj_cam`, `k_proj_cam`, `v_proj_cam` | 3 x 5.0 | TensorRT graph matmul (`_add_linear`) | yes | + +The graph path also has the GDN scan (`SanaWmGdnScan` and its combined +variants), softmax attention in every fourth block (`SanaWmSoftmaxAttention`), +UCPE camera transforms, layer norm, modulation and frame gate as plugins. The +small per-sample linears (timestep, caption and final layer) are graph matmuls, +but run on a few rows only. + +The refiner DiT (37.8 GB, `refiner_dit_builder.py`) packs each of its 48 LTX-2 +blocks into one `SanaWmLtxVideoBlock` plugin (`packed_weights` field; all +linears are `at::linear` and attention is `at::scaled_dot_product_attention`), +so it has no graph GEMM that could take FP8. The refiner text connector is also +built from plugins. + +### Method + +No FP8 or NVFP4 checkpoint of SANA-WM exists on the Hub (the +Efficient-Large-Model releases are bf16 or fp32), and no FP8 build of the family +was made. The full stage-1 DiT plan was not profiled; the GEMMs of one block were +timed in isolation: + +- Shapes from the checkpoint header and `config.yaml` at the manifest config: + 321 frames, 704 x 1280, VAE stride 8 x 32 x 32, so 41 x 22 x 40 = 36,080 latent + tokens per sample, batch 2 (CFG), hidden size 2240. +- Graph GEMMs: one strongly typed TensorRT network per GEMM shape, input + `[72160, K]` bf16 times a constant `[K, N]` weight. FP8: input Quantize and + Dequantize (per-tensor), FP8 weight constant with a per-tensor Dequantize to + bf16, the form used by the other FP8 builds. Random weights; timing only, no + quality measured. +- Plugin GEMMs and convolutions: the same ATen calls as the plugin `enqueue` + code in bf16 (`at::conv2d` in channels-last for the GLUMBConvTemp FFN, + `at::matmul`/`at::linear` for the 2240 x 2240 projections). +- B200 and B300, one GPU per job, TensorRT 11.4.0.106 (target) and 11.1.0.106, + torch 2.14.0 (CUDA 13). CUDA events, median of 9 interleaved rounds of 10 + calls. + +### Results + +GEMM time per call, ms (x is bf16 time over FP8 time): + +| GEMM (count per block) | B200 bf16 | B200 FP8 | x | B300 bf16 | B300 FP8 | x | +|---|---|---|---|---|---|---| +| `attn.qkv` 2240 to 6720 (1) | 1.419 | 0.835 | 1.70 | 1.570 | 0.938 | 1.67 | +| `q/k/v_proj_cam` 2240 to 2240 (3) | 0.521 | 0.364 | 1.43 | 0.551 | 0.399 | 1.38 | + +Plugin GEMMs and convolutions per block, bf16 ATen, ms: + +| Plugin op | B200 | B300 | +|---|---|---| +| FFN `inverted_conv` 1x1, 2240 to 13440 | 3.188 | 3.014 | +| FFN `depth_conv` 3x3 depthwise, 13440 channels | 1.891 | 1.832 | +| FFN `point_conv` 1x1, 6720 to 2240 | 1.602 | 1.536 | +| FFN `t_conv` 3x1, 2240 to 2240 | 1.593 | 1.410 | +| cross attention `q_linear` + `proj` | 1.067 | 0.926 | +| cross attention `kv_linear` (300 text tokens) | 0.017 | 0.015 | +| `output_gate`, `proj`, `out_proj_cam`, `plucker_proj` | 2.076 | 2.005 | +| sum | 11.43 | 10.74 | + +TensorRT 11.1 gives the same picture: graph GEMMs 3.04 ms bf16 and 1.95 ms FP8 +on B200 (3.13 and 2.13 ms on B300), plugin ops 11.48 ms (10.78 ms), FP8 saving +7.5% (7.2%) of the GEMM and convolution time. + +### Why FP8 does not pay here + +- TensorRT Quantize/Dequantize only reaches graph layers. The plugins hold their + weights as bf16 plugin fields and call ATen kernels in `enqueue`, so 79% (B200) + to 77% (B300) of the block GEMM and convolution time cannot become FP8 without + changing the plugins. +- The two graph GEMMs that can be FP8 are 1.4x to 1.7x faster in isolation, + which is about 1 ms per block. The whole DiT call has more non-GEMM work (GDN + scan, attention, UCPE, norms), so the end-to-end gain is below 7%. +- A quality check would need the full 102 GB checkpoint, the 60-step CFG + pipeline and a calibration set at the 36,080-token shape for this small gain. + +### NVFP4 (not tried) + +NVFP4 has the same limit: TensorRT block quantization applies only to graph +matmuls, which are the same two projections. + +### What would change it + +If the FFN, cross-attention and gate projections were graph matmuls (or the +plugins had FP8 GEMM paths), the same isolated timing shows that FP8 could save +about 4.7 ms of the 14.4 ms of block GEMM and convolution time on B200 (about +33%): `inverted_conv` 3.09 to 1.73 ms, `point_conv` 1.70 to 1.15 ms, `t_conv` +1.76 to 1.13 ms, the six 2240 x 2240 projections 0.58 to 0.39 ms each (TensorRT +11.4, as graph GEMMs). The depthwise convolution (1.9 ms) has no FP8 kernel. Such +a change would also need calibration and a video quality study (the +`minimax_h3` FP8 build is an example). + +Revisit FP8 when: + +- the GLUMBConvTemp FFN, cross attention or `SanaWmGateProj` linears are lowered + as TensorRT graph layers, or +- the family plugins gain FP8 GEMM paths.