Modal: decode a batch's prompts together on the GPU

Merged
SI
3m ago
+245-48

Stacked on #1233: the base is codex/batch-full-report-moments, so only this PR's commit shows. Retarget to main after #1233 merges.

What this changes

generate_batch answered its prompts one at a time, which leaves an A100 mostly idle on an 8B model. After #1233, full_report runs 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_p and do_sample are decoded together, up to 8 per generate() call:

  • The processor pads on the left, so every row's new tokens start at the same place.
  • Each answer is cut to its own max_new_tokens.

Billing stays honest

  • Each answer's elapsed_seconds is 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

  • A single prompt takes the same path as /generate.
  • If the GPU runs out of memory, the group is split in half and retried, down to one prompt.
  • Any other error in a grouped call falls back to answering one at a time.
  • The optional request field group_size (1–8, default 8) lets a caller force one-at-a-time decoding (1).

Tests

modal/test_serve_music_flamingo.py adds:

  • grouping by settings and size,
  • grouped decoding with separate clip files that are deleted afterwards,
  • equal time shares,
  • cleanup when the model fails.

35 tests pass locally. A bench entrypoint (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:

  • Run full_report on the same song before and after.
  • Answers from the greedy presets should match or nearly match. Batched bf16 maths can shift a few tokens.
  • Time and billed seconds should drop.

If answers degrade, any caller can send group_size: 1, or revert this PR.

🤖 Generated with Claude Code

Summary by cubic

Speeds up /generate_batch on GPUs by decoding prompts that share temperature, top_p, and do_sample together, up to 8 per generate() call, instead of one at a time.

Behavior

  • Each answer's elapsed_seconds is its share of its group's time, so the answers still add up to the time billed.
  • Prompts pad on the left so every answer starts at the same place; each is cut to its own max_new_tokens.
  • A single prompt takes the same path as /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.
  • The optional request field group_size (1–8, default 8) lets callers force one-at-a-time decoding with 1.
  • Adds a bench entrypoint that compares one-at-a-time and grouped decoding on a real song; this hasn't been verified on a GPU yet.