[wishlist] Batched 128x128 SPD inverse with two epilogue products (Video Delta rule, VDN-H3 / SGLang PR #37903)

[haochengxi]

Project and real-world impact

SGLang diffusion — PR sgl-project/sglang#37903 (VDN-H3: MiniMax-H3 video generation with hybrid window softmax + Video Delta linear attention, 8-NFE distilled). The kernel sits in the linear-attention branch of every DiT block: 50 calls per forward, 8 forwards per clip, one call per rank on an 8x B200 Ulysses deployment. Today it costs 0.935 ms per call (≈47 ms per forward, ~5% of a 0.94 s forward); the FLOP/byte floor is about 0.15 ms. Anyone serving VDN-H3 / MiniMax-H3 on Blackwell benefits.

Kernel summary

For a batch of N symmetric positive semi-definite 128×128 fp32 matrices, compute

• T = diag(alpha) · (I + A)^-1
• J = B · (I + A)^-1

i.e. a batched SPD inverse with two epilogue products (the per-frame transition and injection of the Video Delta rule; the state update downstream is S <- S·T + J). Current implementation is four library calls: torch.linalg.cholesky (cuSOLVER batched potrf), torch.linalg.solve_triangular(L, I) (cuBLAS batched trsm), linv.T @ linv and B @ inv (fp32 SIMT GEMMs) plus a row scaling — about 40 kernel launches, 0.935 ms on one B200. Breakdown per call: trsm 328 µs (2 launches), potrf family 174 µs (22 launches), tril 40 µs, the two GEMMs 116 µs. Library alternatives are slower: cholesky_inverse 1.20 ms, cholesky_solve(I, L) 1.21 ms, linalg.inv 2.78 ms, linalg.solve(M, I) 2.79 ms. The bottleneck is launch count and the inefficiency of batched trsm at 128×128, not arithmetic (≈5 GFLOP fp32 ≈ 0.15 ms of SIMT compute; ≈180 MB of traffic ≈ 30 µs).

FlashInfer Trace definition

https://gist.github.com/haochengxi/a0a970648e0e4aaa8d52040c2cb3e569/b0e96588996536ddf0cfddac0b1848a0b8b016d4 — README.md (definition, contract, acceptance) and vdn_delta_inverse_ref.py (torch-only reference implementation, workload generator with the real statistics, fp64 correctness check, CUDA-event timing and launch counting; candidates plug into CANDIDATES). Upstream reference function: delta_factor_apply (rule vdn_solve) in python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn.py on branch VDN-h3 of https://github.com/kevin-mii/sglang (PR #37903). Not yet converted to the FlashInfer-Trace JSON schema; happy to do so on request.

Input and output contract

Inputs (all contiguous fp32, CUDA):

tensor	shape	properties
A	[N, 128, 128]	symmetric PSD: A = Kᵀ diag(β) K with unit-norm rows of K (SiLU + L2 norm), β ∈ (0, 1), 1008 rows per frame → eigenvalues in [0, ~1000]; I + A is SPD
B	[N, 128, 128]	B = Vᵀ diag(β) K, no special structure
alpha	[N, 128]	values in (0, 1]; row 0 is a virtual "text" frame with alpha = 1

Outputs (contiguous fp32): T [N, 128, 128] = diag(alpha) · (I + A)^-1 (row scaling of the inverse) and J [N, 128, 128] = B · (I + A)^-1.

• The explicit inverse is required (T enters a product chain downstream); (I + A)^-1 is a contraction (eigenvalues in (0, 1]).
• N is arbitrary (it is the batch); d = 128 may be hard-coded (d = 64 support is a bonus).
• Signature: vdn_delta_factors(A, B, alpha) -> (T, J).
• Edge cases: cond(I + A) up to ~1e3 must stay finite and accurate; alpha = 1 rows; T / alpha must remain symmetric.

Representative workloads

One dominant shape: N = 707 (101 frames × 7 heads per rank), d = 128, 1008 tokens per frame — the paper clip (1344×768, 345 frames → 102 latent frames, 8× B200 Ulysses → 7 of 56 heads per rank). It runs 50 times per forward and 8 forwards per clip on every rank, so the whole distribution is this single shape repeated; shorter clips lower N to ~100 (--frames, --heads flags). Conditioning stress cases: --beta-scale 50 and --beta-scale 100. python vdn_delta_inverse_ref.py with no arguments is exactly the default workload.

Correctness requirements

Reference: today() in the script (the current four-call implementation) and an fp64 torch.linalg.inv reference. Acceptance: Frobenius relative error versus fp64 ≤ 1e-6 for both T and J (today: 2.1e-7 and 2.9e-7), no NaN/Inf at --beta-scale 50 and --beta-scale 100, deterministic output. All arithmetic in fp32; TF32 (or lower) tensor-core math must not be used for the factorization or the inverse products — it destroys the conditioning of I + A (observed by the model authors). Validation command: python vdn_delta_inverse_ref.py (prints OK/FAIL per candidate), then the two --beta-scale runs.

Target hardware

NVIDIA B200

Benchmark configuration

python vdn_delta_inverse_ref.py on one B200 (torch 2.13.0 + CUDA 13.0; any recent torch with CUDA ≥ 12.8 works, no other dependencies): 10 warm-up iterations, 50 timed iterations with CUDA events, per-call milliseconds and the CUDA launch count via torch.profiler. Baseline: 0.935 ms per call, 40 launches. Add a candidate to CANDIDATES and the script reports it next to the baseline.

Primary optimization objective

Latency

Implementation constraints

• CUDA C++ preferred (Triton or CuTe DSL acceptable); callable from PyTorch with the signature above (custom op or JIT extension); fp32 in and out; no TF32/bf16 tensor cores for the factorization or the inverse products.
• Target ≤ 0.3 ms per call in a single launch (≈0.15–0.2 ms is the floor). Suggested design: one CTA per matrix; load A into shared memory, add I, in-place Cholesky (lower triangle only), triangular inverse of L against the identity, inv = L^-T L^-1 (lower triangle then mirror), then the two epilogue products; 128 KB shared-memory budget, 256 threads.
• Hardware floor: fp32 SIMT and ≥ 128 KB of shared memory per block. B200 (SM100) is the deployment target and the machine every number here was measured on; H100 should not regress if the code builds there.
• Integration boundary: it will be wired into sglang's delta_factor_apply (rule vdn_solve) following sglang's JIT-kernel conventions (python/sglang/kernels/jit/csrc, load_jit), with the torch path kept as fallback for other shapes and devices. Apache-2.0 compatible code only.

License and public-use terms

Apache-2.0 for the definition, the reference script and the workload generator (SPDX header in the script; same license as sglang: https://github.com/sgl-project/sglang/blob/main/LICENSE).

Submission confirmations

• I can reproduce the reference implementation and benchmark from the linked materials.
• I have provided a correctness-testable definition and representative workloads.
• I understand that submission does not guarantee selection.
• I understand that the definition, implementation, benchmarks, and design summary may be published publicly unless the team explicitly agrees otherwise.

---

Filed from the community wishlist entry humanfia/KDA-wishlist#2 (same author, same definition).