Manually fused padding with convolution to resolve slice_forward compilation issue.

Author: Zars19

examples/stable_diffusion/run_sdxl.py

@@ -10,49 +10,68 @@
 world_size = tensorrt_llm.mpi_world_size()
 rank = tensorrt_llm.mpi_rank()
 
-parser = argparse.ArgumentParser(
-    description='run SDXL with the UNet TensorRT engine.')
-parser.add_argument('--size', type=int, default=1024)
-parser.add_argument('--seed', type=int, default=233)
-parser.add_argument('--num_inference_steps', type=int, default=50)
-parser.add_argument(
-    '--prompt',
-    type=str,
-    default=
-    "masterpiece, gouache painting, 1girl, distant view, lone boat, willow trees"
-)
-parser.add_argument('--model_dir',
-                    type=str,
-                    default=None,
-                    help='model directory')
-
-args = parser.parse_args()
-size = args.size
-seed = args.seed
-prompt = args.prompt
-num_inference_steps = args.num_inference_steps
-model_dir = f'sdxl_s{size}_w{world_size}' if args.model_dir is None else args.model_dir
-
-pipeline = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipeline.set_progress_bar_config(disable=rank != 0)
-pipeline.prepare(f'sdxl_s{size}_w{world_size}', size)
-pipeline.to('cuda')
-
-li = []
-for i in range(10):
-    st = time.time()
-    image = pipeline(num_inference_steps=num_inference_steps,
-                     prompt=prompt,
-                     generator=torch.Generator(device="cuda").manual_seed(seed),
-                     height=size,
-                     width=size).images[0]
-    ed = time.time()
-    li.append(ed - st)
-
-if rank == 0:
-    print(f'Avg latency: {np.sum(li[-7:]) / 7.0}s')
-    image.save(f"output.png")
+
+def parseArgs():
+    parser = argparse.ArgumentParser(
+        description='run SDXL with the UNet TensorRT engine.')
+    parser.add_argument('--size', type=int, default=1024)
+    parser.add_argument('--seed', type=int, default=233)
+    parser.add_argument('--num_inference_steps', type=int, default=50)
+    parser.add_argument(
+        '--prompt',
+        type=str,
+        default=
+        "masterpiece, gouache painting, 1girl, distant view, lone boat, willow trees"
+    )
+    parser.add_argument('--model_dir',
+                        type=str,
+                        default=None,
+                        help='model directory')
+    parser.add_argument('--num-warmup-runs', type=int, default=3)
+    parser.add_argument('--avg-runs', type=int, default=10)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parseArgs()
+    size = args.size
+    seed = args.seed
+    prompt = args.prompt
+    num_inference_steps = args.num_inference_steps
+    model_dir = f'sdxl_s{size}_w{world_size}' if args.model_dir is None else args.model_dir
+    num_warmup_runs = args.num_warmup_runs
+    avg_runs = args.avg_runs
+
+    pipeline = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    pipeline.set_progress_bar_config(disable=rank != 0)
+    pipeline.prepare(f'sdxl_s{size}_w{world_size}', size)
+    pipeline.to('cuda')
+
+    # warm up
+    for i in range(num_warmup_runs):
+        image = pipeline(
+            num_inference_steps=num_inference_steps,
+            prompt=prompt,
+            generator=torch.Generator(device="cuda").manual_seed(seed),
+            height=size,
+            width=size).images[0]
+
+    latency = []
+    for i in range(avg_runs):
+        st = time.time()
+        image = pipeline(
+            num_inference_steps=num_inference_steps,
+            prompt=prompt,
+            generator=torch.Generator(device="cuda").manual_seed(seed),
+            height=size,
+            width=size,).images[0]
+        ed = time.time()
+        latency.append(ed - st)
+
+    if rank == 0:
+        print(f'Avg latency: {np.sum(latency) / avg_runs}s')
+        image.save(f"output.png")

tensorrt_llm/functional.py

@@ -3441,13 +3441,17 @@ def conv1d(input: Tensor,
     return output_1d
 
 
-def conv2d(input: Tensor,
-           weight: Tensor,
-           bias: Optional[Tensor] = None,
-           stride: Tuple[int, int] = (1, 1),
-           padding: Tuple[int, int] = (0, 0),
-           dilation: Tuple[int, int] = (1, 1),
-           groups: int = 1) -> Tensor:
+def conv2d(
+    input: Tensor,
+    weight: Tensor,
+    bias: Optional[Tensor] = None,
+    stride: Tuple[int, int] = (1, 1),
+    padding: Tuple[int, int] = (0, 0),
+    dilation: Tuple[int, int] = (1, 1),
+    groups: int = 1,
+    pre_padding: Optional[Tuple[int, int]] = None,
+    post_padding: Optional[Tuple[int, int]] = None
+) -> Tensor:
     ##
     ## TODO: Document that function!
     ##
@@ -3475,6 +3479,10 @@ def conv2d(input: Tensor,
     layer.dilation_nd = dilation
     layer.num_groups = groups
     layer.dilation_nd = dilation
+    if pre_padding:
+        layer.pre_padding = pre_padding
+    if post_padding:
+        layer.post_padding = post_padding
 
     if not is_weight_constant:
         layer.set_input(1, weight.trt_tensor)

tensorrt_llm/models/unet/embeddings.py

@@ -14,7 +14,9 @@
 # limitations under the License.
 import math
 
-from ..._utils import fp32_array
+import tensorrt as trt
+
+from ..._utils import fp16_array, fp32_array
 from ...functional import concat, constant, cos, exp, silu, sin
 from ...layers import Linear
 from ...module import Module
@@ -43,7 +45,10 @@ def get_timestep_embedding(timesteps,
         for i in range(half_dim)
     ]
 
-    emb = exp(constant(fp32_array(exponent)))
+    if dtype is trt.float16:
+        emb = exp(constant(fp16_array(exponent)))
+    else:
+        emb = exp(constant(fp32_array(exponent)))
 
     ts_shape = list(timesteps.size())
     ts_shape.append(1)

tensorrt_llm/models/unet/pp/conv2d.py

@@ -12,25 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import tensorrt as trt
 
-from ....functional import allgather, concat, conv2d, slice, stack, unsqueeze
+from ....functional import (allgather, concat, conv2d, slice, stack, unsqueeze)
 from ....layers import Conv2d
 from ....mapping import Mapping
 from ....module import Module
 
 
-def pad(input, pad):
-    assert input.ndim() == 4
-    n, c, h, w = input.shape
-    padded_input = slice(input,
-                         starts=[0, 0, -pad[2], -pad[0]],
-                         sizes=[n, c, pad[2] + h + pad[3], pad[0] + w + pad[1]],
-                         mode=trt.SampleMode.FILL,
-                         fill_value=0.0)
-    return padded_input
-
-
 class DistriConv2dPP(Module):
 
     def __init__(self,
@@ -54,20 +42,22 @@ def sliced_forward(self, x):
         idx = mapping.tp_rank
         h_begin = output_h * idx * stride - padding
         h_end = output_h * (idx + 1) * stride + padding
-        final_padding = [padding, padding, 0, 0]
+        pre_padding = [0, padding]
+        post_padding = [0, padding]
         if h_begin < 0:
             h_begin = 0
-            final_padding[2] = padding
+            pre_padding[0] = padding
         if h_end > h:
             h_end = h
-            final_padding[3] = padding
+            post_padding[0] = padding
         sliced_input = slice(x, [0, 0, h_begin, 0], [b, c, h_end - h_begin, w])
-        padded_input = pad(sliced_input, final_padding)
-        return conv2d(padded_input,
+        return conv2d(sliced_input,
                       self.conv.weight.value,
                       None if self.conv.bias is None else self.conv.bias.value,
                       stride=self.conv.stride,
-                      padding=(0, 0))
+                      padding=(0, 0),
+                      pre_padding=tuple(pre_padding),
+                      post_padding=tuple(post_padding))
 
     def forward(self, x, *args, **kwargs):
         mapping = self.mapping
@@ -78,14 +68,16 @@ def forward(self, x, *args, **kwargs):
             boundary_size = self.conv.padding[0]
 
             def create_padded_x(x, boundaries):
+                preH = 0
+                postH = 0
                 if mapping.tp_rank == 0:
                     b = boundaries.select(0, mapping.tp_rank + 1).select(0, 0)
-                    concat_x = concat([x, b], dim=2)
-                    padded_x = pad(concat_x, [0, 0, boundary_size, 0])
+                    padded_x = concat([x, b], dim=2)
+                    preH = boundary_size
                 elif mapping.tp_rank == mapping.tp_size - 1:
                     b = boundaries.select(0, mapping.tp_rank - 1).select(0, 1)
-                    concat_x = concat([b, x], dim=2)
-                    padded_x = pad(concat_x, [0, 0, 0, boundary_size])
+                    padded_x = concat([b, x], dim=2)
+                    postH = boundary_size
                 else:
                     b0 = boundaries.select(0, mapping.tp_rank - 1).select(0, 1)
                     b1 = boundaries.select(0, mapping.tp_rank + 1).select(0, 0)
@@ -97,7 +89,7 @@ def create_padded_x(x, boundaries):
                         ],
                         dim=2,
                     )
-                return padded_x
+                return padded_x, preH, postH
 
             n, c, h, w = x.shape
             b0 = slice(x, [0, 0, 0, 0], [n, c, boundary_size, w])
@@ -107,13 +99,11 @@ def create_padded_x(x, boundaries):
 
             boundaries = allgather(unsqueeze(boundary, 0),
                                    group=mapping.tp_group)
-            padded_x = create_padded_x(x, boundaries)
-            output = conv2d(
-                padded_x,
-                self.conv.weight.value,
-                self.conv.bias.value,
-                stride=self.conv.stride,
-                padding=(0, self.conv.padding[1]),
-            )
-
+            padded_x, preH, postH = create_padded_x(x, boundaries)
+            output = conv2d(padded_x,
+                            self.conv.weight.value,
+                            self.conv.bias.value,
+                            stride=self.conv.stride,
+                            pre_padding=(preH, self.conv.padding[1]),
+                            post_padding=(postH, self.conv.padding[1]))
         return output