[cuda_graph] Size dp-local decode graphs by the runner's own gather requirement

[metamergebot]

Motivation

get_batch_sizes_to_capture filters the decode capture buckets by
get_cuda_graph_batch_size_alignment(), which multiplies the alignment by
attn_tp_size whenever the global require_gathered_buffer() is true.
DecodeCudaGraphRunner already exempts a dp-local forward from the MLP gather
through _forward_is_dp_local (the dense draft that drafts attn-TP-locally
with no cross-DP collective), but the bucket filter never saw that exemption.

So under DP attention a dp-local draft runner keeps only buckets whose
bs * rows_per_request is a multiple of attn_tp_size, pads a one- or
two-request local batch up to that bucket, and runs its attention kernels at a
shape many times the real batch. A small user-supplied bucket list (for
example --cuda-graph-bs 1 2 4) fails startup with
AssertionError: capture_bs=[] for the same reason.

Modifications

• get_cuda_graph_batch_size_alignment and get_batch_sizes_to_capture take
  an optional keyword gathered_buffer_required. Omitting it keeps the global
  require_gathered_buffer() behaviour, so every existing caller is unchanged.
• DecodeCudaGraphRunner.decode_graph_gather_requirements(model_runner)
  (classmethod) owns the (mlp_tp_gather, attn_tp_gather) pair the runner's
  graphs actually need, including the _forward_is_dp_local exemption.
  __init__ derives its instance predicates from it and passes
  require_gathered_buffer into both get_batch_sizes_to_capture calls.
• ModelRunner.decode_graph_gather_requirements() delegates to the runner
  class, so max_decode_logits_rows and BaseRunner._widest_decode_rows size
  the shared logits buffer and the reduction workspace with the same alignment
  the runner captures with, instead of re-deriving a global one.
• get_batch_sizes_to_capture pads num_max_requests with the alignment it
  computed (ceil_align(num_max_requests, mul_base)) rather than the global
  get_cuda_graph_max_batch_size.
• capture_decode_graph reports bs= from the constructed runner's
  capture_bs in the end-of-capture log instead of re-deriving the list with
  global alignment before the runner exists. All in-tree graph runners
  (DecodeCudaGraphRunner and its NPU/XPU subclasses, CPUGraphRunner)
  expose capture_bs.

Behaviour change: a runner whose _forward_is_dp_local is true now captures
the plain decode bucket list (1, 2, 4, ...) and replays the graph sized to its
local batch. Every other runner keeps the global alignment.

Tests:

• test_cuda_graph_setup.py::test_draft_startup_preserves_runner_local_capture_sizes:
  a dp-local draft runner with attn_tp_size=16 and buckets [1, 2, 3, 4, 8]
  keeps all five and the end-of-capture log reports them. On current main the
  same case fails at startup with AssertionError: capture_bs=[].
• test_model_runner_decode_rows.py: the logits-row derivation forwards the
  runner's gather requirement to get_batch_sizes_to_capture.
• Ran with CUDA hidden:
  pytest test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py test/registered/unit/model_executor/test_model_runner_decode_rows.py
  -> 7 passed; the neighbouring graph-runner unit suites under
  test/registered/unit/model_executor/ -> 186 passed, 10 skipped.
• pre-commit run --files <changed files> -> passed.

Accuracy Tests

Graph bucket selection only; no kernel or model-forward change. Not run.

Speed Tests and Profiling

Not run here. The effect is confined to dp-local decode runners, whose
attention now runs at the local batch instead of an attn_tp_size-aligned
padded one.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #37892518588
Latest PR Test (Extra): ❌ Run #37892518374
Latest PR Test (AMD ROCm 10): ❌ Run #37892518684

[metamergebot]

This port clears our bar on 95d8bd63fdd189a19f5ec2f88eb327a55d6449ce: every base-a and base-b job is green.
base-c and the accelerator suites are reported but not gated, because they run
on scarce hardware and upstream merges through them routinely.

gate 21✓✗3✂5 lint✓ build 5✓✗1✂1 A✓ B✓ C 0/10⏳ other 7✓✂7 · accel ✗amd,npu,xpu

@hanming-lu @merrymercy @charlotte12l — could one of you take a look and merge it if you are happy? Nothing else is waiting on our side.