Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Sparse VideoGen (SVG) attention on TPUs for Wan in MaxDiffusion, introducing modules for routing, token placement, and kernel execution. The review feedback highlights several critical issues to address: potential TypeError and AttributeError crashes in the pipeline and attention configuration logic, possible NaN propagation bugs in the attention kernels when all keys are masked out, a missing divisibility check in the Pallas partial kernel, and opportunities to optimize compilation by statically checking step_index.
Perseus14
left a comment
Perseus14
left a comment
There was a problem hiding this comment.
Added some minor comments. PTAL
| def _closest_prefix(items, target): | ||
| ordered = sorted( | ||
| items, | ||
| key=lambda x: (-x.alpha, -x.exact_pairs, x.qi, x.kj), | ||
| ) | ||
| best_k, best_err, running = 0, abs(float(target)), 0 | ||
| for k, tile in enumerate(ordered, 1): | ||
| running += tile.real_pairs | ||
| err = abs(float(running) - float(target)) | ||
| if err < best_err: | ||
| best_k, best_err = k, err | ||
| return {(x.qi, x.kj) for x in ordered[:best_k]} |
There was a problem hiding this comment.
_closest_prefix() minimizes the global pair-budget error without ensuring each query tile retains any support. It can select zero boundary tiles even when there are no full tiles.
When Q = K = 0 and every value equal to one, the kernel returns all zeros. Attention over any nonempty set of these values must return ones.
This also occurs for individual query rows on larger grids at low densities. Preserving the global budget does not guarantee valid attention locally.
There was a problem hiding this comment.
Fixed in e5298b8. If rounding leaves a query tile without any keys, we now retain one boundary tile for that row and include its cost in the budget report. Added coverage checks and a TPU regression test for the constant-value example.
| low_noise_config = getattr(self.low_noise_transformer, "config", None) | ||
| if getattr(self, "use_svg_attention", False) or getattr(low_noise_config, "use_svg_attention", False): | ||
| if use_cfg_cache or use_magcache: | ||
| raise ValueError("SVG sparse attention cannot be combined with CFG cache or MagCache.") | ||
|
|
There was a problem hiding this comment.
use_svg_attention is in attention_config and not transformer.config as per
maxdiffusion/src/maxdiffusion/pipelines/wan/wan_pipeline.py
Lines 354 to 363 in e38cbc4
There was a problem hiding this comment.
Fixed in e5298b8. The low-noise check now reads use_svg_attention from attention_config. Added a regression test for low-noise SVG with incompatible caches.
| with jax.named_scope("svg_route_profile"): | ||
| sequence_length = query.shape[2] | ||
| sample_pool_size = min(max(int(sample_max_row), 1), sequence_length) | ||
| sample_count = min(max(int(query_count), 1), sequence_length) |
There was a problem hiding this comment.
This might be more accurate sample_count = min(max(int(query_count), 1), sample_pool_size)
There was a problem hiding this comment.
Fixed in e5298b8. sample_count is now capped by sample_pool_size.
Preserve nonempty rounded support for each real query tile and account for coverage repair in the pair-budget report. Keep supported selections unchanged. Read low-noise SVG enablement from attention_config, cap routing samples by the configured pool, and remove the unused compute_band_width wrapper. Add coverage, budget, sampling, and nested-configuration regressions, including a TPU constant-value attention check. CPU source-isolated checks pass; TPU execution and the full integration suite remain to be run.
3 participants
Video diffusion processes long sequences of video tokens across many denoising steps, making self-attention expensive. Sparse attention reduces this cost by skipping interactions that contribute little to the output.
This PR introduces an opt-in sparse attention implementation for MaxDiffusion: Sparse VideoGen (SVG) for Wan on TPUs. SVG samples a few queries to choose a spatial or temporal pattern for each head, then computes only the selected mask. Density and active steps/layers are configurable, and SVG is disabled by default.
See SVG.md for an illustrated explanation, configuration examples, and profiling instructions.
The implementation rounds sparse boundaries to hardware tiles, approximately preserving the attention-pair budget. Selected interior tiles share one mask-free kernel; only tiles touching sequence padding need validity masks.
The feature is organized into six commits:
A review follow-up adds configuration and shape guards, a static inactive-step shortcut, and regression tests.
Wan2.2 720p results on TPU v6e-8, using 81 frames and 40 denoising steps:
Timings are medians of three warm runs against same-node optimized fixed-M dense controls. PSNR uses FFmpeg’s aggregate YUV metric against dense outputs.
In an earlier evaluation using the moderate SVG policy, Wan2.2 at 720p retained 98.1% of the dense baseline’s mean VBench dimension score across a 31-prompt, 16-dimension screening subset.
SVG supports inference through the custom Ulysses/ring backends; training, Animate, external masks, periodic support, chunked Ulysses, CFG cache, and MagCache are unsupported.