Wenbo Ji 嵇文博

Cache-safe attention forward

• Sequence-bound nonpaged V-cache TMA loads and zero-fill unused rows to prevent NaN contamination. B200: 146 GPU regressions and memory checks pass; valid-value NaNs and paged KV retain their contracts.

  #2983 Sequence-bounded V-cache loads

Cache locality and launch geometry

• Adapt co-authored CLC L2 swizzling while preserving STATIC 2CTA. B200 BF16 causal, 180 1CTA/CLC shapes: 1.039× geomean speedup; 42 shapes regress.

  #2976 CLC L2 swizzling and STATIC 2CTA

Asymmetric-head backward correctness

• Separate dV TMEM/SMEM staging widths from Q/K for asymmetric head dimensions. Targeted MHA/GQA gradient regressions cover silent backward errors.

  #2975 Asymmetric dV memory staging

Fused sparse-attention data movement

• Fuse VSA64 scatter and layout transforms into one Triton kernel. Wan2.1-VSA-T2V-14B, 50 steps, 1/2 B200s: 1.155×/1.145× speedup with bitwise-equal tested outputs (SM100 BF16, D128, tile64; inference only).

  #1900 Fused VSA64 scatter and layouts

Sequence-parallel control streams

• Shard Wan-VACE 1.3B/14B image, video, and mask control streams across sequence-parallel ranks. Teacher-forced checks compare against Diffusers; FP32 SP2 uses 1e-5 tolerance. BF16 parallel outputs are not bitwise-equivalent.

  #1903 Wan-VACE conditioning and sharding

Long-sequence attention tiling

• Optimize Attn-QAT full-tile/masked-tail loops and 64×128 backward tiling. Fixed-batch Wan2.1-1.3B SP2, two B200s: forward/backward/AdamW step latency falls 10.9%, preserving tested outputs and gradients bit-for-bit.

  #1899 SM100 Attn-QAT tile paths

Native diffusion execution

• Native Ovis-Image 7B execution with checkpoint mapping, CFG, and TP/SP. Single-B200 exact-tier output matches pinned Diffusers bit-for-bit; four-B200 replay is validated within BF16 control error.

  #43152 Native Ovis-Image execution

Sharded attention consistency

• Pad exact-tier SDPA queries to align TP/Ulysses head-shard reduction paths. Operator and two-B200 checks are bit-identical; whole-model differences remain. Default lossless-tier behavior is unchanged.

  #43242 Exact-tier head shards

Request-local cache state

• Isolate TeaCache CFG branches and reset cached steps and residuals per request. Reused Wan instances match fresh execution bit-for-bit on 1/2 B200s under serial/parallel CFG.

  #42422 CFG-aware TeaCache lifecycle

Contiguous patch-embedding layouts

• Restore contiguous Wan patch embeddings for cuDNN execution. Wan1.3B BF16, 480×832, 81 frames, 50 steps: two-B200 Ulysses/Ring median latency −6.3%/−7.0%, with bit-identical outputs.

  #861 Wan patch-layout correction

Valid-token attention boundaries

• Propagate valid-token bounds through padded sequence shards and joint attention, guarding unsupported layouts. Eight-B200-node tests match selected ZImage/Cosmos3/QwenImage/LingBot references exactly; FLUX2 retains a small replay error.

  #851 SP padding and joint tokens

Distributed execution ownership

• Give prompt sharding, output gathering, and parallel VAE decoding explicit ownership. An eight-B200 DP/Ulysses/CFG matrix verifies missing/duplicated-output fixes and coordinated VAE collectives.

  #902 DP, output rank, and parallel VAE

Memory-efficient VAE execution

• Blend overlapping QwenImageVAE tiles for memory-efficient encoding and decoding. One B200, 2048² BF16, tile/stride 128/64: isolated encode/decode peaks 7.77/8.53 → 2.15/2.43 GiB; tiled outputs differ slightly.

  #1720 Overlapping Qwen-Image VAE tiles

Guidance-aware cache ownership

• Separate conditional/unconditional TeaCache state through shared CFG contexts. Regressions cover branch-local updates and fallback without an initialized distributed group.

  #8475 Shared CFG cache contexts

Request failure recovery

• Isolate failed requests and clean up partially submitted batches. GPU regressions preserve engine reuse for later requests, covering list and generator inputs.

  #8450 Request isolation and batch cleanup

Batch-invariant attention kernels

• Fixed-order Triton H3 attention without atomics or split-KV. One H100, BF16 non-causal forward/backward: batch/head invariant, up to 48× vs strict deterministic reference at 8,349 tokens; still 1.6–2.3× slower than SDPA.

  #576 Fixed-order H3 attention

Rollout/recompute reproducibility

TP4 validation → EP4 extension

• Map Qwen3-Next TP4 weights across vLLM rollout and Megatron recompute. Four B200s, EP1: 64 replay cells validate bitwise parity through two updates and checkpoint reload.

  #543 TP4 rollout and recompute

• Extend TP4 replay to EP4 with fixed-order expert routing. One four-B200 node: 16/16 rollout/recompute and 97/97 batch-invariance cases pass; EP4 training and checkpoint export stay disabled.

  #557 EP4 replay extension

Bounded-memory log-probability backward

• Fuse direct-write logp gradients with a bounded chunked fallback. One B200, BF16 32768×151936 logits: peak memory 46.37 → 9.27 GiB including gradients; some tuned cross-entropy alternatives are faster.

  #495 Zero-workspace logp gradients

Branch-aware video caching

• Separate Wan V2V CFG cache contexts for TaylorSeer, FirstBlockCache, and MagCache. One-B200 Wan1.3B smoke tests verify cache-off parity and request reset; speed and generation quality are unqualified.

  #14948 Wan V2V cache contexts