feat(sana): in-tree SANA-Streaming V2V DiT (GDN linear-attention) + C-API

[jcelerier]

Summary

In-tree port of SANA-Streaming (video-to-video, GDN linear-attention DiT) into librediffusion as an independent C++/CUDA implementation, exposed through a librediffusion_sana_* C-API. Enables no-Python streaming V2V within the library.

What's added

• src/librediffusion.sana.cu — GDN (Gated-Delta-Net) linear-attention recurrence kernel (row-split grid + shared-mem state, ~8× over the naive port) + all DiT block kernels (RMSNorm / adaLN / RoPE / softmax-attn / cross-attn / GLU-MBConv temporal FFN) + the full 20-block forward. Namespaced librediffusion::sana (no symbol clash with kernels.cu).
• src/librediffusion.sana.hpp — SanaDiT C++ class (no CUDA types leaked).
• src/librediffusion_c.sana.cpp + librediffusion_c.h — librediffusion_sana_{create, dit_forward, rollout_selftest, free} C-API.
• CMakeLists.txt — compile librediffusion.sana.cu into librediffusion_cuda; link the module + CUDA::cublas into the librediffusion DLL.

Validation (rel-L2 vs PyTorch reference, fp32)

Check	rel-L2
Full DiT model (model_in → model_out)	3.10%
Full 5-chunk streaming rollout vs sampler.sample() (through the DLL C-API)	2.47%
GDN kernel vs fp32 reference	~1e-7
VAE + text encoder (validated separately via TRT)	0.01–0.28%

Per-chunk rollout rel-L2: 3.51 / 1.54 / 1.24 / 1.98 / 2.65%. (rel-L2 is the correct metric — the latents have a large dynamic range, so clamped mean-rel is misleading.)

No regression

• build_dll.bat (Debug) + build_dll_release.bat (Release) → ninja exit 0; DLL builds.
• Existing exports unchanged: rife = 8, img2img_turbo = 5, flux2 = 17.
• gdn-kernel-test ALL PASS; DiT selftest 3.0962%.

Performance (RTX 5090)

DiT denoise step 98.95 ms (bf16); full 121-frame clip ~4.5 s ≈ 27 fps (→ ~3.5 s / ~34 fps with bf16 VAE engines). PyTorch baseline ~2.7 s. Per-chunk cost is constant → genuinely infinite-streamable.

Licensing / weights

The DiT + GDN algorithm is ported from NVIDIA SANA (Apache-2.0); this is an independent re-implementation. Model weights are research-only under the NVIDIA Open Model License and are NOT included or distributed — librediffusion_sana_create() loads them at runtime from a user-provided path.

Not included (follow-ups)

• Self-contained C++ cache production threading (the rollout currently seeds from captured caches; GDN state production is validated ~1e-7 separately).
• librediffusion_sana_run_v2v's in-process TRT VAE / gemma boundary (VAE bf16 + gemma each validated in C++ separately; dit_forward + the 5-chunk rollout are wired and validated).
• Routing through librediffusion's TRT engine manager / pinned-memory allocator / CUDA-graph capture (the module currently uses standalone cudaMalloc).

🤖 Generated with Claude Code

https://claude.ai/code/session_01EPT5weE9MMS8cyxqui83zc

[gemini-code-assist]

Code Review

This pull request introduces an in-tree C++/CUDA implementation of the SANA-Streaming DiT architecture, including the Gated-Delta-Net (GDN) linear-attention recurrence, softmax/cross-attention, and GLU-MBConv FFN, along with a C-API wrapper. The review feedback highlights several critical performance and safety improvements. First, performance can be significantly optimized by replacing expensive synchronous device-to-host round-trips with custom GPU kernels for splitting QKV and CKV tensors, and by utilizing shared memory in the gdn_bidi kernel to avoid redundant global memory accesses. Second, resource management should be improved by adding a reference counter to properly manage and destroy the static cublasHandle_t. Finally, exception safety issues should be addressed by implementing a destructor for the Arena struct to prevent GPU memory leaks, and by reordering allocations in librediffusion_sana_create to avoid leaking the wrapper handle if the SanaDiT constructor throws an exception.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Add a reference counter to track the lifetime of the global static cublasHandle_t so it can be properly destroyed when no longer needed.

static cublasHandle_t HBL = nullptr;
static int HBL_ref_count = 0;

[gemini-code-assist]

The PR description states that the GDN linear-attention recurrence kernel uses a 'shared-mem state' to achieve an ~8x speedup. However, the gdn_bidi kernel currently accesses the recurrent states skv and sz directly from global memory (b.state_kv and b.state_z) on every step of the loop. This results in highly redundant global memory reads and writes across all threads. Consider loading the recurrent states into shared memory at the start of the kernel, performing the loop updates in shared memory, and writing the final states back to global memory at the end.

[gemini-code-assist]

The Arena struct manually manages GPU memory allocations, but does not implement a destructor to automatically free tracked pointers. If an exception is thrown or an early return occurs, any allocated pointers in v will be leaked. Implementing a destructor to automatically call free() ensures exception safety and prevents GPU memory leaks.

struct Arena{ std::vector<float*> v; ~Arena() { free(); } float* a(size_t n){ float* d=dev(n); if(d)v.push_back(d); return d; } void free(){ for(auto p:v)cudaFree(p); v.clear(); } };

[gemini-code-assist]

Add custom CUDA kernels to perform QKV and CKV splitting entirely on the GPU, avoiding expensive device-to-host round-trips.

static float* LDB(const std::string& d,const char* nm,size_t n){ auto h=LB(d,nm,n); return up(h); }

__global__ void k_split_qkv(float* q, float* k, const float* qkv, int N, int C) {
    long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i >= (long)N * C) return;
    int c = i % C;
    int n = i / C;
    q[i] = qkv[(long)n * 3 * C + c];
    k[i] = qkv[(long)n * 3 * C + C + c];
}

__global__ void k_split_ckv(float* ck, float* cv, const float* ckv, int Lk, int C) {
    long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i >= (long)Lk * C) return;
    int c = i % C;
    int l = i / C;
    ck[i] = ckv[(long)l * 2 * C + c];
    cv[i] = ckv[(long)l * 2 * C + C + c];
}

[gemini-code-assist]

The current implementation performs a synchronous device-to-host copy (down), splits the QKV tensor on the CPU, and then copies the split tensors back to the GPU (cudaMemcpy). This introduces massive PCIe transfer overhead (~42 MB per block, over 600 MB per denoise step), which severely degrades performance. We can completely eliminate this round-trip by performing the split entirely on the GPU using a simple custom CUDA kernel.

        k_split_qkv<<<gr((long)Nt*C), 256>>>(q_sec, k_sec, qkv, Nt, C);

[gemini-code-assist]

Replace the host-side copy and split of CKV with the k_split_ckv GPU kernel to avoid expensive device-to-host round-trips.

    k_split_ckv<<<gr((long)Lk*C), 256>>>(ck, cv, ckv, Lk, C);

[gemini-code-assist]

Replace the host-side copy and split with the k_split_qkv GPU kernel to avoid expensive device-to-host round-trips.

        k_split_qkv<<<gr((long)N*C), 256>>>(q_sec, k_sec, qkv, N, C);

[gemini-code-assist]

Replace the host-side copy and split of CKV with the k_split_ckv GPU kernel to avoid expensive device-to-host round-trips.

    k_split_ckv<<<gr((long)Lk*C), 256>>>(ck2, cv2, ckv, Lk, C);

[gemini-code-assist]

Properly manage the lifetime of the static cublasHandle_t using the reference counter to prevent resource leaks.

SanaDiT::SanaDiT(const std::string& weights_dir){ p_=new Impl{weights_dir}; if(HBL_ref_count++ == 0) cublasCreate(&sana::HBL); }
SanaDiT::~SanaDiT(){ delete p_; if(--HBL_ref_count == 0 && sana::HBL) { cublasDestroy(sana::HBL); sana::HBL = nullptr; } }

[gemini-code-assist]

If the SanaDiT constructor throws an exception (e.g., due to missing weight files or allocation failure), the raw pointer h allocated with new will be leaked because the exception is caught in the outer try-catch block. Constructing the SanaDiT instance first before allocating h ensures exception safety.

Suggested change

	librediffusion_sana_handle librediffusion_sana_create(const char* weights_dir)
	{
	try
	{
	auto* h = new librediffusion_sana;
	h->dit = std::make_unique<SanaDiT>(weights_dir ? weights_dir : "");
	return h;
	}
	catch(const std::exception& e)
	{
	fprintf(stderr, "sana_create failed: %s\n", e.what());
	return nullptr;
	}
	}
	librediffusion_sana_handle librediffusion_sana_create(const char* weights_dir)
	{
	try
	{
	auto dit = std::make_unique<SanaDiT>(weights_dir ? weights_dir : "");
	auto* h = new librediffusion_sana;
	h->dit = std::move(dit);
	return h;
	}
	catch(const std::exception& e)
	{
	fprintf(stderr, "sana_create failed: %s\n", e.what());
	return nullptr;
	}
	}