feat(llm-bench): add ragged prompt batches

[janderson4]

Description

Update the generation latency benchmark's unrouted mode to send one completions request containing a deterministic array of distinct, ragged prompts instead of one prompt with n=batch_size.

Prompt lengths use evenly spaced gamma quantiles with configurable shape --gamma-shape-k (default 0.9). Prompt text cycles through the new 256-row JSONL corpus, repeats each selected row, reserves the required analysis suffix, and truncates at token boundaries. Derived lengths are clamped to the effective model context.

Code overview

flowchart TD
    A[Load tokenizer, 256 prompts, and matrix pairs] --> B[Resolve model context and tokenize source rows once]
    B --> C[Compute gamma-quantile lengths for each matrix cell]
    C --> D[Repeat token IDs, truncate, and apply chat template]
    D --> E[Warm the exact constructed prompts]
    E --> F{Routing enabled?}
    F -->|No| G[Send one prompt-array request with n=1]
    F -->|Yes| H[Assign prompts to workers by deterministic round robin]
    H --> I[Send concurrent routed n=1 requests]
    G --> J[Derive latency per target-model forward]
    I --> J

Loading

Type of change

• Feature
• Bug fix
• Breaking change
• Documentation only

Testing

• python3 -m py_compile llm_bench/gen_load_test.py
• git diff --check
• Deterministic prompt construction, row wrapping, suffix preservation, and target lengths checked with a lightweight tokenizer stub
• Independent agent review: PASS
• Black check (Black is not installed in the current environment)

Perf evidence

• Not a perf-sensitive serving change. This changes benchmark workload generation only and makes no performance claim.

Five-run validation: latency consistency and target batch attainment

Validated benchmark commit 44bf4f9 before the corpus-only cleanup (request construction and runtime mechanics are unchanged) on accounts/infra/deployments/ycmsc0uc with target mean sequence lengths 1k–64k, batch sizes 1–64, deployed model context 262,144, fixed 100-token generation, batched prompt warmup, and speculative acceptance disabled while retaining draft/verifier compute.

Runs: 29667639633, 29670952953, 29672329584, 29672766217, 29673223232. Their generation windows did not overlap.

Reported latency consistency

Metric across 98 matrix cells	Result
Median per-cell CV	0.87%
P90 per-cell CV	4.83%
Maximum per-cell CV	18.20%
Cells with CV <= 5%	89 / 98
Cells with CV <= 10%	97 / 98
Whole-run mean latency	23.077–23.401 ms (1.40% span)
Whole-run median latency	17.785–17.970 ms (1.03% span)

Physical forward consistency and target batch size

Pulled 52,606 serving-engine generation step logs containing physical active=N and fwd=Xms values; 49,989 positive forwards aligned to measured cells.

• Requested batch size was reached and was the modal active occupancy in 490/490 run/cells.
• Active occupancy never exceeded target in any cell.
• 48,993/49,989 forwards (98.01%) ran at exact target occupancy.
• 407/490 cells spent every logged positive forward at target.
• Exact-target physical forward latency was more stable than the response-header estimate: median CV 0.66%, P90 1.11%, maximum 4.59%; all 98 cells were within 5% CV.
• All generated sequences produced exactly 100 tokens and zero speculative tokens were accepted.

As an A/B sanity check, exact-occupancy physical fwd_ms versus the clean natural-acceptance run changed by median +0.20% across the 48 common cells (range −1.92% to +1.80%), confirming that disabling acceptance did not reduce verifier-forward latency.

Benchmark Results Before/After this PR

The plot above shows that this PR mostly closes the gap between the per-forward latencies from benchmark vs from serverless traffic. Note that the remaining gap is explainable by the extreme degree of raggedness of serverless traffic--measured per-forward max sequence lengths were still higher for the serverless traffic having average sequence length in this range than the max batch sizes from the simulated ragged traffic from the new benchmark code at high batch sizes. In order to avoid making this benchmark over-tailored to the extreme seen in serverless traffic, k was set at a middle ground of 0.9. This means that the benchmark sends traffic that is close to serverless-level raggedness but definitely more uniform than it, explaining why the full gap is not closed. This is a needed tradeoff to avoid overly-pessimistic estimates for other deployments with less-ragged traffic.

Checklist

• Changes are limited to benchmark workload generation and its input corpus
• No credentials or secrets are included
• New dependency is declared
• Staged diff passed syntax and whitespace checks

Made with Cursor

[cursor]

Stale comment

Risk: medium. Not approving: Cursor Bugbot reported 1 unresolved finding (ragged prompt lengths may exceed the pair's labeled seq_len) and its check completed as skipping. Human review should verify gamma-quantile clamping and benchmark row correctness. Reviewers assigned.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: ragged-batch benchmark generation is a non-trivial behavior change above the low-risk auto-approval threshold. Cursor Bugbot passed (prior seq_len finding resolved); no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: ragged-batch benchmark generation is a substantive behavior change above the low-risk auto-approval threshold. Cursor Bugbot passed with no new findings; no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 7060078. Configure here.}

[cursor]

Stale comment

Risk: medium. Not approving: Cursor Bugbot completed as skipping and reported 1 unresolved finding (ragged warmup request shape mismatch). Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: Cursor Bugbot completed as skipping with 1 unresolved finding (ragged warmup request shape mismatch), and this benchmark behavior change exceeds the low-risk auto-approval threshold. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: Cursor Bugbot completed as skipping with 1 unresolved finding (ragged warmup request shape mismatch), and this benchmark behavior change exceeds the low-risk auto-approval threshold. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: substantive ragged-batch benchmark behavior change exceeds the low-risk auto-approval threshold. Cursor Bugbot passed with prior findings resolved; no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: substantive ragged-batch benchmark behavior change exceeds the low-risk auto-approval threshold. Cursor Bugbot passed with prior findings resolved; no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: medium. Not approving: substantive ragged-batch benchmark generation change exceeds the low-risk auto-approval threshold. Cursor Bugbot passed with prior findings resolved; no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[divchenko]

Can you re-use existing limericks and code datasets instead of adding a new one?

[divchenko]

why do you need this file? can you generate prompts on the fly based on limericks/code datasets?

[janderson4]

The idea of using this set of diverse prompts is to try to get as many experts as possible activated per forward pass. One of the reasons for the lower latency in benchmark compared to prod traffic was that the benchmark used to send a single prompt with n=batch_size, which would lead to far fewer experts being activated on average than in prod traffic. By sending a set of batch_size prompts that are totally distinct from each other, we maximize the breadth of experts that will be activated across the batch, better simulating the average counts of experts activated in prod traffic.

[chualanagit]

i think dmytro is saying why not get a bunch of diverse prompts from the limericks/code dataset. why come up with your own diverse prompts basically

[janderson4]

The issue with using the limericks/code datasets is that they are actually really small: limericks is ~700 lines and code is ~300 lines. We need to send prompts that are diverse in content, and there's no way to construct diverse prompts using just these two short files. Even reordering the same content from these two files will result in prompts that are actually nearly identical, especially because we construct prompts with lengths of 128k and above. The core issue is that we'll be sending the same content in every prompt, and that same content will lead to the same experts being activated.

In order to actually activate a broad set of experts, we need to send prompts that are actually diverse in their content. That's why I made a file with 256 diverse, completely distinct prompts. This way, even at batch size 256 we can be sending 256 prompts that are entirely different from each other and will properly activate a broader set of experts. This is a better simulation of serverless/prod traffic where all prompts are naturally different.

Let me know if this doesn't explain.

[divchenko]

explain more what this gamma is and how it effect produced prompt lens

[janderson4]

This is the gamma distribution which is the probability distribution that most often is the best fit to the prompt lengths we see on prod deployments. It is often used to model waiting times between independent events, but the reason I chose it was just because it was the best fit to the data in most cases (both serverless and customer deployments).

This PR changes the benchmark to model the traffic distribution given any mean prompt length as a gamma distribution with a matching mean. For every cell in the (batch_size, mean_prompt_length) matrix, the benchmark now assumes that the traffic pattern is actually a gamma distribution with mean (k*theta) equal to mean_prompt_length. Making this assumption allows us to reasonably distribute the prompt lengths within each batch size to make them about as ragged as would be seen in prod.

Holding the gamma distribution's mean equal to mean_prompt_length is straightforward, but the trickier part here was choosing the gamma's shape parameter k. k controls how spread out (high variance) vs how narrow (low variance) the distribution is. Since it isn't tractable to sweep over the full (batch_size, mean_prompt_length, k/p90_prompt_length) matrix, we need a reasonable assumption of what that shape parameter k is. By fitting gamma distributions to prod traffic across various deployments, I found that some deployments receive much more ragged traffic than others (especially serverless is more ragged=lower k). I chose a k value somewhere in the middle, keeping it relatively low to be more conservative (lower means more ragged, which slows down forward passes and leads to higher latency estimates, which is more conservative).

So the final method is: assuming a fixed k (shape parameter), for each point in the benchmark we simulate a gamma distribution with the target mean prompt length.

[janderson4]

Following up on the dataset question: I want to use a diverse set of prompts to encourage wide expert activation at large batch sizes rather than reusing the same one or two prompts, which would lead to fewer average experts being activated per layer and slightly underestimate per-forward latency.

[cursor]

Risk: medium. Not approving: substantive ragged-batch benchmark generation change exceeds the low-risk auto-approval threshold. Cursor Bugbot passed with prior findings resolved; no Cursor Security Agent check ran. Human review already requested.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[chualanagit]

i don't think this change is backwards compatible? We would get worse numbers and some workflow would fail? Does shape validation rely on this? Not sure what other stuff relies on this. I think we should have 2 options and allow users to choose ragged versus uniform sequence. Default to uniform for now and perf agent can use the ragged sequence path

[janderson4]

This PR is going to change the generation latency estimates slightly to make them agree more with what we've observed in prod--so yes, it's definitely changing the latency numbers, but that is intended.

Nothing will fail as a result of this change, because there are only 3 workflows that use this code and none of them gate anything on latency regressions:

• Release sweeps: weekly run. This one does call update_mapping.py but it's in dry-run mode so it never updated config_map.json anyway, and this change would not affect its verdict, see explanation below (its verdict is a no-op anyway).
• Nightly sweeps: daily run of kimi
• Single-benchmark workflow: rarely used by anyone, but available for on-demand testing. It's just a single-point run, no comparisons.

Shape validation does not rely on checking regressions in these latency numbers or the trained curves at all (either the validate deployment shape workflow or the sweep to shape workflow). So validation of shapes would be totally unaffected.

The only way the sweep results actually affect CICD in practice is the config_map.json file, which gets manually updated on a weekly basis after a review of the sweep results. Nothing would be blocking the normal manual review and manual update of config_map.json after this change rolls out. There is a regression-check step update_mappings.py that gets run as part of the weekly sweeps, however it's in dry-run mode so whether it passes or fails during the weekly sweeps actually doesn't affect anything. And this workflow only checks two things: whether there was more than a 15% average latency increase across all points checked in a (gpu_count draft_tokens, sequence length, batch size) grid according to the fitted latency model and whether there was over a 50% increase in latency for the estimate at any individual point in that grid. In practice, neither of these things would be reached by this change, because the most extreme difference in latencies I saw across models was for kimi-k2p7 and even then the maximum increase in estimated latency would be only 21% (well below the 50% threshold).

So using ragged batch sizes and diverse prompts will change the generation latency numbers slightly, but there's a good reason for the change and rolling out the change actually won't cause issues in any workflow. I don't think defaulting to uniform batch size with same prompt would be a good idea, because it defeats the purpose of these improvements.