Repository navigation
Modal: decode a batch's prompts together on the GPU - #1234
Open
sidneyswift wants to merge 1 commit into
Open
sidneyswift wants to merge 1 commit into
sidneyswift wants to merge 1 commit into
Conversation
generate_batch answered its prompts one at a time, which leaves an A100 mostly idle on an 8B model; full_report now runs its 13 sections in one batch, so the time adds up. Prompts that share temperature, top_p and do_sample are now decoded together, up to 8 per generate() call (the processor pads on the left; each answer is cut to its own max_new_tokens). - Each answer's elapsed_seconds is its share of its group's time, so the answers still add up to the time billed. - Fails safe: one prompt takes the /generate path; out of GPU memory splits the group in half and retries; any other grouped error falls back to one prompt at a time. - Optional request field group_size (1-8, default 8); 1 restores one-at-a-time decoding. - bench entrypoint compares the two on a real song and audio URL. Not yet run on a GPU (the local Modal workspace is paused); verify on production by comparing a full_report against the one-at-a-time baseline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review: 2
Please check the settings in the CodeRabbit UI or the
 |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
6 issues found across 3 files
Confidence score: 3/5
- OOM recovery in
serve_music_flamingo.pymay fail because the failed call’s tensors can remain referenced by the exception traceback while the handler retries. Move the split and retry outside theexceptblock. max_new_tokens=max(limits)makes short prompts decode up to the longest prompt’s limit before their output is sliced. Group prompts by compatible limits or add per-row stopping.- The new fallback and grouping behavior is not adequately covered in
test_serve_music_flamingo.py: add tests for OOM splitting and fallback, and varytop_pto verify prompts with different values are not grouped together. serve_music_flamingo.pynow combines model serving with a separate benchmark entrypoint in a 713-line module. Consider splitting the benchmark into its own module to keep the code easier to maintain.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="modal/test_serve_music_flamingo.py">
<violation number="1" location="modal/test_serve_music_flamingo.py:192">
P2: The grouping fixture never varies `top_p`, so an implementation that groups prompts with different `top_p` values together would still pass these tests. Include `top_p` in the helper settings and assert that otherwise-identical prompts with different `top_p` values stay in separate groups.</violation>
<violation number="2" location="modal/test_serve_music_flamingo.py:255">
P2: Custom agent: **Flag AI Slop and Fabricated Changes**
The fail-safe paths advertised in the PR description have no test coverage. `run_prompts` in serve_music_flamingo.py splits an OOM group in half and retries, and falls back to one prompt at a time on any other grouped error, but no test exercises either branch: `FakeModel` raises unconditionally and `run_batch` propagates the error, so the retry code never runs. Add regression tests with a fake model that fails on multi-prompt groups and succeeds on singles, and one that raises `torch.cuda.OutOfMemoryError` on the grouped call to verify the half-split retry, since the PR itself notes these paths were never run on a GPU.</violation>
</file>
<file name="modal/serve_music_flamingo.py">
<violation number="1" location="modal/serve_music_flamingo.py:483">
P2: Run the OOM split after the `except` block exits, not inside it. The failed call's tensors stay referenced by the exception traceback while the handler runs, so `empty_cache()` cannot free them and the half-size retries start with that memory still held. Record the OOM, leave the handler, then call `empty_cache()` and recurse.</violation>
<violation number="2" location="modal/serve_music_flamingo.py:505">
P2: `max_new_tokens=max(limits)` makes every row decode to the largest per-prompt limit before slicing. Group compatible limits separately or add per-row stopping so one 2048-token request does not inflate every short prompt's GPU work and memory.</violation>
<violation number="3" location="modal/serve_music_flamingo.py:683">
P2: Custom agent: **Enforce Clear Code Style and Maintainability Practices**
This module is now 713 lines and combines model serving with a separate benchmark entrypoint, far beyond the under-100-line limit. Split the benchmark and independent helpers into focused modules.</violation>
</file>
<file name="modal/README.md">
<violation number="1" location="modal/README.md:47">
P3: The benchmark code fence is not closed because prose follows the backticks on the same line. Put the prose after a standalone closing fence so the following test instructions render normally.</violation>
</file>
Architecture diagram
sequenceDiagram
participant Client as API Client
participant API as FastAPI Modal Endpoint
participant Batch as generate_batch Handler
participant Grouper as group_jobs
participant Runner as run_batch
participant Model as ModelRunner
participant GPU as GPU / generate()
participant Files as Audio Files
Note over Client,Files: Batch Generation Request Flow
Client->>API: POST /generate_batch
API->>API: Parse & validate prompts, defaults, group_size
API->>Batch: generate_batch(request)
Batch->>Batch: Validate group_size in [1, 8], default 8
Batch->>Batch: Download/transcode audio once
Batch->>Grouper: group_jobs(jobs, group_size)
Grouper->>Grouper: Group by (temperature, top_p, do_sample)
Grouper-->>Batch: List of job index groups
loop For each group of jobs
Batch->>Runner: run_batch(jobs, audio_path, run_prompts, group_size)
Runner->>Runner: create clip per job if needed
Runner->>Model: run_prompts(requests) - grouped call
Model->>Model: Check if single request in group
alt Single request in group
Model->>Model: run_prompt(request)
Model->>GPU: generate() with individual settings
GPU-->>Model: Single decoded answer
else Multiple requests in group
Model->>Model: _generate_together(requests)
Model->>Model: Build conversations, apply chat template, pad left
Model->>GPU: generate() with max(limits), shared settings
GPU-->>Model: Batched outputs
alt OOM error
Model->>Model: empty_cache + split group in half
Model->>Model: Recursively run_prompts on halves
else Other error
Model->>Model: Fall back to run_prompt per request
end
Model->>Model: batch_decode each row, cut to own max_new_tokens
end
Model-->>Runner: Answers in group order
Runner->>Runner: Calculate time share per answer
Runner->>Runner: Delete clips for group (finally)
Runner-->>Batch: Answers with elapsed_seconds
end
Batch-->>API: Answers in original job order
API-->>Client: JSON response with answers
Note over Model,GPU: Time share calculation
Note over Runner: Each answer gets elapsed_seconds = group_time / group_size
Note over Runner: Sum of shares equals total time spent
Reply with feedback, questions, or to request a fix.
View guided diff | Turn on auto-fix | Re-trigger cubic
| return [f"answer to {call['prompt']}" for call in group] | ||
|
|
||
|
|
||
| def jobs_with(*settings): |
There was a problem hiding this comment.
P2: The grouping fixture never varies top_p, so an implementation that groups prompts with different top_p values together would still pass these tests. Include top_p in the helper settings and assert that otherwise-identical prompts with different top_p values stay in separate groups.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/test_serve_music_flamingo.py, line 192:
<comment>The grouping fixture never varies `top_p`, so an implementation that groups prompts with different `top_p` values together would still pass these tests. Include `top_p` in the helper settings and assert that otherwise-identical prompts with different `top_p` values stay in separate groups.</comment>
<file context>
@@ -164,21 +165,44 @@ def test_a_clip_window_at_infinity_or_too_large_for_a_float_is_refused():
+ return [f"answer to {call['prompt']}" for call in group]
+
+
+def jobs_with(*settings):
+ """One plain job per settings tuple (temperature, do_sample), ids j0, j1..."""
+ return parse_batch_prompts([
</file context>
| except torch.cuda.OutOfMemoryError: | ||
| torch.cuda.empty_cache() | ||
| half = len(requests) // 2 | ||
| return self.run_prompts(requests[:half]) + self.run_prompts(requests[half:]) |
There was a problem hiding this comment.
P2: Run the OOM split after the except block exits, not inside it. The failed call's tensors stay referenced by the exception traceback while the handler runs, so empty_cache() cannot free them and the half-size retries start with that memory still held. Record the OOM, leave the handler, then call empty_cache() and recurse.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/serve_music_flamingo.py, line 483:
<comment>Run the OOM split after the `except` block exits, not inside it. The failed call's tensors stay referenced by the exception traceback while the handler runs, so `empty_cache()` cannot free them and the half-size retries start with that memory still held. Record the OOM, leave the handler, then call `empty_cache()` and recurse.</comment>
<file context>
@@ -423,6 +459,62 @@ def run_prompt(self, prompt, audio_path, max_tokens, temperature, top_p, do_samp
+ except torch.cuda.OutOfMemoryError:
+ torch.cuda.empty_cache()
+ half = len(requests) // 2
+ return self.run_prompts(requests[:half]) + self.run_prompts(requests[half:])
+ except Exception as err:
+ print(f"[generate_batch] grouped generation of {len(requests)} failed ({err!r}); "
</file context>
| _, _, _, temperature, top_p, do_sample = requests[0] | ||
| outputs = self.model.generate( | ||
| **inputs, | ||
| max_new_tokens=max(limits), |
There was a problem hiding this comment.
P2: max_new_tokens=max(limits) makes every row decode to the largest per-prompt limit before slicing. Group compatible limits separately or add per-row stopping so one 2048-token request does not inflate every short prompt's GPU work and memory.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/serve_music_flamingo.py, line 505:
<comment>`max_new_tokens=max(limits)` makes every row decode to the largest per-prompt limit before slicing. Group compatible limits separately or add per-row stopping so one 2048-token request does not inflate every short prompt's GPU work and memory.</comment>
<file context>
@@ -423,6 +459,62 @@ def run_prompt(self, prompt, audio_path, max_tokens, temperature, top_p, do_samp
+ _, _, _, temperature, top_p, do_sample = requests[0]
+ outputs = self.model.generate(
+ **inputs,
+ max_new_tokens=max(limits),
+ temperature=temperature,
+ top_p=top_p,
</file context>
| assert answers[0]["elapsed_seconds"] == answers[2]["elapsed_seconds"] == answers[3]["elapsed_seconds"] | ||
|
|
||
|
|
||
| def test_run_batch_deletes_the_clips_when_the_model_fails(tmp_path): |
There was a problem hiding this comment.
P2: Custom agent: Flag AI Slop and Fabricated Changes
The fail-safe paths advertised in the PR description have no test coverage. run_prompts in serve_music_flamingo.py splits an OOM group in half and retries, and falls back to one prompt at a time on any other grouped error, but no test exercises either branch: FakeModel raises unconditionally and run_batch propagates the error, so the retry code never runs. Add regression tests with a fake model that fails on multi-prompt groups and succeeds on singles, and one that raises torch.cuda.OutOfMemoryError on the grouped call to verify the half-split retry, since the PR itself notes these paths were never run on a GPU.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/test_serve_music_flamingo.py, line 255:
<comment>The fail-safe paths advertised in the PR description have no test coverage. `run_prompts` in serve_music_flamingo.py splits an OOM group in half and retries, and falls back to one prompt at a time on any other grouped error, but no test exercises either branch: `FakeModel` raises unconditionally and `run_batch` propagates the error, so the retry code never runs. Add regression tests with a fake model that fails on multi-prompt groups and succeeds on singles, and one that raises `torch.cuda.OutOfMemoryError` on the grouped call to verify the half-split retry, since the PR itself notes these paths were never run on a GPU.</comment>
<file context>
@@ -206,15 +230,38 @@ def test_run_batch_answers_in_order_with_each_prompts_own_settings_and_clip(tmp_
+ assert answers[0]["elapsed_seconds"] == answers[2]["elapsed_seconds"] == answers[3]["elapsed_seconds"]
+
+
+def test_run_batch_deletes_the_clips_when_the_model_fails(tmp_path):
+ source = tone(str(tmp_path / "song.wav"), 3)
+ jobs = parse_batch_prompts([
</file context>
| @@ -31,6 +31,10 @@ | |||
| MAX_NEW_TOKENS = 512 # default max tokens to generate | |||
There was a problem hiding this comment.
P2: Custom agent: Enforce Clear Code Style and Maintainability Practices
This module is now 713 lines and combines model serving with a separate benchmark entrypoint, far beyond the under-100-line limit. Split the benchmark and independent helpers into focused modules.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/serve_music_flamingo.py, line 683:
<comment>This module is now 713 lines and combines model serving with a separate benchmark entrypoint, far beyond the under-100-line limit. Split the benchmark and independent helpers into focused modules.</comment>
<file context>
@@ -574,3 +678,36 @@ def main():
print(f"Took {result['elapsed_seconds']}s")
+
+
+@app.local_entrypoint()
+def bench(audio_url_file: str, prompts_file: str, out_file: str = ""):
+ """
</file context>
|
|
||
| ```bash | ||
| modal run modal/serve_music_flamingo.py::bench --audio-url-file url.txt --prompts-file prompts.json | ||
| ``` The validation and the batch loop are pure Python and tested |
There was a problem hiding this comment.
P3: The benchmark code fence is not closed because prose follows the backticks on the same line. Put the prose after a standalone closing fence so the following test instructions render normally.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At modal/README.md, line 47:
<comment>The benchmark code fence is not closed because prose follows the backticks on the same line. Put the prose after a standalone closing fence so the following test instructions render normally.</comment>
<file context>
@@ -31,7 +31,20 @@ both or neither) to hear only that stretch, cut fresh from the one download
+
+```bash
+modal run modal/serve_music_flamingo.py::bench --audio-url-file url.txt --prompts-file prompts.json
+``` The validation and the batch loop are pure Python and tested
locally (ffmpeg and ffprobe must be on the PATH):
</file context>
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #1233: the base is
codex/batch-full-report-moments, so only this PR's commit shows. Retarget tomainafter #1233 merges.What this changes
generate_batchanswered its prompts one at a time, which leaves an A100 mostly idle on an 8B model. After #1233,full_reportruns its 13 sections in one batch, so the full report's total time is the sum of its sections.Now prompts that share
temperature,top_panddo_sampleare decoded together, up to 8 pergenerate()call:max_new_tokens.Billing stays honest
elapsed_secondsis its share of its group's time. The answers still add up to the time actually spent, which is what the API bills.Fails safe
/generate.group_size(1–8, default 8) lets a caller force one-at-a-time decoding (1).Tests
modal/test_serve_music_flamingo.pyadds:35 tests pass locally. A
benchentrypoint (modal run modal/serve_music_flamingo.py::bench ...) compares one-at-a-time and grouped decoding on a real song.Not verified on a GPU
The local Modal workspace is paused, so this hasn't run on an A100. Verify after deploy:
full_reporton the same song before and after.If answers degrade, any caller can send
group_size: 1, or revert this PR.🤖 Generated with Claude Code
Summary by cubic
Speeds up
/generate_batchon GPUs by decoding prompts that sharetemperature,top_p, anddo_sampletogether, up to 8 pergenerate()call, instead of one at a time.Behavior
elapsed_secondsis its share of its group's time, so the answers still add up to the time billed.max_new_tokens./generate; an out-of-memory error splits the group in half and retries, and any other grouped error falls back to one at a time.group_size(1–8, default 8) lets callers force one-at-a-time decoding with1.benchentrypoint that compares one-at-a-time and grouped decoding on a real song; this hasn't been verified on a GPU yet.Written for commit 84d6a4d. Summary will update on new commits.