From 1f008dc226cd25b0b6cb1ab068249c03adccd7af Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <1182563586@qq.com>
Date: Wed, 12 Aug 2026 16:28:27 +0800
Subject: [PATCH] [Diffusion][LTX-2] Allocate AdaLN outputs from one contiguous
 slab (#34508)

---
 .../ops/diffusion/triton/ltx2_ada_values.py   |  8 ++++---
 .../ops/diffusion/test_ltx2_ada_values.py     | 22 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py b/python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py
index f1d3e81ba..138934d36 100644
--- a/python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py
+++ b/python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py
@@ -165,10 +165,12 @@ def ltx2_ada_values9(
 
     batch, seq, _ = timestep.shape
     rows = int(batch * seq)
-    outs = tuple(
-        torch.empty((batch, seq, hidden), device=timestep.device, dtype=timestep.dtype)
-        for _ in range(9)
+    # Each returned output is a disjoint, contiguous view, so one allocation
+    # avoids nine allocator round trips per transformer block.
+    output_storage = torch.empty(
+        (9, batch, seq, hidden), device=timestep.device, dtype=timestep.dtype
     )
+    outs = tuple(output_storage.unbind(dim=0))
     _ltx2_ada_values9_kernel[(rows,)](
         timestep,
         scale_shift_table,
diff --git a/test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py b/test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py
index aba8f2137..3f780b2df 100644
--- a/test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py
+++ b/test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py
@@ -54,6 +54,28 @@ def test_ltx2_ada_values9(
 
     assert len(actual) == 9
     for actual_value, expected_value in zip(actual, expected):
+        assert actual_value.is_contiguous()
+        torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)
+
+
+@torch.no_grad()
+def test_ltx2_ada_values9_torch_compile_fullgraph() -> None:
+    hidden = 4096
+    scale_shift_table = torch.randn(
+        9, hidden, device=DEVICE, dtype=torch.bfloat16
+    ).contiguous()
+    timestep = torch.randn(
+        1, 1, 9 * hidden, device=DEVICE, dtype=torch.bfloat16
+    ).contiguous()
+
+    actual = torch.compile(ltx2_ada_values9, fullgraph=True)(
+        scale_shift_table, timestep
+    )
+    expected = _reference(scale_shift_table, timestep)
+
+    assert len(actual) == 9
+    for actual_value, expected_value in zip(actual, expected):
+        assert actual_value.is_contiguous()
         torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)