[KDA-Pilot] Add diffusion residual-gate CUDA fast path for LTX2

[BBuf]

Motivation

Add a native CUDA JIT fast path for the LTX-2.3 residual-gate update pattern from KDA-Pilot task b200_diffusion_residual_gate_add__multi_shape / BBuf/KDA-Pilot#134.

The hot pattern is equivalent to:

hidden_states = hidden_states + update * gate

for contiguous [B, L, D] tensors with either a full gate or a broadcast gate shaped [B, 1, D]. The LTX-2.3 HQ two-stage path hits large broadcast rows such as [1, 32640, 4096] and full-gate rows such as [1, 8160, 4096] across the DiT residual, attention, and MLP updates.

Modifications

• Add diffusion_residual_gate_add lightweight JIT CUDA kernel.
• Mark KDA-developed CUDA csrc files with the MIT HAN Lab Kernel Design Agents provenance comment.
• Register it as an eager SGLang custom op with a fake impl so direct torch.compile tracing sees an opaque op instead of tracing JIT/module loading.
• Support both full gates and row-broadcast gates, including the large LTX-2.3 HQ broadcast shape with a row-tiled CUDA path.
• Wire LTX2 transformer residual-gate updates to try CUDA by default, then fall back to the existing Triton fuse_scale_shift_kernel implementation if the JIT load/launch fails.
• Add a concise correctness test over broadcast/full gates, dtype coverage, tail shapes, unsupported input rejection, plus a direct torch.compile(fullgraph=True) custom-op test.
• Add a standalone benchmark for the traced LTX-style and neighboring diffusion shapes.

Accuracy Tests

python3 -m py_compile \
  python/sglang/jit_kernel/diffusion/residual_gate_add.py \
  python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py \
  test/registered/jit/diffusion/test_residual_gate_add.py \
  test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py
python3 -m ruff format --check ...
python3 -m ruff check ...
git diff --check

B200:

CUDA_VISIBLE_DEVICES=7 \
PYTHONPATH=/tmp/sglang_ltx2_rga_pr/python:/tmp/sglang_ltx2_rga_pr \
TVM_FFI_CACHE_DIR=/tmp/tvm-ffi-rga-b200-yuanluo \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
python3 -m pytest -q test/registered/jit/diffusion/test_residual_gate_add.py -s

Result: 12 passed.

H100:

CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH=/tmp/sglang_ltx2_rga_pr/python:/tmp/sglang_ltx2_rga_pr \
TVM_FFI_CACHE_DIR=/tmp/tvm-ffi-rga-h100-pr \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
python3 -m pytest -q test/registered/jit/diffusion/test_residual_gate_add.py -s

Result: 12 passed.

Speed Tests and Profiling

Kernel benchmark command:

CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH=/tmp/sglang_ltx2_rga_pr/python:/tmp/sglang_ltx2_rga_pr \
TVM_FFI_CACHE_DIR=/tmp/tvm-ffi-rga-<gpu> \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
python3 test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py

B200 kernel benchmark:

Workload	Gate	Torch us	Triton us	CUDA us	CUDA / Triton
ltx2_bcast_s32640_c4096	bcast	415.24	132.98	120.04	1.108x
ltx2_full_s8160_c4096	full	65.46	47.07	41.65	1.130x
ideogram4_bcast_s4096_c4608	bcast	64.71	34.94	19.92	1.754x
flux2_bcast_s4608_c3072	bcast	47.46	34.49	15.69	2.198x
flux2_bcast_s4096_c3072	bcast	41.41	34.53	15.56	2.219x
flux2_bcast_s512_c3072	bcast	13.41	34.45	15.51	2.221x
ltx2_full_s126_c2048	full	13.17	35.02	13.54	2.587x

H100 kernel benchmark:

Workload	Gate	Torch us	Triton us	CUDA us	CUDA / Triton
ltx2_bcast_s32640_c4096	bcast	598.64	295.20	264.68	1.115x
ltx2_full_s8160_c4096	full	136.22	91.67	88.52	1.036x
ideogram4_bcast_s4096_c4608	bcast	91.05	43.35	41.30	1.050x
flux2_bcast_s4608_c3072	bcast	69.64	32.81	31.45	1.043x
flux2_bcast_s4096_c3072	bcast	61.24	29.83	27.50	1.085x
flux2_bcast_s512_c3072	bcast	9.21	22.33	10.26	2.176x
ltx2_full_s126_c2048	full	8.41	20.50	8.43	2.431x

B200 Nsight Compute check for the large broadcast workload:

Implementation	NCU duration us	DRAM %	SM %	Global load inst	L1/TEX load sectors	Regs/thread	Achieved occupancy
Original flat CUDA candidate	204.99	56.79	70.54	1,566,720	25,067,520	44	51.08%
Triton BLC fallback	157.60	73.34	16.07	1,566,720	20,889,600	112	22.79%
CUDA row-tile 4x	132.54	87.06	42.66	1,175,040	18,800,640	30	89.93%

The row-tiled CUDA path reuses each broadcast gate vector across 4 rows, removes the expensive flat-index modulo from the hot loop, and reduces repeated gate memory traffic.

B200 model E2E benchmark:

CUDA_VISIBLE_DEVICES=7 \
PYTHONPATH=/tmp/sglang_ltx2_rga_<impl>/python:/tmp/sglang_ltx2_rga_<impl> \
TVM_FFI_CACHE_DIR=/tmp/tvm-ffi-rga-b200-real-<impl> \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
HF_HUB_OFFLINE=1 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
sglang generate \
  --backend=sglang \
  --model-path=Lightricks/LTX-2.3 \
  --pipeline-class-name=LTX2TwoStageHQPipeline \
  --ltx2-two-stage-device-mode original \
  --prompt="A beautiful sunset over the ocean" \
  --seed=42 \
  --width=1920 --height=1088 --num-frames=121 \
  --num-inference-steps=15 \
  --save-output --warmup \
  --perf-dump-path /tmp/ltx2_rga_real/<impl>/perf.json \
  --output-file-path /tmp/ltx2_rga_real/<impl>/<impl>.mp4

• Host: ion-b200, NVIDIA B200, CUDA_VISIBLE_DEVICES=7.
• Model path: Lightricks/LTX-2.3, native LTX2TwoStageHQPipeline; no diffusers fallback observed.
• Shape: T2V 1920x1088x121f, recommended HQ num_inference_steps=15, seed=42.
• Measurement: sglang generate --warmup; metric is perf.json total_duration_ms, with the warmup request excluded.

Implementation	E2E ms	LTX2AVDenoisingStage ms	LTX2RefinementStage ms	Peak reserved MB	Speedup
Main torch expression	46644.08	29839.68	12758.12	67902	1.00x
CUDA JIT fast path	45198.37	28816.24	12419.63	67642	1.032x (+3.20%)

B200 Visual Output Check

Visual equivalence check uses the same full LTX-2.3 HQ command as the model benchmark above.

• Host/GPU: ion-b200, NVIDIA B200, CUDA_VISIBLE_DEVICES=7.
• Main commit: 4d06d4c97f; PR commit: b4dfe508ca45.
• Prompt/shape: A beautiful sunset over the ocean, 1920x1088x121f, recommended HQ num_inference_steps=15, seed=42.
• Native LTX2TwoStageHQPipeline was used for both runs; no diffusers fallback was observed.
• Decoded-video comparison: SSIM All 0.984868, PSNR average 41.339961 dB.

Left is main; right is this PR:

MP4 artifacts: main.mp4, pr.mp4, side-by-side.mp4.

Checklist

• Default path tries CUDA first and falls back to Triton on failure.
• No environment variable is required to enable the CUDA kernel.
• Direct custom op registration is eager and has a fake impl for torch.compile compatibility.
• Direct torch.compile(fullgraph=True) custom-op unit test included.
• B200 kernel benchmark and unit test included.
• H100 kernel benchmark and unit test included.
• LTX-2.3 HQ native model benchmark and visual output check completed on B200.

Review and Merge Process

This is a focused KDA-Pilot kernel integration PR. The expected useful effect is the faster residual-gate update group in LTX-2.3 HQ DiT blocks; whole-model speedup is visible but modest because attention, MLP, and decode still dominate the full video generation runtime.

---

CI States

Latest PR Test (Base): ❌ Run #28272388900
Latest PR Test (Extra): ❌ Run #28272388830

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/28272388900/job/83780924164?pr=29361

[BBuf]