Skip to content

[Bugfix] Graph-capture trace export crashes on DSA dual-graph models - #37780

Closed
fengshaoyi-amd wants to merge 1 commit into
sgl-project:mainfrom
fengshaoyi-amd:fix/graph-capture-trace-naming
Closed

fengshaoyi-amd wants to merge 1 commit into
sgl-project:mainfrom
fengshaoyi-amd:fix/graph-capture-trace-naming

Conversation

@fengshaoyi-amd

@fengshaoyi-amd fengshaoyi-amd commented Sep 3, 2026 •

Copy link
Copy Markdown

Motivation

With SGLANG_GRAPH_BATCH_CAPTURE=1 and --enable-profile-cuda-graph, starting
a server for a DSA model on ROCm dies with an IndexError partway through CUDA
graph capture, before the server ever becomes ready.

_init_profile_context_and_memory_record names each capture trace by walking a
list sized for batch sizes only, advancing one slot per profiler callback:

self._profile_bs_list = list(reversed(self.capture_bs))
self._profile_bs_idx = 0

def on_trace_ready(prof):
    bs = self._profile_bs_list[self._profile_bs_idx]   # IndexError
    ...
    self._profile_bs_idx += 1

That assumes exactly one callback per batch size. But capture() loops over
three dimensions, and each innermost iteration runs its own full profiler cycle
(FullCudaGraphBackend.capture_one steps the wait=2, warmup=0, active=1
schedule three times, so one capture_one_shape produces exactly one
callback):

for bs in capture_range:
    for variant_label, _variant_has_lora in lora_variants:
        for attention_variant in attention_variants:
            self.capture_one_shape(...)

So the callback count is len(capture_bs) * len(lora_variants) * len(attention_variants) while the list holds len(capture_bs) entries. Any
model that turns on an extra dimension overruns the list.

Both extra dimensions are live on main today, for different reasons.

create_attention_graph_variants returns DsaGraphVariants for every HIP host
running one of the architectures in is_deepseek_dsa (DeepSeek-V3/V3.2,
GLM-MoE-DSA, MistralLarge3, Pixtral, LongcatFlash, Dots3Note, ...), whose
capture_labels is ("dense", "sparse") — a dense and a sparse graph per
bucket. It just happens to be off for the models this path was originally
exercised against, which is why it has looked healthy.

The LoRA dimension is a second trigger, and it is not HIP-gated:
UnoDecodeCudaGraphRunner passes record_nolora_graph=forward_width > 1 to
DecodeCudaGraphRunner.__init__, so linear UNO with forward_width > 1
expands lora_variants to [("lora", True), ("nolora", False)] and overruns
the list on any accelerator. It inherits capture() and
_init_profile_context_and_memory_record unchanged, so it hits exactly the
crash below.

The traces written before the crash are mislabeled

The crash is the visible half of the problem. Because the filename comes from
the callback's position in the list rather than from the shape being
captured, every trace after the first carries a bs that does not match the
graph inside it. For 8 buckets with dual-graph on:

Callback Shape captured Filename before Filename after
1 bs=8 dense bs_8 bs_8_dense
2 bs=8 sparse bs_7 (wrong bs) bs_8_sparse
3 bs=7 dense bs_6 (wrong bs) bs_7_dense
4 bs=7 sparse bs_5 (wrong bs) bs_7_sparse
... ... ... ...
8 bs=5 sparse bs_1 (wrong bs) bs_5_sparse
9 bs=4 dense IndexError bs_4_dense

So of the eight files each rank writes before dying, only the first is labeled
correctly. Anything that consumes these traces per batch size — roofline
analysis, kernel attribution — silently analyzes the wrong shape. That is
harder to notice than the crash.

Observed failure

Seen on 8x MI355X, TP8, a GlmMoeDsaForCausalLM checkpoint
(index_topk=2048), with --enable-profile-cuda-graph and
SGLANG_GRAPH_BATCH_CAPTURE=1:

[TP0] Scheduler hit an exception: Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 5048, in run_scheduler_process
    scheduler = Scheduler(
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 518, in __init__
    self.init_model_worker()
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 1006, in init_model_worker
    self.init_all_cuda_graphs()
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 989, in init_all_cuda_graphs
    self.tp_worker.init_cuda_graphs()
  File "/sgl-workspace/sglang/python/sglang/srt/managers/tp_worker.py", line 426, in init_cuda_graphs
    self.model_runner.init_cuda_graphs(
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py", line 985, in init_cuda_graphs
    capture = capture_cuda_graphs(
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py", line 179, in capture_cuda_graphs
    decode = capture_decode_graph(model_runner=model_runner)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py", line 475, in capture_decode_graph
    runner = graph_runners[model_runner.device](model_runner)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py", line 461, in __init__
    self.capture()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py", line 1040, in capture
    self._capture_one_stream()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py", line 1113, in _capture_one_stream
    self.capture_one_shape(
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py", line 1226, in capture_one_shape
    self.backend.capture_one(
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py", line 132, in capture_one
    profiler.step()
  File "/opt/venv/lib/python3.10/site-packages/torch/profiler/profiler.py", line 838, in step
    self._transit_action(prev_action, self.current_action)
  File "/opt/venv/lib/python3.10/site-packages/torch/profiler/profiler.py", line 872, in _transit_action
    action()
  File "/opt/venv/lib/python3.10/site-packages/torch/profiler/profiler.py", line 866, in _trace_ready
    self.on_trace_ready(self)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py", line 781, in on_trace_ready
    bs = self._profile_bs_list[self._profile_bs_idx]
IndexError: list index out of range

followed by Received sigquit from a child process and teardown of the process
tree — the server never reaches ready.

The stack pins the mechanism down without any guesswork. On the v0.5.18 tag
this was captured on (71de97b), decode_cuda_graph_runner.py reads:

1108|                        if dsa_variant is None:
1109|                            self.capture_one_shape(
1110|                                bs, forward, stream_idx, variant_label
1111|                            )
1112|                        else:
1113|                            self.capture_one_shape(
1114|                                bs, forward, stream_idx, variant_label, dsa_variant
1115|                            )

The frame is at line 1113, inside the else, which is only reachable when
dsa_variant is not None. So dual-graph capture was active and dsa_variants
had two entries: 8 buckets x 2 variants = 16 callbacks into an 8-element list.

(#38993 has since generalized dsa_variant / dsa_variants into
attention_variant / attention_variants and moved the label set behind
AttentionGraphVariants.capture_labels. The loop structure, the callback
count, and the overrun are unchanged; this PR is written against the renamed
form, the trace above against v0.5.18.)

Measured against that: capture_bs was [1..8] (8 buckets), 64 sidecar traces
were written (8 buckets x 8 TP ranks) before the failure, and every rank raised
on its 9th callback with _profile_bs_idx == 8. Each rank keeps its own index,
so all of them fail the same way.

Note that the speculative-decoding config in that run (EAGLE, 5 steps) is not
involved: EAGLEDraftCudaGraphRunner is a separate instance with its own
_profile_bs_list, so it cannot advance this runner's index. Consistent with
that, only DecodeCudaGraphRunner traces were produced.

Reproduced in three profiling rounds across two independent sessions, so it is
not intermittent. Environment: sglang v0.5.18 (71de97b), torch
2.9.1+rocm7.2.0, ROCm 7.2.0, Python 3.10.12, 8x MI355X (gfx950).

Reproduction

SGLANG_GRAPH_BATCH_CAPTURE=1 python3 -m sglang.launch_server \
  --model-path <DeepSeek-V3.2 or another is_deepseek_dsa model> \
  --enable-profile-cuda-graph --cuda-graph-max-bs 8 --tp 8

Models with neither variant dimension produce exactly len(capture_bs)
callbacks, which is why this path has looked healthy so far.

Modifications

One file, python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py.

Name the trace after the shape actually being captured instead of deriving it
from a monotonic counter. capture() records the current (bs, lora, attention) combination right before each capture_one_shape, and
on_trace_ready reads
it. There is no index left to overrun, and each trace is labeled with the shape
it actually contains. A sequence-number fallback keeps the callback safe if it
ever fires without a tag set.

Choosing a tag over "make the list the right length" is deliberate: the list
approach would need the variant lists computed twice, in
_init_profile_context_and_memory_record and again in capture(), and the two
would have to stay in agreement as dimensions are added. That is the failure
mode this bug already is.

Backward compatibility

For models with neither variant dimension the emitted names are byte-identical
to before (DecodeCudaGraphRunner_bs_52_rank0.json.gz), so tooling that parses
these filenames is unaffected. Names only gain a suffix in the configurations
that previously crashed:

Model shape Before After
52 bs, no DSA, no LoRA 52 traces, correct names 52 traces, identical names
8 bs, DSA dual-graph IndexError on callback 9, 8 written traces mislabeled 16 traces, bs_8_dense / bs_8_sparse
8 bs, linear UNO (forward_width > 1) IndexError on callback 9 16 traces, bs_8_lora / bs_8_nolora
8 bs, DSA + UNO IndexError on callback 9 32 traces, all unique

DSA models now emit twice as many trace files as there are buckets. That is the
point of the fix — previously the run crashed, so there was no smaller-output
alternative — but it is a visible change in disk usage for anyone using this
flag on a DSA model.

Scope

  • Both SGLANG_GRAPH_BATCH_CAPTURE and SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE
    default to False in environ.py, and the block additionally requires
    --enable-profile-cuda-graph. Serving without profiling never reaches this
    code.
  • No code in the tree parses these trace filenames; the export site is the only
    producer and there is no consumer.
  • _set_profile_capture_tag only reads the env flags and assigns, so it is safe
    on every runner that inherits capture(). The EAGLE / frozen-KV /
    multi-layer draft runners all set record_nolora_graph = False and
    attention_graph_variants = None, so both dimensions collapse to a single
    entry and their filenames stay bs_{bs} exactly as before.
    UnoDecodeCudaGraphRunner is the one inheritor that opts a dimension in, and
    it is the one that needs the fix.
  • PrefillCudaGraphRunner derives from BaseCudaGraphRunner and has no
    graph-capture profiling block, so it is untouched.
  • MultiLayerEagleMultiStepDraftExtendCudaGraphRunner calls
    backend.capture_one() directly, bypassing capture(). It never sets
    _profiler, so it produced no traces before this change and produces none
    after — unchanged behavior.

Accuracy Tests

Not applicable: the change only affects trace filenames on the graph-capture
profiling path. No kernel or model forward code is touched, and the modified
block does not execute unless profiling is explicitly enabled.

Speed Tests and Profiling

Not applicable to serving throughput, for the same reason. The added work is one
env-flag read and a string join per captured shape (tens of calls per server
start), inside a profiling-only path that already runs a full torch profiler
cycle per shape.

Testing

  • The failure above is a real crash from a v0.5.18 deployment, reproduced in
    three profiling rounds across two independent sessions.
  • Replayed the exact callback sequence of capture() for all four shapes in the
    compatibility table and confirmed the new scheme yields one unique, correctly
    attributed filename per callback, and that the no-variant case reproduces the
    old names byte for byte.
  • Confirmed on MI355X that a GlmMoeDsaForCausalLM checkpoint enters the
    dual-graph path: startup logs [dense-decode] DSA dual-graph enabled before
    capture begins.
  • pre-commit run --files python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py passes
    with no modifications.

No unit test is added: exercising this path needs a multi-GPU ROCm host and a
DSA-architecture checkpoint, which CI does not cover. Happy to add one if
reviewers would like the naming logic factored out far enough to test on CPU.

Checklist

  • Format your code according to Format code with pre-commit.
  • Add unit tests according to Run and add unit tests. (see note above)
  • Update documentation according to Write documentations. (no docs affected)
  • Provide accuracy and speed benchmark results. (not applicable, see above)
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #35088934593
Latest PR Test (Extra): ❌ Run #35088934548
Latest PR Test (AMD ROCm 10): ❌ Run #35088934880

The bs-indexed trace exporter assumed one profiler callback per batch
size, but capture() runs one per (bs, lora, attention) combination. DSA
dual-graph models therefore overran the list and crashed partway
through capture, before the server became ready.

The traces written before the crash were mislabeled too: the name came
from the callback's position in the list rather than from the shape
being captured, so all but the first carried a bs that did not match
the graph inside.

Record the combination being captured and name the trace after it.
Models with neither variant dimension keep byte-identical filenames.
@fengshaoyi-amd
fengshaoyi-amd force-pushed the fix/graph-capture-trace-naming branch from bc65b79 to f4fb632 Compare September 16, 2026 11:10
@fengshaoyi-amd

Copy link
Copy Markdown
Author

Closing in favor of #39840, which fixes the same _profile_bs_list overrun the same way (naming each trace after the shape being captured) and is more complete: it also covers the pdmux stream_idx axis and updates the existing unit tests. The stack trace above is an independent repro on v0.5.18 (8x MI355X, TP8, GlmMoeDsaForCausalLM), for reference.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant