[4/N] extend CUDA-specific graph API to pytorch's accelerator graph backend, to make graph supported on Intel GPU (XPU) by yao-matrix · Pull Request #262 · mstar-project/mstar

yao-matrix · 2026-09-16T21:25:51Z

Summary

Replace CUDA-specific graph plumbing with an accelerator graph backend for CUDA and Intel GPU.
Migrate existing model integrations, graph lifecycle hooks, and tests to the accelerator graph API.
Use device-specific pinned staging, events, output copies, and graph-stable XPU paged-attention plans. Upload captured XPU block tables, query indptrs, and KV lengths directly from pinned host staging into their existing per-slot buffers.
Dispatch existing Chatterbox S3Gen captures through the accelerator graph backend.
Use accelerator_graph deployment switches and environment names in BAGEL, Waypoint, and Cosmos3.
Keep BAGEL timestep sinusoidal frequencies in FP32 across module dtype conversions.
Preserve Pi05's FP32 vision encoder and connector with the submodule's disable_autocast policy on every device. Add eager single/batched regressions and include the Pi05 CPU test module in test/cpu-core.txt.
Use mixed CPU/XPU process groups (cpu:gloo,xpu:xccl) only on XPU. Its TP/SP subgroups inherit both backends and share the handle for CPU object and XPU device collectives. CUDA uses the original NCCL world/device groups and separate Gloo object-collective subgroups, including the original local-synchronization option.
Migrate Nemotron Duplex NanoLM, EarTTS talker, and LTX-2.5 capture recipes and tests to the accelerator graph API.
Preserve BAGEL's default CUDA graph captures; its XPU YAMLs explicitly set model_kwargs.accelerator_graph: false.
Use AcceleratorGraphPool throughout capture and recovery, including tests. Its lifetime hooks retain failed graphs only in the XPU pool; CUDA hooks are no-ops.
Retain CGSlotKey, CGSlotSpec, and cg_ names with CG meaning "captured graph" across accelerator backends; the definition is in the stacked comment PR.
Share a cached DeviceCapability for availability, pinned-memory, and graph-support queries from the worker through the engine and runners; staging utilities reuse the same cache. Streams, events, and synchronization use PyTorch's generic APIs directly. Model loading rebinds a capability whose device differs from the loaded device, and preserves a matching supplied handle.
Let resources report can_capture(); whole-forward and piecewise runners reject unsupported recipes before buffer allocation and preserve MSTAR_REQUIRE_ACCELERATOR_GRAPHS=1. Qwen 3.5 recipes no longer branch on device type.
Keep environment-variable documentation matching the renamed controls, essential explanations of new abstractions, and repository skill guidance with the implementation. Update AGENTS.md and the model, resource, diffusion, and benchmarking skills to reference the accelerator graph APIs, config classes, runner paths, and capture tests. Broader documentation and comment wording updates remain in the stacked PR: [5/N] docs: accelerator graph documentation and terminology #475.

Verification

CPU Core portion of the combined suite at ef538648, with PyTorch 2.9.1+cpu and OMP_NUM_THREADS=2: 2,574 passed, 482 skipped, no failures.
Four edited modules outside test/cpu-core.txt were also run. With the documented qwen-tts==0.1.1 package installed, the follow-up selection covering them and CPU ragged-attention tests returned 90 passed, 57 skipped, 3,039 deselected. CUDA ragged-attention/GDN kernel tests were skipped.
Native graph/resource tests on Intel Arc Pro B60 with PyTorch 2.13.0+xpu: 39 passed, 7 CUDA-only skips. This includes uneven packed-query replay and graph-pool recovery.
Three-XPU mixed-group test: 1 passed. Production TCP rendezvous includes a non-member worker; CPU object gathering/reduction and eager XPU all-reduce/all-gather use the same mixed subgroup.
Whole-repository ruff check ., PR diff whitespace checks, Python parsing, and guidance-link checks pass.

The same combined command completed on stacked documentation PR #475 (a4244f8f): 2,647 passed, 539 skipped. That runtime matched #262 at ef538648 across all 107 differing Python files. The current stack (7c35e6e6) matches current #262 (b954445b) across all 111 differing Python files; YAML values and shell commands also match.

Graph API consistency verification at ab97c1d2: test_step_prep_host_syncs.py, test_cosmos3_graph_switch.py, and test_orpheus_snac_stream.py returned 75 passed. The Cosmos3 fixture verifies that its environment settings disable prefill capture while preserving image/video capture; a negative control using the old environment names fails this check. An executable-code audit finds no legacy graph identifiers or CUDA-named graph environment controls, and no raw capture calls under model packages. Whole-repository Ruff and whitespace checks pass.

Focused device-binding and XPU staging verification at b954445b:

CPU graph runtime, device/resource capability, XPU attention planning, resource lifecycle, and shared piecewise pool tests: 83 passed, 2 XPU hardware skips, 3,001 deselected on PyTorch 2.9.1+cpu. Nine capability-binding cases cover early queries and matching/mismatching supplied handles on CPU/CUDA/XPU. The metadata regression rejects temporary tensor allocations during replanning and checks values, padding, and plan promotion; the unfixed revision fails five selected cases.
Native test_xpu_attention.py on Intel Arc Pro B60 with PyTorch 2.13.0+xpu: 18 passed. The new test queues 20 replays on alternating slots with a dedicated preplan stream and reuses the pinned host ring; all replayed metadata matches the expected layouts. Uneven-query attention replay remains bitwise equal to eager attention.
Isolated upload profile, five warmed preplans with one label, batch size 8, and 1,024 page columns on B60: each preplan changes from 6 memory copies / 3 zeEventHostSynchronize calls / 3 temporary device conversions to 3 copies / 0 upload waits / 0 temporary device conversions. PinnedStager still waits when recycling an in-flight host source. This is a staging measurement, with no end-to-end speedup claim.
Whole-repository Ruff and whitespace checks pass. CUDA hardware was not available for this verification.

Focused CPU command (same environment as the full suite below):

python -m pytest -q -ra --strict-markers @test/cpu-core.txt \
  -k 'xpu_attention or accelerator_graph_runtime or device_capability or resource_graph_capability or resource_runner or piecewise_shared_pool'

Native command (PyTorch XPU build; import real Triton before pytest on this host): python -m pytest -q -ra --strict-markers test/modular/test_xpu_attention.py, with OMP_NUM_THREADS=2 MKL_NUM_THREADS=2.

CPU suite commands:

OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 HF_HUB_OFFLINE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
python -m pytest -q -ra --strict-markers --durations=10 --basetemp=/tmp/p262final \
  @test/cpu-core.txt test/modular/test_qwen3omni_vision_mixed.py \
  test/modular/test_qwen3_tts_model.py test/modular/test_gdn_mixed_split.py \
  test/modular/test_ragged_attention.py
# Follow-up selection from the same collection:
# -k 'qwen3_tts_model or qwen3omni_vision_mixed or gdn_mixed_split or ragged_attention'

Native XPU tests: test_xpu_attention.py, test_accelerator_graph_backend.py, test_resource_graph_capability.py, and test_xpu_parallel_groups.py.

CUDA hardware and Rust-runtime tests were not executed locally; the CPU run skips GPU-, checkpoint-, and Rust-dependent cases. The stacked documentation/comment PR remains #475.

AI assistance was used; tests and results are listed above.

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

yao-matrix · 2026-10-02T20:56:56Z

@merceod @NSagan271 in this PR, I want to abstract out the AcceleratorGraph to make graph mode work both on CUDA and other devices which support graph mode(e.g. Intel GPU). It's a breaking change and about device choice. So, before I enhance the PR, I want to knock your door and get your comments on this. Could you let me what's your preferred design to expand graph mode to other devices? I'll follow your design on this. Thx very much.

Preserve new Qwen3-TTS packed prefill capture, 1.7B/voice clone support, and streaming codec optimizations using the accelerator graph API. Update the newly added diffusion capture hook and its tests to that API. Resolve documentation and Qwen3-TTS conflicts without legacy aliases.

yao-matrix · 2026-10-05T21:46:01Z

@NSagan271 @merceod gentle ping. It's maybe too early to get a thorough review on this PR. I just want to align to your design direction. For Graph support in mstar, do you prefer to abstract out an AcceleratorGraph as PyTorch does (https://docs.pytorch.org/docs/2.14/generated/torch.accelerator.Graph.html) or you prefer to still go with CUDAGraph but redirect it to XPUGraph as vLLM does(https://github.com/vllm-project/vllm/blob/2e3154aa18b3d2df2253c3e55f8b32fe9af5cfe8/vllm/v1/worker/xpu_model_runner.py#L56)? I will follow up on the PR per your design decision, thx.

NSagan271 · 2026-10-05T22:45:05Z

@yao-matrix Sorry for the delay! I'd generally prefer an AcceleratorGraph-type abstraction; I believe that would be more generalizable to different backends instead of being tied to the CUDAGraph API. I generally like the abstraction you made in AcceleratorGraphBackend, but I'd prefer an abstract base class with CUDA and XPU versions overriding it, in order to avoid adding more if statements in the graph backend if we add support for more hardware types beyond GPU and XPU.

Sync with upstream main 166ddb0, including per-batch loop metadata and threaded speech request mapping. Update remaining generic graph wording in resources, models, documentation, benchmarks, and tests. Adapt speculative loop counter tests to the metadata API and run them in CPU CI, preserving in-flight counter snapshots after shared state changes.

yao-matrix · 2026-10-06T17:52:55Z

@yao-matrix Sorry for the delay! I'd generally prefer an AcceleratorGraph-type abstraction; I believe that would be more generalizable to different backends instead of being tied to the CUDAGraph API. I generally like the abstraction you made in AcceleratorGraphBackend, but I'd prefer an abstract base class with CUDA and XPU versions overriding it, in order to avoid adding more if statements in the graph backend if we add support for more hardware types beyond GPU and XPU.

@NSagan271 @merceod , glad to know that I am aligned to your design choice, pls help review the PR, thx.

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

Preserve batched model outputs and size-class pinned buffers with device-specific CUDA/XPU streams. Adapt the new Qwen3.5 submodules to accelerator graph hooks and update output-contract regression coverage. Signed-off-by: Matrix Yao <matrix.yao@intel.com>

#262 renamed use_cuda_graph to accelerator_graph on the FlashInfer decode and prefill wrappers; FlashInferMixedWrapper still passed the old name, so every mixed[...] capture raised TypeError and those steps ran eager. Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com>

…lves, instead of a hand-kept list (#581) * CI: delete test/cpu-core.txt and run test/modular directly; broken tests skip in-file Every PR that added a CPU test appended to test/cpu-core.txt, so each merge put other open PRs into conflict, and modules not on the list never ran. The CPU Core job and the Mac cpu-smoke check now run test/modular plus the cosmos3 windowed-decode test. Tests that need flashinfer, qwen-tts, mstar_rust or soundfile skip without them, and the kokoro oracle tolerates a missing misaki. Tests broken on main skip in their own files with the reason: the bit-exact check. The FlashInfer guard in test_kv_page_ownership.py matches AGENTS.md and the writing-prs skill describe the skip rule. * test_nemotron_duplex_talker_pool: import its sibling module directly 'from test.modular...' resolves to CPython's own test package on interpreters that ship it (CI's does), so collection failed once CI ran all of test/modular. * test_z_image_native_equivalence: skip only the padding-independence check Split it out of the caption-encoder test so the HF hidden_states[-2] parity asserts keep running on CPU. * test_openai_transcriptions: skip long-form tests only when no decoder can read a duration Probe audio_duration_seconds on a short WAV instead of checking for soundfile, so a torchcodec-only install runs them too. * test_attn_mixed_split: skip the prompt-row cap check without CUDA * CI: run test/distributed/test_fanout_shapes.py with the CPU tests * CI: select every test directory, model-local tests by glob, and guard the selection CPU_TEST_PATHS now names directories (test/modular, test/distributed, test/chatterbox, test/kokoro, test/omnivoice, mstar/model/*/tests, benchmark/worker_phases) instead of single files, so mstar/model/cosmos3/tests/test_dit_step_guard.py runs again. Every other test module is listed in NOT_CPU_CORE with where it runs instead. test/ci/test_cpu_selection.py fails when a module the deleted test/cpu-core.txt ran (frozen in the test) is no longer selected, and when a test module is neither selected nor excluded with a reason. CPU Core runs on 4 pytest-xdist workers with worksteal, so the 100 s TP lockstep simulations no longer add up serially. * test_openai_router: skip the omnivoice clip cases unless torchcodec can decode Newer torchcodec imports without FFmpeg and fails only on the first decode, so the import-only guard let the two 61 s clip cases run and fail. * test_z_image_native_equivalence: check padding independence at one shape, and across lengths on CUDA A batch without padding takes the is_causal SDPA path and a padded one the masked path, so comparing them bit for bit compared two kernels. The same-shape check (padding contents) runs everywhere; the shorter padded batch is compared on CUDA, where it is bit-exact. * Reference-parity tests: skip on a submodule only the reference package has mstar/model is a namespace package, so pytest's prepend import of mstar/model/*/tests puts mstar/model on sys.path, and then importorskip("omnivoice"), ("chatterbox") and ("kokoro") find mstar's own model packages. The in-test omnivoice guards ran after that and failed in CPU Core; the module-level chatterbox and kokoro guards passed only because their directories come first on the command line. * CI: run the TP lockstep simulation in its own CPU Core shard Its three simulations take 1.5 to 3.5 minutes each on the GitHub runner and held the 4-worker run's critical path (CPU Core step 235 to 380 s over four runs). A matrix runs them beside the rest with the same install; the main shard keeps the check name CPU Core. test/ci checks the main shard ignores the file and the second shard runs it. * test/omnivoice: rename test_parity.py to test_omnivoice_parity.py; guard unique basenames test/kokoro/test_parity.py has the same basename, neither directory is a package, and CPU Core now collects both: with the kokoro reference importable the run stopped at collection with "import file mismatch". test/ci fails on two selected modules outside a package with one basename. CONTRIBUTING.md and AGENTS.md say why reference guards import a submodule. * CI: one OpenMP thread per CPU Core worker On a host held to 4 cores the command took about 1,550 s with default threads (the cosmos3 windowed-decode tests stall under 4 workers x 4 threads) and about 210 s with OMP_NUM_THREADS=1 and MKL_NUM_THREADS=1. * CI: run test/sampling_test with the CPU tests main's list now runs test_penalty_gate.py on CPU, and the XPU sampler tests skip without an XPU, so the directory joins CPU_TEST_PATHS instead of NOT_CPU_CORE. flashinfer_batch_test.py, a GPU script with one pytest-collected check, skips without CUDA or flashinfer. * test_kv_prefix_root: give the hand-built Engine its device _compile_submodules now asks the engine's device capability whether an accelerator is available (#262), which reads Engine._device; the stub built with Engine.__new__ had none and failed with AttributeError. main's list never ran this module, so main's CI did not catch it. --------- Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…gave get_cuda_graph_configs

#262 renamed use_cuda_graph to accelerator_graph on the FlashInfer decode and prefill wrappers; FlashInferMixedWrapper still passed the old name, so every mixed[...] capture raised TypeError and those steps ran eager. (cherry picked from commit e7b3a08)

Conflicts, each resolved to main's version plus this PR's own change: - whisper_model: main's #518 sized the decoder caches and checks max_concurrent_requests against them at startup; kept main's _cache_pages / validate_config_yaml, plus this PR's two extra self-attention pages per request for the decode captures' padding rows, and both sides' new methods (load_audio, prompt_tokens; max_input_files). - #262 renamed the CUDA-graph API (AcceleratorGraphConfig, get_accelerator_graph_configs); the Whisper and Qwen3-ASR submodules and their tests use the new names. - worker: main moved store-and-route into _store_and_route; the host-copy drop after this step's sends (this PR's SHM fix) now reads stored.flat_uuids. pre_plan from the serve yaml is kept beside main's launch_wait_s and device_capability. - engine / engine_manager: enable_pre_plan beside main's capture_runner, gpu_memory_fraction and device_capability. - test/cpu-core.txt: deleted, as main did (#581). - configs, cli, docs: both sides' model lists. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Takes main's deletion of test/cpu-core.txt (mstar-project#581 runs test/modular on CPU), and moves the new tests to mstar-project#262's accelerator_graph names. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… graph config

…gave get_cuda_graph_configs

* qwen3-omni: start streamed audio after 4 Code2Wav frames instead of 25 The first audio waited for a 25-frame (2 s) Code2Wav chunk. The Talker -> Code2Wav edge now uses ScheduledLeftContextChunkPolicy with codec_chunk_schedule=(4, 8, 13), then 25-frame chunks behind 25 frames of left context. The ramp sums to the left context, so every ramp window holds the whole stream so far and every later window is the one 25-frame first chunks pop: the streamed audio does not change. An empty schedule restores 25-frame first chunks. The trim reads each window's context from the stream metadata instead of assuming 0 for the first window and codec_left_context_frames after it, and every window still pads to one shape (the schedule's widest window), so the vocoder keeps a single CUDA graph. The schedule, chunk and left context can be set from the deployment YAML's model_kwargs. Values that would drop audio are rejected at startup with an error naming the key: a left context under 1 frame (a window decodes 555 samples short and only a context frame re-decodes them), non-positive or non-integer chunk sizes. * test: Code2Wav first-chunk ramp stitches, batches and keeps boundary quality Streamed length and samples equal the full decode for several ramps, reply lengths (shorter than the first chunk included) and every accepted knob combination over a small grid; batches mix ramp stages; the default ramp plays the audio of 25-frame first chunks; ramp boundaries are no worse than 25-frame chunks; bad knobs are rejected naming their key. The real-weights integration test follows the configured schedule. * Qwen3-Omni: start streamed audio after 12 frames, then 13, by default (4, 8, 13) needs the Talker at 2x real time and stalled playback from concurrency 16 on a B200; (12, 13) needs 1.5x and did not stall up to concurrency 32. * Qwen3-Omni: document the ramp's real-time need, drop the unused consumed_tokens The EOS comment in prepare_inputs now says an EOS frame can be a context item of the final flush and the clamp handles it. docs/models.rst gets Qwen3-Omni notes on the 1.5x real-time requirement and the empty-schedule escape hatch. A test runs a trailing codec_eos frame through the default ramp. * test: call get_accelerator_graph_configs, the name #262 gave get_cuda_graph_configs * test: give the remote-code check a default Qwen3-Omni config, which load reads for the Code2Wav knobs

Main gained, from this PR, bounded KV (#816) and RecurrentStatePool gather / scatter_ (#891), and renamed CUDA graphs to accelerator graphs (#262): - bounded KV and the pool take main's versions; - SlotRows reads and writes through the pool's label-based gather / scatter_, slicing (and padding, onto the sink) for head(); - the model and its tests use the AcceleratorGraph names; - the engine keeps a piecewise region that falls back to eager, under main's capture plan; - the position resource keeps main's prefix seeding, with the fused path's range check in its own admit (fused refuses rope_table_dtype); - configs/minicpm_o.yaml sets main's `mps: true` instead of documenting a hand-started daemon; test/cpu-core.txt is gone on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* engine: MLA, fp8 MoE, step budget, symm all-reduce attn: MLA on a latent KV layout (#321). fused_moe: fp8 block-scale experts (#270). worker: per-step token budget; api_server: role on the first token. distributed: group timeout, compile_mode, symmetric-memory all-reduce. * glm52: GLM-5.2 with MLA and fp8 MoE GLM-5.2 at k=0: MLA with the DSA indexer, fp8 MoE, weight loader, configs. Replaces #272. Open: the configs' new max_concurrent_requests caps. * glm52: refuse a prompt without room, as a 400 A prompt of limit - 1 tokens failed its whole decode batch (the speculative step overran the window), and a long-context prompt past index_topk failed its prefill batch. process_prompt now refuses both; prepare_inputs agrees. * serving_chat: decode text across chunk boundaries Byte-level BPE models emit each token's raw bytes, which can end inside a character; decoding per chunk turned those into U+FFFD. * glm52: emit each token's raw bytes Per-token decode turned a character split across byte-level BPE tokens into U+FFFD; the 10-07 GSM8K answers have them at ÷ and √. * attn: MLA wrappers per slot, graph indices per row As FlashInferManager since #243/#280: an eager plan staged into the one wrapper per label while the last step's copy could still be queued, and a captured step whose rows share a cached prefix overran the kv_indices buffer. * glm52: ask for the MoE op pre-build The conductor now builds it only for a model that sets prebuild_fused_moe. * fused_moe: prefill experts as one custom op As fused_experts_fp8: dynamo traced the launch path inline, which breaks the graph and fails on a symbolic token count under torch 2.9.1. * glm52: fail loading loudly, not serve garbage An fp8 checkpoint whose quant config went undetected loaded clean with unscaled e4m3 weights; a failed Hub fetch came back as a local path. * micro_scheduler: step budget takes the oldest The runtime's ready order is handle order in Rust, and handles are recycled: a large prefill could be passed over indefinitely. Backlogged rows count too, since the batch builder tops a backlog up with fresh rows. * api_server: TextChatAdapter checks what it passes A non-numeric max_output_tokens from extra_body reached the conductor's min() and took its message batch down; image or audio parts were decoded and dropped. Both are 400s now. * fused_moe: decode router takes at most 64 tokens Its logits kernel holds every token in one tile; past 128 the first launch raised OutOfResources. The models never send more; now it says so. * glm52: smaller review fixes top_k and min_p reach the sampler; a fused-kernel import failure is logged; MSTAR_GLM52_GRAPH_COMPILE=0 also leaves uncaptured steps uncompiled; the single-rank test config is dummy mode, not a 750 GB pull. * attn: FlashMLA only for eager MLA plans Chosen per plan, a captured bucket could re-plan into the other kernel while the graph replayed the captured one on stale buffers. * glm52: refuse a scaled RoPE from_hf_config read only rope_theta, so a YaRN or dynamic variant ran as plain RoPE, its frequencies and softmax scale wrong. * fused_moe: an empty step skips the align op The CUDA op launched a 0-block grid unchecked; the latched error surfaced in the next unrelated kernel. * test: CI runs the new CPU tests Each new test file with a CPU case, listed in test/cpu-core.txt. * distributed: symm all-reduce on every rank or none A rank whose setup failed fell back to NCCL alone, and its peers' first all-reduce waited on it forever. tp_allreduce_max_kb: 0 now turns it off instead of reading as unset (512). * glm52: tokenizer registers special tokens Built from tokenizer.json alone (transformers 4.x), the tokenizer had no special tokens, so the detokenizer printed <|user|> at the end of replies. * cuda_graph_runner: refuse empty/cudagraph modes An empty MSTAR_GRAPH_COMPILE_MODE compiled in Inductor's default mode silently; reduce-overhead and max-autotune capture inside the runner's own graph, fail it, and the bucket serves eager. * glm52: a prefill stop token ends the reply The prefill walk has no decode_loop: the runtime dropped (and logged) its stop, and decode wrote on past the stop token. The first decode step now stops and emits nothing; the CPU stream harness runs through the runtime. * fused_moe: check block scales and inputs The decode and prefill kernels index scales, x and shared as dense rows; a strided view was mis-scaled or mis-added silently. The prefill test now covers GLM-5.2's hidden 6144, the fused sum-and-add branch. * glm52: a bad sampling kwarg is a 400 None now means unset. A string or null reached the conductor's ingest and raised there, outside any request; process_prompt reads the same kwargs first, where a bad one is the client's 400. * kv: SHM transfer reads an MLA cache by token The reader sliced NHD's K/V axis first; on an MLA cache that took latent columns, not tokens. * glm52: refuse min_p as a 400 The sampler is built without min-p, so the worker refused it at ingest as a 500, after every rank had the request. * conductor: pre-build the MoE op only when asked Every deployment JIT-built the fused-MoE align op before spawning its workers: non-MoE models paid the nvcc build, or hung on a stuck build lock. A model now asks with prebuild_fused_moe. * api_server: TextChatAdapter refuses what it drops n, stop, logprobs, logit_bias and the presence and frequency penalties returned 200 and were dropped. Left off they pass; set, they are a 400. * attn: refuse what only the MLA backend reads sm_scale on another backend was dropped silently; an MLA cache without num_qo_heads defaulted to the one latent KV head, and the kernel planned 1. * glm52: no max_tokens runs to the window The default budget was 1024 on every window, so a thinking reply was cut with length; vLLM runs it to what the prompt leaves, as the submodule's cap now does. * test: CI runs the new CPU tests Each new test file with a CPU case, listed in test/cpu-core.txt. * cli: mstar serve glm52 It answered unknown model; the default is the TP8 yaml. * engine: select TP async scheduling from config The deployment config's tp_async_sched key (true, false or a node list) turns it on so a shipped yaml can opt in; MSTAR_TP_ASYNC_SCHED still overrides it. Default stays off. * pyproject: glm extra Every other model family has one; GLM's runtime came only with all. * api: warmup_requests can repeat their text A long prompt for a path only long prompts take, such as a long-context prefill, without pasting it into the yaml. * glm52: an oversized top_k or seed is a 400 The sampler keeps them in int32 / int64 rows. Past that its step raised on every rank, failing the whole batch; under TP async the followers then hung. * worker: tp_async_sched reads boolean words A quoted yaml true or false arrives as a string and was taken as a node name, so the switch silently did nothing. * glm52: an oversized max_tokens is a 400 Past 2**64 it failed the wire to the conductor as a 500. The bound is int64, what the Rust server's max_tokens holds. * cuda_graph_runner: empty compile mode env is unset Raising at import failed every model's boot on an empty export. * glm52: an out-of-range sampling value is a 400 A negative temperature or a top_p of 0 returned 200 with garbage: GLM-5.2's eager prefill inverted the logits or sampled token 0. * distributed: load timeout ends after warmup dist_timeout_s held for every serving collective, so a wedged rank hung its peers for the whole load timeout. Serving gets torch's default. * api_server: three more dropped fields are a 400 best_of, response_format and stop_token_ids returned 200 and were dropped, as n and stop were. * micro_scheduler: drop the step token budget Chunked prefill (#365) brings a per-(node, walk) max_batch_tokens, and the two would fight over the same batch; GLM moves to it once it lands. The round-robin pass over walks with no room stays, and so do its tests. * attn: refuse context_only on the backends that don't read it Every backend but MLA planned it as an ordinary step, so the queries also attended to the tokens the step appends. * attn: drop FlashMLA Off by default and measured no faster than FlashInfer's MLA kernel; it needs 64-token pages, which no GLM config uses, and served eager plans only. * cuda_graph_runner: drop MSTAR_GRAPH_COMPILE_MODE It set the mode of every capture of every model at once. A capture config's compile_mode, declared per (node, walk), is now checked when the config is built, so a bad one fails warmup instead of serving its buckets eager. * distributed: tp_allreduce and dist_timeout_s from the config only MSTAR_TP_ALLREDUCE, MSTAR_TP_SYMM_AR_MAX_KB and MSTAR_DIST_TIMEOUT_S were a second way to set three serving-config keys, and nothing sets them. * distributed: all-reduce backends as classes AllReduce is NCCL; SymmMemAllReduce takes the small messages it can and hands the rest to it; AllReduce.new picks one from tp_allreduce. CommGroup holds one. resolve_dist_timeout moves to distributed/utils.py. * distributed: one multimem kernel on every rank The one-shot kernel ran only where the caller's tensor started 16-byte aligned, which each rank decided for itself, and ranks on different kernels hang the collective. A misaligned tensor gets the sum in an aligned one and a copy. * docs: the MLA attention resource in adding_models The MLA cache layout, the FLASHINFER_MLA backend and AttentionStep's context_only, beside the resources they extend. * distributed: a class per symmetric-memory kernel SymmMemAllReduce keeps the buffer and the routing; OneShotAllReduce and MultimemAllReduce are its two kernels, in place of a mode string checked on every call. AllReduce.new refuses a tp_allreduce it doesn't know, which used to serve on NCCL behind a warning. * glm52: the accelerator graph names from #262, and capture takes #341's planned * test: the GLM-5.2 tests install the CPU FlashInfer per test, not at import * test: GLM-5.2's attention spec takes #603's AUTO backend * pyproject: the glm extra takes main's Mooncake floor * distributed: compiled one-shot all-reduce stays on NCCL Inductor gives one_shot_all_reduce's input a persistent symmetric buffer sized by the first call, so BAGEL at TP2 with symm_oneshot failed decode capture at its second batch size. * test: #755's looping worker has the timeout restore warmup ends on * test: rust-transport lists TextChatAdapter's chat as not ported mstar-server does not serve GLM's flattened text chat yet (#798). * test: the GLM-5.2 serve harnesses stop as the runtime does They ran decode on the prefill walk's info, whose check_stop ends no loop since 0dedee4, so the loop never stopped. They now check the prefill's stop, then decode on the decode walk, as the capture test does. * glm52: request kwargs through normalize_generation_kwargs The model's own sampling_kwargs duplicated main's check. As main has it, a number sent as a string is a 400, and a top_k past int32 is served as off. * glm52: checkpoints resolve through hf_snapshot.resolve_snapshot_dir The model's copy is gone; a repo id the Hub cannot fetch still raises. * glm52: a request's token budget and prefill stop in its request_state The engine drops the state with the request, so cleanup_request only evicts the DSA k-store. * glm52: positions from a PositionSpec The step's tokens land where the position resource's counter says, not at the KV stream's stored length, as every other model does. The layers still rotate q_pe / k_pe themselves. KVManager.stored_len had no other user. * fp8 block-scale checkpoints in mstar/utils/fp8_block.py Moved from mstar/model/glm52/quantization.py: nothing in it is GLM's. quant_fp8 takes its FP8_DTYPE from there, and the post-load hook walk is in mstar/model/utils.py. * moe: the noaux_tc router and fp8 expert loaders are shared Moved from mstar/model/glm52/components/moe.py into components/moe.py, nothing in them is GLM's. The router keeps main's router contract (a third value, None: it is stateless). * rope: the interleaved fp32 RoPE is shared Moved from mstar/model/glm52/components/rope.py into components/rope.py as InterleavedRotaryEmbedding, same math: one fp32 cos/sin table per forward. * mla: the shared MLAAttention; GLM-5.2's layer adds DSA MLA with weight absorption (projections, the w_kc / w_vc fold, the naive and absorbed paths) moves to components/mla.py. Glm52MLAAttention is that plus its DSA, through two hooks: _select and _attend_selected. * glm52: stage host tensors through main's PinnedStager mstar/utils/pinned_staging.py duplicated mstar.utils.h2d.PinnedStager with a fresh pinned allocation per call. GLM-5.2's preprocess now uses a PinnedStager like Qwen3-TTS; PinnedStager.to_device returns a new device tensor of the values' shape, for the callers that build one. * dsa: GLM-5.2's DSA as an engine resource mstar/engine/resources/dsa holds what glm52/dsa.py and the attention layer did: the index key store, each step's rows off the KV plan, the indexer's scores and top-k, IndexShare's shared pick and sparse MLA over the picks. The model keeps the indexer projections and declares a DSASpec under dsa_long_context. Still eager only; the pages are staged pinned and the selection no longer syncs on positions.max(). * glm52: drop the unmeasured MoE fused all-reduce switch MSTAR_GLM52_MOE_FUSED_ALLREDUCE came with no measurement, which AGENTS.md invariant 9 asks for. The block reduces the routed and the shared expert's outputs separately, as before the switch; ParallelGatedMLP is main's again. * docs: DSASpec in adding_models.rst What the DSA resource builds, what it depends on and what a step declares, with its limits in this PR: eager only, prompts within topk, keys outside the paged pool. * position: refuse rope_table_dtype widths the CUDA kernel can't rotate FlashInfer's cos/sin-cache kernel moves 16 elements a thread; at rotary_dim 8 it returned garbage, different from run to run, instead of failing. Build now refuses a CUDA table whose width is not a multiple of 16. The comment no longer calls the in-kernel angles exact: they come from a fast sin/cos, about a bf16 ulp from a float32 table. * glm52: positions and RoPE from the position resource, softmax_scale from config The forward reads PositionManager.plan's ids (pos_ids("main")) instead of building its own in preprocess. The layers and the DSA indexer rotate with apply_qk, which takes explicit positions (pos_ids=) and rotates q alone (k=None); GLM-5.2's PositionConfig is interleaved over qk_rope_head_dim from an fp32 table, the angles as before. components/rope.py is gone. MLAAttention reads config.softmax_scale. The reduced test config rotates 16 dims (8 before), the narrowest the CUDA table kernel takes. * dsa: a selecting step's page tables go to the device once, in plan attend copied each request's page list to the device in every layer, and the stager's ring waited on a copy eight calls back, so the CPU ran at most eight layers ahead of the GPU. plan() now stages the tables once per step and attend gathers every row's picks in one batch. The decode-only check moves to plan, before any layer runs. The docs row says what the paged follow-up changes. * api_server: TextChatAdapter refuses tool calls A chat with tools reached the model as plain text and came back with a 200, as if the model had declined every tool; _refuse_tool_calls makes it a 400, as for the other chat adapters. A test pins that each turn keeps its role. * glm52: chat turns in their roles, main's intake checks, batched check_stop - process_prompt rendered every message's text as one user turn, the system prompt and earlier replies included; it now builds one message per turn from prompt_parts, as Qwen3.5 does, and a chat the template refuses is a 400. - The sampling-kwarg and min_p checks in process_prompt repeated what the conductor already does at intake since #402 (SamplingReqConfig.validate, and a ValueError before dispatch is a 400); they are gone, and the tests drive the conductor's own methods. - get_max_output_tokens_limit publishes the context window, so a larger max_tokens is the conductor's 400 rather than a silent clamp. - check_stop_batched reads the step's tokens in one tolist; the sampler returns them as row outputs, one clone and one D2H per step. - MSTAR_GLM52_GRAPH_COMPILE becomes model_kwargs.graph_compile. * glm52: GLM-5.2 declares its KV layout Since #832 a paged pool needs its labels declared, and each request's layout returned from process_prompt; without them every GLM-5.2 request was refused. - kv_labels: the one "main" label, growing on "decode" under decode_loop. - process_prompt returns the prompt's ids span and one slot per decode iteration under "kv_layout". - get_loop_caps: max_output_tokens counts the token the prefill emits, so decode_loop runs max_output_tokens - 1 iterations, where check_stop ends it. #808 refuses a cap past the loop's max_iters (context_limit - 1), which is the bound get_max_output_tokens_limit published; #808 removed that hook, so GLM-5.2's override goes too. - The prefill's prompt-aware clamp stays: the conductor holds prompt plus max_tokens to the KV pool's max_seq_len, the clamp to GLM-5.2's own limit, one row short for a step still in flight when the stop lands. - The walk harness passes the model's loop caps to the graph runtime. * glm52: no token-budget clamp; the KV pool is the context window Since #808 and #832 the conductor holds every request to the room its prompt leaves in the KV pool (it refuses a max_tokens past it, and fits a request that names none), and the decode loop's cap ends decode where check_stop does, so no step runs past it. GLM-5.2's own prompt-aware clamp duplicated that, one row short, so a request sized exactly to the window got one token fewer than it asked for. check_stop now reads the request's max_tokens. The KV pool's max_seq_len is now Glm52Model.context_limit(), so the conductor sizes requests against the window the preprocess guard holds contexts to. Before, the two agreed only by convention (index_topk equals max_seq_len in the checkpoint's config); the reduced variant's did not (64 against 512). Claude-Session: https://claude.ai/code/session_01FeseQ9qqtZRk2CqVpUJ6MH * loader: pread into private buffers, an index for stacked rules (from #389) Tensors are read with pread by a small thread pool a bounded window ahead of the consumer instead of faulting through safetensors' mmap, which moves a few pages per fault on ZFS and network filesystems. Per-expert stacked rule lists are matched through a token index instead of trying every rule on every key. Moved here from #389 (caeea87, 1ded60c) so #387 boots without it. * configs: glm52_tp8_longctx points its DSA limits at the resource * test: the GLM-5.2 layout test asks main's validation checks #640 moved the paged-pool and prefix-walk refusals out of mstar/worker/engine_manager.py into mstar/validation/resources.py, as generators of problems, so test_glm52_kv_layout.py failed to import after the merge of main. --------- Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com>

yao-matrix changed the title ~~Add BAGEL XPUGraph capture on Intel GPU (XPU)~~ [4/N] extend CUDA-specific graph API to pytorch's accelerator graph backend, to make graph supported on Intel GPU (XPU) Sep 16, 2026

yao-matrix marked this pull request as draft September 16, 2026 21:28

merceod mentioned this pull request Sep 22, 2026

FLUX.2-klein PR 1: Scaffold for flow matching shared by several models, first of which is flux2-klein #286

Merged

2 tasks done

yao-matrix marked this pull request as ready for review October 2, 2026 16:23

yao-matrix added 2 commits October 2, 2026 16:28

Add accelerator graph abstraction and BAGEL XPUGraph capture

441020c

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

Preserve BAGEL timestep frequencies in FP32

84ba91c

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

yao-matrix force-pushed the codex/bagel-xpu-graph branch from f7512df to e62e7f0 Compare October 2, 2026 16:32

yao-matrix added 6 commits October 2, 2026 16:33

Remove CUDA graph compatibility API

e62e7f0

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

Document accelerator graph APIs and device support

b0cdfa3

Remove CUDA assumptions from shared accelerator graph execution

4b9f62f

Add BAGEL accelerator graph deployment switch

9924d2d

Use accelerator graph names throughout BAGEL

eee3abc

Unify model deployment switches under accelerator graph names

38d8f32

yao-matrix mentioned this pull request Oct 2, 2026

Batch XPU token sampling with per-request device RNG #347

Merged

yao-matrix added 9 commits October 6, 2026 14:58

Separate CUDA and XPU graph backend implementations

8c7ca46

Update bagel_xpu_cfg_tp2.yaml

ac047f9

Size XPU graph attention plans for resident KV capacity

49f7d45

Update bagel_xpu_cfg_tp2.yaml

56d63f5

Disable BAGEL XPU CFG TP2 accelerator graphs by default

22a2599

Default BAGEL accelerator graph capture to disabled

f91d9d2

Update bagel_xpu_cfg_tp2.yaml

609a892

Clarify BAGEL graph capture opt-in in documentation

ba6a6eb

yao-matrix added 2 commits October 7, 2026 00:25

Fix graph backend mock for CPU-only PyTorch wheels

17dc7d1

Signed-off-by: Matrix Yao <matrix.yao@intel.com>

Gaurav-Shah05 mentioned this pull request Oct 19, 2026

engine, communication: capture only the CUDA graph buckets that leave room to serve #341

Merged

2 tasks done

This was referenced Oct 20, 2026

[Model] Add VoxCPM2: tokenizer-free TTS with voice design, cloning and streamed 48 kHz audio #496

Closed

[Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495

Merged

arun-elr added a commit to arun-elr/mstar that referenced this pull request Oct 20, 2026

test: call get_accelerator_graph_configs, the name mstar-project#262 …

92a58cf

…gave get_cuda_graph_configs

Gaurav-Shah05 mentioned this pull request Oct 20, 2026

Qwen3.5: video input, matching transformers' Qwen3VLProcessor #427

Merged

vasilevklart added a commit that referenced this pull request Oct 22, 2026

glm52, glm5_next: the sampled verify capture takes #262's accelerator…

b957cf9

… graph config

arun-elr added a commit to arun-elr/mstar that referenced this pull request Oct 24, 2026

test: call get_accelerator_graph_configs, the name mstar-project#262 …

e6198da

…gave get_cuda_graph_configs

arun-elr mentioned this pull request Oct 25, 2026

Add an opt-in ragged FA3 varlen backend for cacheless attention. #863

Merged

2 tasks done

Navigation Menu

Repository navigation

[4/N] extend CUDA-specific graph API to pytorch's accelerator graph backend, to make graph supported on Intel GPU (XPU) - #262

[4/N] extend CUDA-specific graph API to pytorch's accelerator graph backend, to make graph supported on Intel GPU (XPU)#262
Gaurav-Shah05 merged 52 commits into
mstar-project:mainfrom
yao-matrix:codex/bagel-xpu-graph

yao-matrix commented Sep 16, 2026 •

edited

Loading

Uh oh!

yao-matrix commented Oct 2, 2026 •

edited

Loading

Uh oh!

yao-matrix commented Oct 5, 2026

Uh oh!

NSagan271 commented Oct 5, 2026

Uh oh!

yao-matrix commented Oct 6, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

yao-matrix commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Verification

Uh oh!

yao-matrix commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yao-matrix commented Oct 5, 2026

Uh oh!

NSagan271 commented Oct 5, 2026

Uh oh!

yao-matrix commented Oct 6, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yao-matrix commented Sep 16, 2026 •

edited

Loading

yao-matrix commented Oct 2, 2026 •

edited

Loading