[shortfin-sd] Reusable service command buffers #963

monorimet · 2025-02-13T05:29:27Z

This PR non-trivially restructures the SDXL inference process to interface with reusable command buffers.

At service startup, we attach to each fiber a set of reusable command buffers (a set of empty device arrays to be reused throughout the inference process), and upon assigning each exec request a fiber (now the "meta_fiber" equipped with command buffers), we pop an appropriate command buffer from the meta fiber, mount the command buffer on the request, and return the command buffer to the meta_fiber for reuse.

Each meta_fiber is equipped with an explicit number of command buffers, and if this count is exceeded by the number of InferenceExecProcesses assigned to the fiber, a new command buffer will be initialized ad-hoc. This shouldn't happen if service is managed properly, as we shouldn't oversubscribe fibers too heavily (thundering herd).

These changes also simplify the request structure throughout the inference processor, operating on a single request as opposed to a bundle. This greatly simplifies the work required to enable larger batch sizes, and reduces how much we need to handle program I/O during execution.

We also remove all await device usage where it wasn't required. It should only be used after VAE decode, where we are done with the GPU and need to use the result contents of a device->host transfer. fyi @daveliddell

IanNod

few nit comments but otherwise I think it looks good

IanNod · 2025-02-13T18:10:12Z

shortfin/python/shortfin_apps/sd/components/messages.py

        self.print_debug = True
-
+        self.batch_size = 1


Why are we hardcoding this? Also I don't see it used anywhere

self.exec_request.batch_size is used in the inference process. I'll fix the hardcode once larger batch sizes are actually possible.

IanNod · 2025-02-13T18:10:56Z

shortfin/python/shortfin_apps/sd/components/messages.py

        self.input_ids = input_ids

        # Denoise phase.
        self.prompt_embeds = prompt_embeds
        self.text_embeds = text_embeds
        self.sample = sample
        self.steps = steps
+        self.steps_arr = None


Doesn't look to be used anywhere

IanNod · 2025-02-13T18:14:36Z

shortfin/python/shortfin_apps/sd/components/messages.py

+        guidance_host = cb.guidance_scale.for_transfer()
+        with guidance_host.map(discard=True) as m:
+            # TODO: do this without numpy
+            np_arr = np.asarray(self.guidance_scale, dtype="float16")


is guidance_scale always fp16 or only for fp16/punet

Oops. left hardcoded to avoid parsing the sfnp.DType to np.dtype. Will fix.

daveliddell

Very elegant solution to the management of the device arrays!

All of my comments are at your discretion to disregard, as always. :-)

daveliddell · 2025-02-13T20:05:59Z

shortfin/python/shortfin_apps/sd/components/messages.py

+        # Input IDs for CLIP if they are used as inputs instead of prompts.
+        if self.input_ids is not None:
+            # Take a batch of sets of input ids as ndarrays and fill cb.input_ids
+            host_arrs = [None] * 4


Where did this 4 come from? Should it be len(cb.input_ids)?

daveliddell · 2025-02-13T21:02:35Z

shortfin/python/shortfin_apps/sd/components/service.py

@@ -350,366 +359,293 @@ def __init__(
        service: GenerateService,
        fiber,


A meta fiber is the same thing as an EquippedFiber, right? If so, can you use one or the other name consistently? Also, it would be nice if function arguments taking the thing would be called meta_fiber or equipped_fiber to distinguish it from raw sf fibers.

Yeah... regrets. Best to nip this in the bud.

daveliddell · 2025-02-13T21:11:07Z

shortfin/python/shortfin_apps/sd/components/service.py

+                await self._postprocess(device=device)
+            await device
+            self.exec_request.done.set_success()
+            self.meta_fiber.command_buffers.append(self.exec_request.command_buffer)


Does returning the command buffer to the meta fiber (equipped fiber) need to be done outside the try block, so that an exception doesn't cause the command buffer to be lost?

daveliddell · 2025-02-13T21:16:10Z

shortfin/python/shortfin_apps/sd/components/service.py

+                self.assign_command_buffer(self.exec_request)
+
+            device = self.fiber.device(0)
+            await device


What is this wait for? If it's for metrics, we might need to break the actual "inference process" out of the run function:

async def run(self): ... device = self.fiber.device(0) await device run_inference(device) @measure(...) aync def run_inference(device): phases = ... _prepare, _encode, etc

When we assign a command buffer to the request in the beginning of the inference process, we copy some host arrays (anything user-submitted like input_ids as numpy arrays) and await here for the copy to device to finish.
The await after the phases is extraneous and will be removed.

Do we need to wait for the copy to device to finish before starting a program? I thought that as long as they're sequenced correctly on a fiber, shortfin and/or the GPU command buffer guarantees correct order of execution

You know, that makes sense to me.... let's try it.

You shouldn't need to wait for it to finish. The only time you should need to await the device is if you copy back to the host and want to inspect the results (on the host).

daveliddell · 2025-02-13T21:23:24Z

shortfin/python/shortfin_apps/sd/components/service.py

-            sfnp.device_array.for_device(
-                device, [req_bs, self.service.model_params.max_seq_len], sfnp.sint64
-            ),
-        ]
        host_arrs = [None] * 4


another hard-coded 4. Should it be len(cb.input_ids}?

For SDXL inference this is a very specific shape that isn't tied to a config variable (basically it's 4 because CLIP takes positive, negative, positive pooled, and negative pooled ids from the tokenizers.). I don't think it should be driven by the shape of a command buffer element -- it won't ever change for SDXL

Meh... since we're doing h2d here it probably makes sense just to align them.

daveliddell · 2025-02-13T21:32:29Z

shortfin/python/shortfin_apps/sd/components/service.py

+        sfnp.fill_randn(sample_host, generator=generator)
+
+        cb.sample.copy_from(sample_host)
+        await device


This is a new await device that wasn't in the original code. Is it needed?

yes, as it succeeds a host to device copy.

Shouldn't need await on h2d so long as you are not mutating the host side buffer while the pipeline is progressing.

daveliddell · 2025-02-13T21:36:10Z

shortfin/python/shortfin_apps/sd/components/service.py

-        (image,) = await fn(latents, fiber=self.fiber)
-
+        (cb.images,) = await fn(cb.latents, fiber=self.fiber)
+        cb.images_host.copy_from(cb.images)
        await device


I pulled this await device out of my (now obsolete) PR, as we're not accessing the actual bytes of the buffer, only the buffer metadata, which you can do without waiting.

The decode step transfers the VAE output to host, then creates a numpy array with the result. Doesn't the numpy array instantiation need the buffer contents and not just its metadata? Hence the await device

Looks like np.frombuffer does a shallow copy, which would mean the data is free to come in later. Anyway, Stella was the one who pointed out that none of the await devices should be necessary until postprocess, where you actually use the data. I had consistent results doing just one await device before preprocess, not that I was pounding on the server particularly hard. :-)

Yeah, you only need to await when observing the bytes on the host. Just juggling views aren't touching the contents. With that said, if you were trying to have a "safety await" and didn't know for sure that the numpy API you are using only takes a view, then it would be a reasonable thing to do.

daveliddell · 2025-02-13T21:42:25Z

shortfin/python/shortfin_apps/sd/components/service.py

            if phases[InferencePhase.DECODE]["required"]:
-                await self._decode(device=device0, requests=self.exec_requests)
+                await self._decode(device=device)


Right in here is where I had added (in my now obsolete PR):

# Postprocessing needs the output data to be on the host. Even # without postprocessing, we're done with the GPU, so we wait for # it to finish here. await device

And I also removed the await device from below and from _decode. Basically, none of the phases except _postprocess are using the data from the device buffers, so the whole thing can free run up to this point.

daveliddell · 2025-02-14T00:23:06Z

shortfin/python/shortfin_apps/sd/components/service.py

-            worker_idx = self.service.get_worker_index(fiber)
-            logger.debug(f"Sending batch to fiber {fiber_idx} (worker {worker_idx})")
-            self.board(batch["reqs"], fiber=fiber)
+            fiber = self.service.idle_meta_fibers.pop(0)


Rename this fiber variable to meta_fiber? I thought (improbably) that I had caught a type mismatch on line 307 when I saw meta_fiber=fiber. :-)

Thanks for the catch!

@daveliddell

This PR non-trivially restructures the SDXL inference process to interface with reusable command buffers. At service startup, we attach to each fiber a set of reusable command buffers (a set of empty device arrays to be reused throughout the inference process), and upon assigning each exec request a fiber (now the "meta_fiber" equipped with command buffers), we pop an appropriate command buffer from the meta fiber, mount the command buffer on the request, and return the command buffer to the meta_fiber for reuse. Each meta_fiber is equipped with an explicit number of command buffers, and if this count is exceeded by the number of InferenceExecProcesses assigned to the fiber, a new command buffer will be initialized ad-hoc. This _shouldn't_ happen if service is managed properly, as we shouldn't oversubscribe fibers too heavily (thundering herd). These changes also simplify the request structure throughout the inference processor, operating on a _single_ request as opposed to a bundle. This greatly simplifies the work required to enable larger batch sizes, and reduces how much we need to handle program I/O during execution. We also remove all `await device` usage where it wasn't required. It should only be used after VAE decode, where we are done with the GPU and need to use the result contents of a device->host transfer. fyi @daveliddell --------- Co-authored-by: Ean Garvey <ean.garvey@amd.com>

eagarvey-amd added 4 commits February 12, 2025 23:28

(WIP) Reusable service command buffers

bfeeae1

Batched preprocessing fixes, command buffer tweaks

32d0a6c

Fix per_fiber and loosen up batcher.

909dc53

Set some logs to debug level.

5800327

monorimet marked this pull request as ready for review February 13, 2025 17:49

monorimet changed the title ~~(Draft) [shortfin-sd] Reusable service command buffers~~ [shortfin-sd] Reusable service command buffers Feb 13, 2025

monorimet requested review from stellaraccident, daveliddell, IanNod and vinayakdsci February 13, 2025 17:59

monorimet and others added 2 commits February 13, 2025 11:04

Merge branch 'main' into sdxl_prealloc_buffers

cfc4f69

Run pre-commit

f9b206a

IanNod approved these changes Feb 13, 2025

View reviewed changes

Fix conditionals in command buffer set method

c7c01f0

daveliddell approved these changes Feb 13, 2025

View reviewed changes

eagarvey-amd and others added 4 commits February 13, 2025 17:26

Address comments

92b9e39

Infer input ids host array list length

cf1e860

Pre-commit

f986eac

Merge branch 'main' into sdxl_prealloc_buffers

2395ff2

monorimet requested review from daveliddell and IanNod February 14, 2025 00:04

eagarvey-amd added 3 commits February 13, 2025 18:07

Remove await after command buffer application

845c66b

Fix service init fiber->meta_fiber

6277c6d

Fixups

501b154

daveliddell approved these changes Feb 14, 2025

View reviewed changes

eagarvey-amd added 3 commits February 14, 2025 11:09

Align more instances of fiber -> meta_fiber

796f83f

Only wait on device after decode.

9502843

Change fibers -> meta_fibers in ClientGenBatchProcess

d14352c

monorimet merged commit 16099b1 into main Feb 14, 2025
37 of 38 checks passed

monorimet deleted the sdxl_prealloc_buffers branch February 14, 2025 18:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[shortfin-sd] Reusable service command buffers #963

[shortfin-sd] Reusable service command buffers #963

monorimet commented Feb 13, 2025 •

edited

Loading

IanNod left a comment

IanNod Feb 13, 2025

monorimet Feb 13, 2025

IanNod Feb 13, 2025

IanNod Feb 13, 2025

monorimet Feb 13, 2025

daveliddell left a comment

daveliddell Feb 13, 2025

daveliddell Feb 13, 2025

monorimet Feb 13, 2025

daveliddell Feb 13, 2025

daveliddell Feb 13, 2025

monorimet Feb 13, 2025

daveliddell Feb 14, 2025 •

edited

Loading

monorimet Feb 14, 2025

stellaraccident Feb 14, 2025

daveliddell Feb 13, 2025

monorimet Feb 13, 2025 •

edited

Loading

monorimet Feb 13, 2025

daveliddell Feb 13, 2025

monorimet Feb 13, 2025

stellaraccident Feb 14, 2025

daveliddell Feb 13, 2025

monorimet Feb 13, 2025

daveliddell Feb 14, 2025

stellaraccident Feb 14, 2025

daveliddell Feb 13, 2025

daveliddell Feb 14, 2025

monorimet Feb 14, 2025

		@@ -350,366 +359,293 @@ def __init__(
		service: GenerateService,
		fiber,

[shortfin-sd] Reusable service command buffers #963

[shortfin-sd] Reusable service command buffers #963

Conversation

monorimet commented Feb 13, 2025 • edited Loading

IanNod left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

daveliddell left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

daveliddell Feb 14, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

monorimet Feb 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

monorimet commented Feb 13, 2025 •

edited

Loading

daveliddell Feb 14, 2025 •

edited

Loading

monorimet Feb 13, 2025 •

edited

Loading