Repository navigation
fix(tbo): decode ubatch metadata follows the micro-batch split; record TBO decode as request graphs - #2692
Merged
Conversation
…batch runs #2688 dropped the TBO decode CUDA graph (`capture_tbo_graph`), so a decode step the DP group splits into micro-batches now always runs eagerly. Its micro-batches are cut from the scheduled requests (`maybe_create_ubatch_slices`: 25 scheduled -> [0, 12) / [12, 25)), but the attention backends still laid each ubatch's metadata out for half the padded batch, the split the removed graph had baked in (32 running -> [0, 16) / [16, 32)). Micro-batch 1's rows then attended with, and wrote KV slots for, other requests' context. gpt-oss-120b, TP2 + DP attention + EP + `--enable-tbo all`, GSM8K 3-shot flexible: 0.8901 (CI, 2c772a5) -> 0.5951 (CI, 9a0a520). Locally on the CI image: 2c772a5 0.8992, 9a0a520 0.6080; on 9a0a520, `--enable-tbo prefill` 0.8901 and TBO off 0.8931, while disabling prefill token graphs (0.5929) and `--enforce-eager` (0.7619) do not recover it. - `decode_ubatch_request_ranges` is the one decode split: the slices (`maybe_create_ubatch_slices`), the per-ubatch token counts the DP sync agrees on (`_precompute_decode`) and every backend's per-ubatch metadata (`_prepare_ubatch_decode` in aiter_attention, aiter_mla and deepseek_v4_attn) all take it. - Each backend keeps its padded per-ubatch width, which always covers the ubatch's requests; only the request ranges change. DeepSeek-V4's non-eager branch cut ranges at the DP-agreed width and is aligned too. - aiter_mla: `build_ubatch_metadata` sliced `dcp_local_context_lens` from half the padded batch as well; it now reads the ubatch's own start (`_tbo_full_running_bs` -> `_tbo_ubatch_request_ranges`). - tbo/__init__: sort the imports and `__all__` the new export touches. Tests: tests/test_tbo_decode_ubatch_split.py (CPU) checks slices, DP sync and backend widths agree; the H2D GPU tests that hard-coded the padded split now expect the scheduled one and check each ubatch's context_lens; the MLA dcp_local_context_lens test covers 3 scheduled requests. Validation (rocm/atom-dev:nightly_202610231620, MI355X): - gpt-oss-120b TP2 DP-attn EP `--enable-tbo all`, GSM8K 3-shot flexible 0.8825 / strict 0.1766 (threshold 0.87) - black, ruff clean on the changed files; CPU + RUN_H2D_GPU_TESTS=1 attention tests: 1217 passed. tests/test_h2d_v4_{,indexer_}publication.py fail the same 28 cases on main and on this branch (stale fakes, e.g. a config without max_num_seqs), so DeepSeek-V4's non-eager branch has no GPU test coverage here. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
#2688 dropped `capture_tbo_graph` and left every decode the DP group splits into micro-batches to run eagerly. Record them on `ReqGraph` instead: with TBO decode on, a bucket that splits (`decode_bucket_splits`: running_bs > 2) is recorded split at the bucket, and a decode step replays it. - `ModelRunner.capture_cudagraph`: such a bucket's capture context carries its micro-batch slices (`decode_ubatch_request_ranges(bs)`) and per-rank rows -- every rank records the same bucket, so each micro-batch's rows are its own, repeated -- and `req_graph.record` records the UBatchWrapper forward as is: its persistent workers issue onto streams forked from the capture stream and `_run_ubatches` joins the comm stream before returning. TR_HYBRID has no split token graph; there a split decode stays eager. - `ForwardMode.decide(tbo_decode_graphs=...)`: a split decode no longer forces eager when the buckets were recorded split; prefill TBO still does. `ForwardMode.ubatch_split_bs` is the batch the step's micro-batches are cut from: the bucket on a replay (whatever each rank scheduled into it), the scheduled requests on an eager split. - Attention builders take it through `build(ubatch_split_bs=)` and every `_prepare_ubatch_decode` cuts its per-ubatch metadata at it (`split_bs`). DeepSeek-V4's read its widths from the forward context, which at build time is the previous forward's; it now takes the split explicitly, and its token offsets are indexed from the scheduled requests a padded ubatch can start past. - `decode_bucket_splits` is the one bs > 2 gate (`_precompute_decode`, the recording, V4's prepare and capture). - docs/compilation_cudagraph_guide.md: how TBO decode records and replays. Validation, gpt-oss-120b TP2 + DP attention + EP (rocm/atom-dev nightly_202610231620, MI355X): - `--enable-tbo all`: every decode bucket records; GSM8K 3-shot flexible 0.8984 / strict 0.1804 (2c772a5: 0.8992). - Replay vs 2c772a5's `capture_tbo_graph`, ISL/OSL 1024, mean TPOT / output tok/s: conc 16 8.85 ms / 1711 vs 8.71 / 1563; conc 64 11.59 / 4999 vs 14.57 / 3940; conc 128 14.43 / 8028 vs 14.35 / 7875. The eager split this replaces did not finish its conc-16 point in 30 minutes. - On this model TBO decode is slower than none with either recording: `--enable-tbo prefill` gives 5.78 / 8.09 / 10.33 ms TPOT (2c772a5 5.49 / 8.02 / 10.25). - tests/test_forward_mode.py: which decodes replay split and what they cut. CPU tests 1364 passed; RUN_H2D_GPU_TESTS=1 attention tests pass, and the V4 H2D tests fail the same 28 cases as main. black, ruff clean. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Add `GPU host health check` as the first step of the atom-test accuracy job on the DO and predownload runners, before any Docker cleanup or image pull. It fails the job, with an `Unhealthy GPU host` annotation, when: - the host has no rocm-smi. Runners registered under linux-atom-do-mi350x-8 without ROCm tools on the host (spur-*, crsuse2-m2m-*) failed 51 jobs at `rocm-smi --showmemuse`, 2026-10-23; - rocm-smi fails, times out (60 s), or prints no GPU rows; - any GPU reports N/A temperature or power. Of the failed accuracy jobs since 2026-10-21, all 30 that hung at server start with `HW Exception ... GPU Hang` (about 10 minutes in) showed this in the rocm-smi table the job printed before starting, on five hosts. None of the 83 other failures and none of 80 successful jobs (32 hosts) did. The GPUs that hung are exactly the ones reporting no power; one host had a single such GPU, and four hosts reported edge temperature N/A on every GPU with one or two lacking power. At least 9 of 16 GPU preflight timeouts started within 11 s of a hang ending, on a runner that replaced the hung one, matching a wedged GPU on the same host. For the GPUs reporting no power it prints `rocm-smi --showuniqueid`, which is rocminfo's GPU UUID, so the host can be found and drained. tests/test_gpu_host_health_check.py runs the step through a stand-in rocm-smi on the three tables CI hosts printed (healthy, one dead GPU, edge-only host with two dead GPUs) and on a missing, failing and row-less rocm-smi. Replaying the step on the rocm-smi tables of 121 historical CI jobs: all 30 that hung fail it, all 91 others pass. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Jasen2201
added a commit
that referenced
this pull request
Oct 30, 2026
…d publication test - KimiK3Parser.required_call_channel is cached per tokenizer: the channel is fixed for a server, so required requests no longer re-encode it. - test_prepare_model_publishes_sampling_and_query_prefix_before_token_consumer: required_channel_states is optional, so a step without a channel leaves it on an older epoch. The packed check now names exactly which members stay unpublished (decode_src outside deferred steps, the channel states without a channel), and the steps alternate no channel / channel / no channel and check the staged states. - The gated runner tests' hand-built runner sets tbo_decode_graphs, which prepare_model has read since #2692; without it all 7 failed on main too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
10 of 11 tasks
yitingw1
added a commit
that referenced
this pull request
Oct 30, 2026
The file failed 20 of 22 cases on main with RUN_H2D_GPU_TESTS=1, all from fixtures that fell behind the code: - set_pcp patched pcp_round_robin_query_indices (removed in #2688); the builder reads get_pcp_rank. - make_runner lacked forward_vars["batch_info"], config.max_num_seqs and the builder's csa_layers, max_model_len_idx, _local_heads and index_logits_workspace. - The padded decode case modelled the bucket width through UBatchWrapper._decode_ub_running_bs; since #2692 the split comes from _prepare_ubatch_decode's split_bs, so the test passes the bucket. - prepare_prefill was called with a total_tokens_num the fake batch lacks. 22/22 pass on this branch and on its merge base. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 task
valarLip
pushed a commit
that referenced
this pull request
Oct 31, 2026
… GPU (#2849) * fix(kimi-k3): hold tool_choice="required" to the tools channel on the GPU Kimi-K3 ends its response channel with `<|close|>response<|sep|>` and then either opens the tools channel or closes the message. ATOM only told it "You MUST call tools" in the template, so under `tool_choice="required"` it still closed the message at that token on a sampled share of requests and wrote the arguments as its answer: about 90% of the failures on the Kimi-Vendor-Verifier tool-call suite, where vLLM -- which holds the output to a K3 structural tag -- has none of that kind. A request now carries a `RequiredChannel` (SamplingParams.required_channel) when it says `tool_choice="required"`, has tools, and the server's tool-call format defines one. Only Kimi-K3 does: once the output ends with `<|close|>response<|sep|>` the next tokens are `<|open|>tools<|sep|><|open|>call tool="`, and until then `<|end_of_msg|>` and EOS cannot be sampled. The tool name and arguments stay the model's. Every other request, and every other format, samples exactly as before. The constraint acts on the target logits on the GPU, before sampling, because nothing can be changed after: the next forward has already taken the sampled ids as input. Each logits row is constrained by the state its sequence has reached by that row: - a decode or prefill row by the state after its anchor; - a verification row as if the drafts before it were accepted, so a draft that breaks the channel is rejected and greedy verification compares against the masked argmax; the bonus row comes after every draft. The host advances each sequence's state over the tokens the scheduler commits (and over a P/D consumer's T0 at admission) and stages it with the other sampling parameters. Under deferred output the host is a forward behind, so the runner also folds each step's emitted ids into the states on the GPU and reads them back through `decode_src` the next step, like the anchor itself. Both kernels run only when the batch holds a request with a channel. Tests: the state machine, the K3-only gating over every registered format, the scheduler's state across commits and P/D admission (CPU), and the masks, the draft walk, the carry and the fold against pure-torch references, with the Triton kernels held to them on a GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(kimi-k3): cache the required channel per tokenizer; fix the packed publication test - KimiK3Parser.required_call_channel is cached per tokenizer: the channel is fixed for a server, so required requests no longer re-encode it. - test_prepare_model_publishes_sampling_and_query_prefix_before_token_consumer: required_channel_states is optional, so a step without a channel leaves it on an older epoch. The packed check now names exactly which members stay unpublished (decode_src outside deferred steps, the channel states without a channel), and the steps alternate no channel / channel / no channel and check the staged states. - The gated runner tests' hand-built runner sets tbo_decode_graphs, which prepare_model has read since #2692; without it all 7 failed on main too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(kimi-k3): cover the chat endpoints' channel wiring; fix two postprocess test doubles - test_required_channel: the openai_server chat endpoint and both AToMesh standalone chat entrypoints (non-stream, stream) hand a required K3 request its RequiredChannel and an `auto` one none. Dropping any one site's `required_channel=` fails exactly its case. - test_postprocess_width: the prepare_sampled_ids double accepts the new `channel_states` keyword. - test_output_transfer_runner: the stub runner binds the real `_constrain_required_channel`, which postprocess now calls. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(kimi-k3): hold back every engine stop id; report pre-mask logprobs - RequiredChannel.ending_on: InputOutputProcessor.preprocess_fanout adds the engine's eos_token_id and stop_token_ids, the ids Scheduler.postprocess ends a sequence on, to the channel's end_ids. A stop id the format did not name could otherwise end the message before the channel opens. The K3 parser now names only its own <|end_of_msg|>. - postprocess: the no-spec path masks the logits in place, so when a batch with a channel asks for logprobs, gather them from a pre-mask copy. A forced token reports its real logprob, as the verification path (which masks copies) and unconstrained requests do, instead of 0.0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(kimi-k3): make the P/D T0 test catch a double count; clarify two comments - test_pd_first_token_counts_once: T0 now completes `after`, so a T0 re-fed by postprocess lands one state past the real one. With the old toy channel a double count also ended at 2, so the test could not tell. - llm_engine: name all three ways a sequence ends; the request's own stop strings are deliberately left to end it. - postprocess: say the mask below is in place, so the logprob copy comes before it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(kimi-k3): an opened channel is UNCONSTRAINED; derive the channel at startup - RequiredChannel.step: writing the last forced id returns UNCONSTRAINED instead of parking at open_state. ScheduledBatch only carries a channel (and only builds the per-row states) while some row has one left to open, so a batch whose requests have all opened theirs stages nothing and launches no channel kernel. - Kernels and references treat an UNCONSTRAINED row as unmasked mid-walk too. A forced id whose own logit is -inf is written as 0, so the row keeps a finite logit instead of a NaN softmax. - The channel is derived once beside the tool-call parser at startup (registry.required_call_channel); each request only gates it (channel_for_request). This drops the per-tokenizer functools.cache. - RapidServe prefill: a comment on why it masks without a pre-mask copy. - Tests: kernels vs references on K3's real ids and vocabulary (> BLOCK_V) without a carry, the -inf forced id, opened rows and all-opened batches. - Docs: model_ops_guide (required_channel.py) and serving_benchmarking_guide (tool_choice="required" on Kimi-K3). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(kimi-k3): open the channel on its own start; spread forced rows over programs Addresses valarLip's review of a764950 and valar-cop's two notes. - RequiredChannel.opens_at: opening[:opens_at] (K3: <|open|>tools<|sep|>) written by the model itself opens the channel, so a request that reaches the tools channel another way is no longer held to max_tokens. `step` follows `after` and that start at once; construction refuses shapes the restart-at-1-or-0 rule cannot follow, and end ids that are channel ids. - constrain_logits runs VOCAB_PROGRAMS (16) programs per request, each writing every 16th 1,024-column block of a forced row without reading it. MI355X, K3 vocabulary: one forced row 58 -> 4 us (bs 1), a DSpark7 verify step with seven forced rows 400 -> 15 us (bs 1) and 482 -> 52 us (bs 64); a step with no forced row stays at 5-7 us. - InputOutputProcessor builds the engine's channel once and shares it, and refuses a request naming a different channel at admission; ScheduledBatch only asserts that invariant. - Logprobs under a channel: a [bs] logsumexp before the mask, not a copy. - tokenIDProcessor records carries_over; the runner reads decode_src itself. - Smaller: a startup warning when K3's control tokens split, the named tool_choice noted as not covered, AToMesh derives the channel in one place, plain batch.required_channel reads, UNCONSTRAINED as a kernel constexpr. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(kimi-k3): check the opening's head before counting opens_at from it - KimiK3Parser.required_call_channel: opens_at counts the ids of <|open|>tools<|sep|> at the head of the forced opening, which holds only if the tokenizer encodes that head as it encodes TOOLS_START alone. A tokenizer that does not now constrains nothing and warns at startup, like one that splits a control token (both warnings are tested). - required_channel: state why the end-id ban needs no all -inf guard. It leaves every column but the end ids as the LM head wrote them, and a guard would read the whole row on the step a request takes until its channel opens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
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.
Problem
After #2688, the
gpt-oss-120b (2 GPUs)accuracy job on main (-tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all) dropped from 0.8901 to 0.5951 GSM8K (run 39986960948). There is no error and no crash; the outputs are just wrong.#2688 removed the TBO decode CUDA graph (
capture_tbo_graph), so every decode step that the DP group splits into micro-batches now runs eagerly. In that eager path, the micro-batches are cut from the scheduled requests (maybe_create_ubatch_slices: 25 →[0,12)/[12,25)). The attention backends, however, still lay out each ubatch's metadata over half of the padded batch (32 →[0,16)/[16,32)). That was the split the removed graph had baked in. As a result, micro-batch 1's rows attend with another request's context, and they write KV slots that belong to other requests.Bisected on the CI image (gpt-oss TP2 DP-attn EP, GSM8K 3-shot flexible):
--enforce-eager--enable-tbo prefillI also instrumented
UBatchWrapperto check whether the per-ubatch DP sizes match the slices. They did on every step (40/40), so the bug is in the per-ubatch attention metadata, not in the DP sizes.Changes
1.
fix(tbo): lay decode ubatch metadata out for the requests each micro-batch runsdecode_ubatch_request_rangesis now the single definition of the decode split. The slices, the DP-synced per-ubatch token counts (_precompute_decode), and every backend's_prepare_ubatch_decode(aiter_attention, aiter_mla, deepseek_v4_attn) all use it.build_ubatch_metadataalso sliceddcp_local_context_lensat half the padded batch. It now uses the ubatch's actual start.2.
feat(tbo): record TBO decode buckets as request graphsdecode_bucket_splits:running_bs > 2) is recorded onReqGraphwith the bucket already split, and a split decode step replays it. Nothing new was needed in the recording mechanism:UBatchWrapper's persistent workers issue work onto streams forked from the capture stream, and_run_ubatchesjoins the comm stream back before it returns.ForwardMode.ubatch_split_bsrecords which batch the step's micro-batches are cut from: the bucket on a replay, or the scheduled requests on an eager split. Builders receive it throughbuild(ubatch_split_bs=)._prepare_ubatch_decodeused to read its ubatch widths from the forward context. At build time that context still belongs to the previous forward, so the widths could be stale. It now receives the split explicitly, and it no longer indexes token offsets past the scheduled requests.Validation
gpt-oss-120b TP2 + DP-attn + EP +
--enable-tbo all, on CI imagerocm/atom-dev:nightly_202610231620, MI355X:--enable-tbo all:capture_tbo_graph--enable-tbo prefillReplay is on par with the old TBO graph (better at conc 64). On this model, TBO decode is slower than no decode TBO with either recording, so this PR does not change that trade-off. That question is separate from this fix.
tests/test_tbo_decode_ubatch_split.py(CPU) checks that the slices, the DP sync, and the backend widths all agree.tests/test_forward_mode.pycovers which decodes replay split, and which batch they are cut from.context_lens.RUN_H2D_GPU_TESTS=1attention tests pass. black and ruff are clean.tests/test_h2d_v4_{,indexer_}publication.pyfail the same 28 cases on main and on this branch. Their fakes are stale (e.g. a config withoutmax_num_seqs), so DeepSeek-V4's ubatch path has no GPU test coverage here.🤖 Generated with Claude Code