From 1759021946f9544e2afda5afe68d59ef33f9ae03 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 11:18:14 -0700 Subject: [PATCH 01/59] Overnight perf push: the code changes, without the run artifacts The code of the overnight multi-agent perf push (dghosef/mstar#13 at d248477ec), on the same base (main abdfd5802), as one commit. Byte-identical to #13's tree outside overnight_artifacts/, which stays on #13 with the benchmark results, profiles and harnesses. Co-Authored-By: Claude Opus 5.5 --- benchmark/waypoint/benchmark_streaming.py | 4 +- configs/bagel_cfg_parallel_2gpu.yaml | 36 ++ configs/chatterbox_multilingual_split.yaml | 31 ++ configs/chatterbox_split.yaml | 30 ++ configs/chatterbox_turbo_split.yaml | 28 ++ configs/cosmos3_droid.yaml | 44 ++ configs/cosmos3_edge_dp2.yaml | 79 ++++ configs/cosmos3_edge_droid.yaml | 38 ++ configs/cosmos3_edge_sp2.yaml | 84 ++++ configs/cosmos3_nano_dp2.yaml | 15 + configs/fish_s2_pro.yaml | 5 + configs/kokoro_bf16.yaml | 13 + configs/lingbot_world_1_3b_fast.yaml | 22 + configs/omnivoice_split.yaml | 14 + configs/qwen3_5_0.8b.yaml | 9 +- configs/qwen3_5_27b_tp4.yaml | 13 +- configs/qwen3_5_2b.yaml | 9 +- configs/qwen3_5_4b.yaml | 9 +- configs/qwen3_5_9b.yaml | 9 +- configs/qwen3_5_9b_tp2.yaml | 13 +- configs/qwen3omni_pd_2gpu.yaml | 18 + configs/qwen3tts_split.yaml | 26 + configs/wan22.yaml | 14 +- configs/wan22_dp2.yaml | 16 + configs/wan22_sp2.yaml | 29 ++ configs/z_image_turbo.yaml | 17 +- docs/environment_variables.rst | 5 - docs/installation.rst | 4 +- docs/models.rst | 152 +++++- docs/serving.rst | 13 +- mstar/api_server/entrypoint.py | 59 ++- mstar/communication/tensor_store.py | 9 + mstar/communication/tensors.py | 5 + mstar/conductor/conductor.py | 34 +- mstar/engine/accelerator_graph_config.py | 4 + mstar/engine/accelerator_graph_runner.py | 5 + mstar/engine/engine.py | 2 +- mstar/engine/resources/attn/ragged/base.py | 2 +- mstar/engine/resources/attn/ragged/config.py | 7 + mstar/engine/resources/convenience.py | 41 +- mstar/model/bagel/bagel_model.py | 22 +- mstar/model/bagel/components/vit_encoder.py | 8 + mstar/model/bagel/config.py | 3 + mstar/model/bagel/submodules.py | 446 ++++++++++++++---- mstar/model/base.py | 4 +- mstar/model/chatterbox/components/t3.py | 8 +- mstar/model/chatterbox/submodules.py | 8 +- .../components/diffusion/denoise_loop.py | 13 +- .../components/diffusion/qwen3/encoder.py | 21 +- mstar/model/components/moe.py | 22 + mstar/model/cosmos3/config.py | 9 +- mstar/model/cosmos3/cosmos3_model.py | 9 +- mstar/model/cosmos3/submodules.py | 31 +- mstar/model/cosmos3/tests/test_serving.py | 44 +- .../cosmos3/tests/test_windowed_decode.py | 22 + .../fish_s2_pro/components/dac_layers.py | 42 +- mstar/model/fish_s2_pro/submodules.py | 44 +- mstar/model/flux2_klein/flux2_klein_model.py | 20 +- mstar/model/flux2_klein/submodules.py | 57 ++- .../higgs_audio/components/audio_tower.py | 34 +- mstar/model/higgs_audio/higgs_audio_model.py | 60 ++- mstar/model/higgs_audio/submodules.py | 314 +++++++++--- mstar/model/kokoro/components/lstm.py | 40 +- mstar/model/kokoro/components/synthesizer.py | 6 + mstar/model/kokoro/config.py | 7 + mstar/model/kokoro/kokoro_model.py | 4 +- mstar/model/kokoro/submodules.py | 19 +- mstar/model/lingbot_world/components/dit.py | 117 ++++- mstar/model/lingbot_world/config.py | 18 +- .../lingbot_world/lingbot_world_model.py | 7 + mstar/model/lingbot_world/submodules.py | 15 +- mstar/model/media_limits.py | 5 + .../omnilingual_asr/components/wav2vec2.py | 98 +++- mstar/model/omnilingual_asr/config.py | 2 + .../omnilingual_asr/omnilingual_asr_model.py | 34 +- mstar/model/omnilingual_asr/submodules.py | 185 +++++++- mstar/model/omnivoice/components/backbone.py | 204 +++++--- mstar/model/omnivoice/components/codec.py | 9 +- mstar/model/omnivoice/components/silence.py | 146 ++++++ mstar/model/omnivoice/config.py | 15 + mstar/model/omnivoice/omnivoice_model.py | 45 +- mstar/model/omnivoice/submodules.py | 157 +++++- mstar/model/orpheus/submodules.py | 51 +- mstar/model/qwen3_5/qwen3_5_model.py | 116 ++++- mstar/model/qwen3_5/submodules.py | 197 +++++--- .../model/qwen3_omni/components/cp_kernels.py | 161 +++++++ mstar/model/qwen3_omni/components/talker.py | 76 ++- mstar/model/qwen3_omni/components/thinker.py | 1 + mstar/model/qwen3_omni/config.py | 12 +- mstar/model/qwen3_omni/qwen3_omni_model.py | 7 +- mstar/model/qwen3_omni/submodules.py | 338 ++++++++----- mstar/model/qwen3_tts/components/talker.py | 39 +- mstar/model/qwen3_tts/submodules.py | 141 ++++-- mstar/model/vjepa2/components/ac_predictor.py | 27 +- mstar/model/vjepa2/components/rope_utils.py | 30 +- mstar/model/vjepa2/config.py | 4 + mstar/model/vjepa2/submodules.py | 56 ++- mstar/model/vjepa2/vjepa2_model.py | 19 +- mstar/model/wan22/components/dit.py | 108 ++++- mstar/model/wan22/config.py | 9 + mstar/model/wan22/submodules.py | 118 ++++- mstar/model/wan22/wan22_model.py | 56 ++- mstar/model/z_image/submodules.py | 109 ++++- mstar/model/z_image/z_image_model.py | 8 +- mstar/utils/fused_moe/kernels.py | 47 ++ mstar/utils/fused_moe/router.py | 64 +++ mstar/utils/fused_moe/runner.py | 7 +- mstar/utils/knobs.py | 8 - mstar/worker/worker.py | 6 + test/chatterbox/test_t3_components.py | 3 +- .../test_qwen3_omni_code2wav_real_weights.py | 82 +++- test/modular/test_bagel_cfg_prune.py | 315 +++++++++++++ test/modular/test_bagel_postprocess.py | 21 + test/modular/test_chatterbox_model.py | 4 +- .../test_conductor_replica_assignment.py | 52 ++ test/modular/test_fish_s2_pro_model.py | 8 +- test/modular/test_flux2_klein_model.py | 5 +- test/modular/test_higgs_audio_decode_hooks.py | 78 +++ test/modular/test_higgs_audio_tower.py | 30 ++ test/modular/test_kokoro_model.py | 5 + test/modular/test_moe_router.py | 88 ++++ test/modular/test_omnilingual_asr_model.py | 58 +++ test/modular/test_orpheus_audio_vocab.py | 97 ++++ test/modular/test_piecewise_can_lease.py | 22 + test/modular/test_qwen3_5_device_loopback.py | 20 +- .../modular/test_qwen3_5_prompt_interleave.py | 29 +- test/modular/test_qwen3_5_video.py | 21 +- .../test_qwen3_omni_code2wav_stream.py | 78 ++- test/modular/test_qwen3_omni_fused_kernels.py | 124 +++++ test/modular/test_qwen3_tts_model.py | 146 +++++- test/modular/test_ragged_attention.py | 45 ++ test/modular/test_ragged_attention_cpu.py | 23 +- test/modular/test_serving_metrics.py | 2 +- test/modular/test_shm_tensor_comm.py | 28 ++ test/modular/test_startup_warmup.py | 47 ++ test/modular/test_step_prep_host_syncs.py | 20 +- test/modular/test_wan22_model.py | 60 ++- test/modular/test_wan_vae_compile_options.py | 17 +- test/modular/test_worker_drain.py | 1 + test/modular/test_z_image_batch_invariance.py | 4 +- .../test_omnivoice_backbone_graphs.py | 101 ++++ .../omnivoice/test_omnivoice_fused_weights.py | 31 ++ test/omnivoice/test_omnivoice_parity.py | 2 +- test/omnivoice/test_omnivoice_silence.py | 46 ++ test/waypoint/serve_rollout.py | 11 +- 145 files changed, 5922 insertions(+), 958 deletions(-) create mode 100644 configs/bagel_cfg_parallel_2gpu.yaml create mode 100644 configs/chatterbox_multilingual_split.yaml create mode 100644 configs/chatterbox_split.yaml create mode 100644 configs/chatterbox_turbo_split.yaml create mode 100644 configs/cosmos3_edge_dp2.yaml create mode 100644 configs/cosmos3_edge_sp2.yaml create mode 100644 configs/cosmos3_nano_dp2.yaml create mode 100644 configs/kokoro_bf16.yaml create mode 100644 configs/lingbot_world_1_3b_fast.yaml create mode 100644 configs/omnivoice_split.yaml create mode 100644 configs/qwen3omni_pd_2gpu.yaml create mode 100644 configs/qwen3tts_split.yaml create mode 100644 configs/wan22_dp2.yaml create mode 100644 configs/wan22_sp2.yaml create mode 100644 mstar/model/omnivoice/components/silence.py create mode 100644 mstar/model/qwen3_omni/components/cp_kernels.py create mode 100644 mstar/utils/fused_moe/router.py create mode 100644 test/modular/test_bagel_cfg_prune.py create mode 100644 test/modular/test_bagel_postprocess.py create mode 100644 test/modular/test_conductor_replica_assignment.py create mode 100644 test/modular/test_higgs_audio_decode_hooks.py create mode 100644 test/modular/test_higgs_audio_tower.py create mode 100644 test/modular/test_moe_router.py create mode 100644 test/modular/test_orpheus_audio_vocab.py create mode 100644 test/modular/test_qwen3_omni_fused_kernels.py create mode 100644 test/omnivoice/test_omnivoice_backbone_graphs.py create mode 100644 test/omnivoice/test_omnivoice_fused_weights.py create mode 100644 test/omnivoice/test_omnivoice_silence.py diff --git a/benchmark/waypoint/benchmark_streaming.py b/benchmark/waypoint/benchmark_streaming.py index a622f27c9d..282895bef1 100644 --- a/benchmark/waypoint/benchmark_streaming.py +++ b/benchmark/waypoint/benchmark_streaming.py @@ -558,7 +558,9 @@ def _dit_step_timestamps(log_text: str, request_ids: set[str]) -> list[float]: batch = ast.literal_eval(line.split(marker, 1)[1].strip()) except (SyntaxError, ValueError): continue - if not isinstance(batch, (list, tuple)) or not any(rid in request_ids for rid in batch): + if not isinstance(batch, (list, tuple)) or not any( + str(rid).split("#", 1)[0] in request_ids for rid in batch + ): continue match = _LOG_TIMESTAMP_RE.match(line) if match is None: diff --git a/configs/bagel_cfg_parallel_2gpu.yaml b/configs/bagel_cfg_parallel_2gpu.yaml new file mode 100644 index 0000000000..345a983f00 --- /dev/null +++ b/configs/bagel_cfg_parallel_2gpu.yaml @@ -0,0 +1,36 @@ +model: "bagel" +max_seq_len: 32768 +max_concurrent_requests: 16 + +resources: + kv: + cpu_offload_pages: 2048 + +# CFG-parallel text-to-image on two GPUs: the main branch (and the combine) on rank 0, +# both guidance branches on rank 1. With the default t2i settings (cfg_img_scale 1.0) the +# image branch never contributes, so the step is bounded by one branch per GPU. +node_groups: + - node_names: + - vit_encoder + ranks: + - 0 + + - node_names: + - vae_encoder + - vae_decoder + ranks: + - 0 + + - node_names: + - LLM + - combine_cfg + ranks: + - 0 + + - node_names: + - LLM_cfg_text + - LLM_cfg_img + ranks: + - 1 + graph_walks: + - image_gen_cfg diff --git a/configs/chatterbox_multilingual_split.yaml b/configs/chatterbox_multilingual_split.yaml new file mode 100644 index 0000000000..7d9d9f9449 --- /dev/null +++ b/configs/chatterbox_multilingual_split.yaml @@ -0,0 +1,31 @@ +# configs/chatterbox_multilingual.yaml with S3Gen on its own worker. Both workers run on GPU 0 +# (rank_devices), so T3's decode steps no longer alternate with S3Gen's flow +# solves on one worker loop; run a CUDA MPS daemon for the two workers' +# kernels to overlap. Same graph, resources and checkpoint. +model: "chatterbox_multilingual" +# cond (34) + text (<= 512) + speech (<= 4096) tokens per stream +max_seq_len: 4680 +# A CFG request at full text and max_new_tokens_limit speech tokens holds +# 2 x 25 pages; 30 such requests and the sink page fit in 1536 (checked at +# startup). A full pool would stall decode. +max_concurrent_requests: 30 +resources: + t3_kv: + # 30 layers x 16 kv heads x 64 dims x bf16 = 120 KB per token and stream + page_size: 64 + max_num_pages: 1536 +node_groups: + - node_names: [voice_encoder, T3] + ranks: [0] + graph_walks: [prefill, prefill_voice, decode] + - node_names: [s3gen] + ranks: [1] + graph_walks: [s3gen_chunk, s3gen_chunk_voice] +rank_devices: {1: 0} +# Two workers share the GPU; each plans its captures within its share. +gpu_memory_fraction: 0.45 +model_kwargs: + max_new_tokens_limit: 1000 # sized with max_concurrent_requests to fit t3_kv +# default_language: en # language_id used when a request gives none (23 codes, see docs) +# voices_dir: /path/to/preset/wavs # voice="" loads .wav +# watermark: true # Perth watermark on every output diff --git a/configs/chatterbox_split.yaml b/configs/chatterbox_split.yaml new file mode 100644 index 0000000000..526ba11263 --- /dev/null +++ b/configs/chatterbox_split.yaml @@ -0,0 +1,30 @@ +# configs/chatterbox.yaml with S3Gen on its own worker. Both workers run on GPU 0 +# (rank_devices), so T3's decode steps no longer alternate with S3Gen's flow +# solves on one worker loop; run a CUDA MPS daemon for the two workers' +# kernels to overlap. Same graph, resources and checkpoint. +model: "chatterbox" +# cond (34) + text (<= 512) + speech (<= 4096) tokens per stream +max_seq_len: 4680 +# A CFG request at full text and max_new_tokens_limit speech tokens holds +# 2 x 25 pages; 30 such requests and the sink page fit in 1536 (checked at +# startup). A full pool would stall decode. +max_concurrent_requests: 30 +resources: + t3_kv: + # 30 layers x 16 kv heads x 64 dims x bf16 = 120 KB per token and stream + page_size: 64 + max_num_pages: 1536 +node_groups: + - node_names: [voice_encoder, T3] + ranks: [0] + graph_walks: [prefill, prefill_voice, decode] + - node_names: [s3gen] + ranks: [1] + graph_walks: [s3gen_chunk, s3gen_chunk_voice] +rank_devices: {1: 0} +# Two workers share the GPU; each plans its captures within its share. +gpu_memory_fraction: 0.45 +model_kwargs: + max_new_tokens_limit: 1000 # sized with max_concurrent_requests to fit t3_kv +# voices_dir: /path/to/preset/wavs # voice="" loads .wav +# watermark: true # Perth watermark on every output diff --git a/configs/chatterbox_turbo_split.yaml b/configs/chatterbox_turbo_split.yaml new file mode 100644 index 0000000000..3afe546d99 --- /dev/null +++ b/configs/chatterbox_turbo_split.yaml @@ -0,0 +1,28 @@ +# configs/chatterbox_turbo.yaml with S3Gen on its own worker. Both workers run on GPU 0 +# (rank_devices), so T3's decode steps no longer alternate with S3Gen's flow +# solves on one worker loop; run a CUDA MPS daemon for the two workers' +# kernels to overlap. Same graph, resources and checkpoint. +model: "chatterbox_turbo" +# cond (376) + text (<= 512) + speech (<= 4096) tokens; single stream (no CFG) +max_seq_len: 5000 +# A request at full text and max_new_tokens_limit speech tokens holds 30 +# pages; 50 such requests and the sink page fit in 1536 (checked at startup). +# A full pool would stall decode. +max_concurrent_requests: 50 +resources: + t3_kv: + # 24 layers x 16 kv heads x 64 dims x bf16 = 96 KB per token + page_size: 64 + max_num_pages: 1536 +node_groups: + - node_names: [voice_encoder, T3] + ranks: [0] + graph_walks: [prefill, prefill_voice, decode] + - node_names: [s3gen] + ranks: [1] + graph_walks: [s3gen_chunk, s3gen_chunk_voice] +rank_devices: {1: 0} +# Two workers share the GPU; each plans its captures within its share. +gpu_memory_fraction: 0.45 +model_kwargs: + max_new_tokens_limit: 1000 # sized with max_concurrent_requests to fit t3_kv diff --git a/configs/cosmos3_droid.yaml b/configs/cosmos3_droid.yaml index a7a06d0aa4..35a5e9b5a6 100644 --- a/configs/cosmos3_droid.yaml +++ b/configs/cosmos3_droid.yaml @@ -29,3 +29,47 @@ node_groups: # (the only consumer of that node) is never built. - node_names: ["vae_encoder", "vae_decoder"] ranks: [0] +# The first policy request of a fresh server compiles the action denoise step, +# and concurrent policy requests batch into a step that compiles once per batch +# size: warm the single step and each batch size up to the cap, at the +# documented 32-step chunk. +warmup_requests: + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 2 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 3 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 4 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 5 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 6 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 7 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 8 diff --git a/configs/cosmos3_edge_dp2.yaml b/configs/cosmos3_edge_dp2.yaml new file mode 100644 index 0000000000..0f2e3792fd --- /dev/null +++ b/configs/cosmos3_edge_dp2.yaml @@ -0,0 +1,79 @@ +model: "cosmos3_edge" +# Sequence-length hint for the scheduler. The conductor only asserts its +# presence; the real per-request capacity is the KV pool below. +max_seq_len: 8192 +# KV pool sizing. Edge is 28 layers x 8 KV heads x 128 (half of Nano's per-page +# footprint); one pool serves both the DiT's frozen text prefixes (a few pages +# per guidance branch) and the reasoner's chat context (up to 131k tokens per +# request). 1024 pages x 128 tokens is ~7.5 GB in bf16. +resources: + kv: + max_num_pages: 1024 +model_kwargs: + # Denoise CUDA-graph capture knobs (serving). cuda_graph=false serves eager; + # graph_max_latent_area caps which resolutions are captured (latent H*W). + # The COSMOS3_DISABLE_CUDA_GRAPH / COSMOS3_GRAPH_MAX_LATENT_AREA env vars + # override. + accelerator_graph: true + graph_max_latent_area: 2000 + # Video denoise steps stay eager on the dense FA3 path: capturing them + # (gen_capture_video) forces the paged attention into the graph, which + # measured 3-6% slower at 832x480 for both the 121-frame clip and the + # 29-frame window (H100, 2026-09-17). Add tiers here only for small, + # launch-bound shapes. + gen_capture_video: [] + # Edge is 480p-native: the model card's generator recipe is 832x480 video + # (121 frames, 24 fps, 20 UniPC steps, guidance 6.0, flow shift 12.0 on the + # native flow schedule) and 640x640 images. Requests override any of these. + image_size_default: [640, 640] + video_size_default: [832, 480] + num_frames_video: 121 + num_inference_steps_video: 20 + guidance_scale: 6.0 + flow_shift_video: 12.0 + # Image-to-video conditioning follows the model card's diffusers recipe: + # cover-scale, antialiased resize, center crop, 8-bit rounding. + conditioning_resize: aspect_crop + # Action requests (forward/inverse dynamics, policy) follow the Edge model + # card: 30 steps, guidance 1.0, flow shift 10.0. + num_inference_steps_action: 30 + guidance_scale_action: 1.0 + flow_shift_action: 10.0 + # Streaming rollout: windowed autoregressive video (the video_gen_ar walk, + # the vae_decoder_ar node and its streaming decoder partition). Requests opt + # in per call with window_mode ("chained" | "kv") and stream each finished + # window with stream_video; the kv-mode context horizon (context_frames) + # bounds the committed K/V a long rollout keeps. + enable_windowed_video: true + # The streaming decoder keeps each live stream's VAE conv cache (1.33 GiB + # of tensors per stream at 832x480 in bf16; it grows with the frame area and + # doubles in fp32). Uncomment to cap the warm streams per decoder copy; + # streams past the cap take the stateless context re-decode. + # windowed_decode_max_warm_streams: 0 # 0: no cap + # Latents per decoder call in the streaming decode (default 1). 0 decodes a + # steady window's new latents in one call: the decoder alone is ~6% faster at + # 832x480, but served end to end it is no faster and takes about 15 GiB more + # peak (832x480 bf16, measured on B200). + # windowed_decode_latents_per_call: 1 +node_groups: + # Two GPUs, throughput-oriented: a whole replica per GPU (the DiT and the + # reasoner share one transformer and KV pool, so they stay in one group with + # the encoders and decoders); the conductor sends each request to the replica + # serving fewer requests. + - node_names: ["dit", "reasoner", "vision_encoder", "vae_encoder", "vae_decoder", "vae_decoder_ar"] + ranks: [0, 1] +# Requests the server runs after the workers are ready and before it binds, +# so no client pays a shape's first torch.compile of the denoise step (15-25 s +# to the first frame on a fresh server against 1.9 s warm). One per default +# video shape: the 121-frame clip and the 29-frame kv window (57 frames is two +# windows) and an image-conditioned rollout (the conditioning encode and the +# i2v prefill are their own first-time cost). The compile keys on shapes, not +# steps, so four steps do. Startup grows by their wall time. +warmup_requests: + - text: warmup + output_modalities: [video] + model_kwargs: {num_frames: 121, num_inference_steps: 4} + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [video] + model_kwargs: {window_mode: kv, num_frames: 57, num_inference_steps: 4} diff --git a/configs/cosmos3_edge_droid.yaml b/configs/cosmos3_edge_droid.yaml index dc10915045..2632aa69e2 100644 --- a/configs/cosmos3_edge_droid.yaml +++ b/configs/cosmos3_edge_droid.yaml @@ -38,3 +38,41 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10} + # Concurrent policy requests batch, and the batched step compiles once per + # batch size and shape: warm each size up to the batch cap at the documented + # 32-step chunk. + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 2 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 3 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 4 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 5 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 6 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 7 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 8 diff --git a/configs/cosmos3_edge_sp2.yaml b/configs/cosmos3_edge_sp2.yaml new file mode 100644 index 0000000000..bd68266586 --- /dev/null +++ b/configs/cosmos3_edge_sp2.yaml @@ -0,0 +1,84 @@ +model: "cosmos3_edge" +# Sequence-length hint for the scheduler. The conductor only asserts its +# presence; the real per-request capacity is the KV pool below. +max_seq_len: 8192 +# KV pool sizing. Edge is 28 layers x 8 KV heads x 128 (half of Nano's per-page +# footprint); one pool serves both the DiT's frozen text prefixes (a few pages +# per guidance branch) and the reasoner's chat context (up to 131k tokens per +# request). 1024 pages x 128 tokens is ~7.5 GB in bf16. +resources: + kv: + max_num_pages: 1024 +model_kwargs: + # generation only: the reasoner shares the transformer and has no SP path + enable_reasoner: false + # Denoise CUDA-graph capture knobs (serving). cuda_graph=false serves eager; + # graph_max_latent_area caps which resolutions are captured (latent H*W). + # The COSMOS3_DISABLE_CUDA_GRAPH / COSMOS3_GRAPH_MAX_LATENT_AREA env vars + # override. + accelerator_graph: true + graph_max_latent_area: 2000 + # Video denoise steps stay eager on the dense FA3 path: capturing them + # (gen_capture_video) forces the paged attention into the graph, which + # measured 3-6% slower at 832x480 for both the 121-frame clip and the + # 29-frame window (H100, 2026-09-17). Add tiers here only for small, + # launch-bound shapes. + gen_capture_video: [] + # Edge is 480p-native: the model card's generator recipe is 832x480 video + # (121 frames, 24 fps, 20 UniPC steps, guidance 6.0, flow shift 12.0 on the + # native flow schedule) and 640x640 images. Requests override any of these. + image_size_default: [640, 640] + video_size_default: [832, 480] + num_frames_video: 121 + num_inference_steps_video: 20 + guidance_scale: 6.0 + flow_shift_video: 12.0 + # Image-to-video conditioning follows the model card's diffusers recipe: + # cover-scale, antialiased resize, center crop, 8-bit rounding. + conditioning_resize: aspect_crop + # Action requests (forward/inverse dynamics, policy) follow the Edge model + # card: 30 steps, guidance 1.0, flow shift 10.0. + num_inference_steps_action: 30 + guidance_scale_action: 1.0 + flow_shift_action: 10.0 + # Streaming rollout: windowed autoregressive video (the video_gen_ar walk, + # the vae_decoder_ar node and its streaming decoder partition). Requests opt + # in per call with window_mode ("chained" | "kv") and stream each finished + # window with stream_video; the kv-mode context horizon (context_frames) + # bounds the committed K/V a long rollout keeps. + enable_windowed_video: true + # The streaming decoder keeps each live stream's VAE conv cache (1.33 GiB + # of tensors per stream at 832x480 in bf16; it grows with the frame area and + # doubles in fp32). Uncomment to cap the warm streams per decoder copy; + # streams past the cap take the stateless context re-decode. + # windowed_decode_max_warm_streams: 0 # 0: no cap + # Latents per decoder call in the streaming decode (default 1). 0 decodes a + # steady window's new latents in one call: the decoder alone is ~6% faster at + # 832x480, but served end to end it is no faster and takes about 15 GiB more + # peak (832x480 bf16, measured on B200). + # windowed_decode_latents_per_call: 1 +node_groups: + # Two GPUs, latency-oriented: the DiT runs Ulysses sequence-parallel over + # both (see cosmos3_nano_sp2.yaml). The reasoner is not sequence-parallel and + # shares the DiT's transformer, so this deployment serves generation only + # (enable_reasoner: false); cosmos3_edge_dp2.yaml keeps chat. + - node_names: ["dit"] + ranks: [0, 1] + sp_size: 2 + - node_names: ["vae_encoder", "vae_decoder", "vae_decoder_ar"] + ranks: [0] +# Requests the server runs after the workers are ready and before it binds, +# so no client pays a shape's first torch.compile of the denoise step (15-25 s +# to the first frame on a fresh server against 1.9 s warm). One per default +# video shape: the 121-frame clip and the 29-frame kv window (57 frames is two +# windows) and an image-conditioned rollout (the conditioning encode and the +# i2v prefill are their own first-time cost). The compile keys on shapes, not +# steps, so four steps do. Startup grows by their wall time. +warmup_requests: + - text: warmup + output_modalities: [video] + model_kwargs: {num_frames: 121, num_inference_steps: 4} + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [video] + model_kwargs: {window_mode: kv, num_frames: 57, num_inference_steps: 4} diff --git a/configs/cosmos3_nano_dp2.yaml b/configs/cosmos3_nano_dp2.yaml new file mode 100644 index 0000000000..195e197361 --- /dev/null +++ b/configs/cosmos3_nano_dp2.yaml @@ -0,0 +1,15 @@ +model: "cosmos3" +# See cosmos3_nano.yaml. Two GPUs, throughput-oriented: every node is +# replicated on both and each request runs end to end on one of them (the +# conductor sends it to the replica serving fewer requests). For one request's +# latency, cosmos3_nano_sp2.yaml splits the DiT over both GPUs instead. +max_seq_len: 8192 +resources: + kv: + max_num_pages: 1024 +model_kwargs: + accelerator_graph: true + graph_max_latent_area: 2000 +node_groups: + - node_names: ["dit", "vae_encoder", "vae_decoder", "audio_decoder"] + ranks: [0, 1] diff --git a/configs/fish_s2_pro.yaml b/configs/fish_s2_pro.yaml index d5ffd4e98e..cc871a0815 100644 --- a/configs/fish_s2_pro.yaml +++ b/configs/fish_s2_pro.yaml @@ -35,3 +35,8 @@ warmup_requests: files: [hf://hexgrad/Kokoro-82M/samples/af_heart_0.wav] output_modalities: [audio] model_kwargs: {max_new_tokens: 32} +# The codec's conv stack is ~0.3 ms per frame and row in exact fp32, most of +# the GPU time once requests overlap. TF32 halves it; the audio stays within +# ~60 dB SNR of the exact decode. Drop this to get the exact (TF32-off) codec. +model_kwargs: + codec_exact_fp32: false diff --git a/configs/kokoro_bf16.yaml b/configs/kokoro_bf16.yaml new file mode 100644 index 0000000000..1d741b3803 --- /dev/null +++ b/configs/kokoro_bf16.yaml @@ -0,0 +1,13 @@ +# configs/kokoro.yaml with the vocoder trunk in bfloat16 (model_kwargs decoder_dtype; +# the harmonic source and the STFT head stay float32). On an H200 it served +# +20% audio-s/s at concurrency 32 and +25% at 64 with the same Whisper WER over +# 156 clips (1.24%) and the same waveform agreement with the float32 path as two +# float32 runs have with each other; float32 stays the parity-tested default. +model: "kokoro" +max_seq_len: 512 +node_groups: + - node_names: [kokoro] + ranks: [0] + graph_walks: [synth] +model_kwargs: + decoder_dtype: "bfloat16" diff --git a/configs/lingbot_world_1_3b_fast.yaml b/configs/lingbot_world_1_3b_fast.yaml new file mode 100644 index 0000000000..2cf4ca4b2e --- /dev/null +++ b/configs/lingbot_world_1_3b_fast.yaml @@ -0,0 +1,22 @@ +model: "lingbot_world_1_3b" + +# Speed over bit-parity with the reference (~1.5x at 480x832 on an H200): +# flash-attn-4 attention, compiled block pointwise ops, an autotuned condition +# encoder and a bf16 decode. Rollouts drift from the reference's about as far as +# under any other attention kernel. Needs flash-attn-4. +model_kwargs: + reference_numerics: false + +# No node holds an engine KV cache (the rolling K/V is per-request state in +# the dit); max_seq_len only satisfies the conductor's config check and matches +# the UMT5 text length. +max_seq_len: 512 + +# One rollout per GPU: the dense K/V and the VAE caches are per request, so the +# model refuses any other value (LingBotWorldModel.validate_config_yaml). +max_concurrent_requests: 1 + +# Single-GPU colocated; the Wan VAE encode and decode are fused into the dit forward. +node_groups: + - node_names: [text_encoder, dit] + ranks: [0] diff --git a/configs/omnivoice_split.yaml b/configs/omnivoice_split.yaml new file mode 100644 index 0000000000..04aeb10583 --- /dev/null +++ b/configs/omnivoice_split.yaml @@ -0,0 +1,14 @@ +# configs/omnivoice.yaml with code2wav on its own worker. Both workers run on +# GPU 0 (rank_devices), so the backbone's unmask loop no longer stops for each +# finished request's codec decode and host-side post-processing; run a CUDA +# MPS daemon for the two workers' kernels to overlap. +model: "omnivoice" +max_seq_len: 8192 +node_groups: + - node_names: [ref_encoder, backbone] + ranks: [0] + - node_names: [code2wav] + ranks: [1] +rank_devices: {1: 0} +# Two workers share the GPU; each plans its captures within its share. +gpu_memory_fraction: 0.45 diff --git a/configs/qwen3_5_0.8b.yaml b/configs/qwen3_5_0.8b.yaml index 35366e8290..2fca227719 100644 --- a/configs/qwen3_5_0.8b.yaml +++ b/configs/qwen3_5_0.8b.yaml @@ -1,5 +1,8 @@ model: "qwen3_5_0.8b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. @@ -10,9 +13,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3_5_27b_tp4.yaml b/configs/qwen3_5_27b_tp4.yaml index 6c34bab5c0..06c4470a99 100644 --- a/configs/qwen3_5_27b_tp4.yaml +++ b/configs/qwen3_5_27b_tp4.yaml @@ -1,5 +1,12 @@ model: "qwen3_5_27b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 +# Speculate and pre-plan step N+1 of the TP LLM during step N, as single-GPU +# nodes do (see MSTAR_TP_ASYNC_SCHED); without it each step's schedule and +# plan run serially between the ranks' forwards. +tp_async_sched: true node_groups: - node_names: [vision_encoder] ranks: [0] @@ -12,9 +19,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3_5_2b.yaml b/configs/qwen3_5_2b.yaml index 306a1c9e6d..e8b6ff7f3a 100644 --- a/configs/qwen3_5_2b.yaml +++ b/configs/qwen3_5_2b.yaml @@ -1,5 +1,8 @@ model: "qwen3_5_2b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. @@ -10,9 +13,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3_5_4b.yaml b/configs/qwen3_5_4b.yaml index ea2a32dba4..bf14f9cc66 100644 --- a/configs/qwen3_5_4b.yaml +++ b/configs/qwen3_5_4b.yaml @@ -1,5 +1,8 @@ model: "qwen3_5_4b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. @@ -10,9 +13,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3_5_9b.yaml b/configs/qwen3_5_9b.yaml index 8982639c82..ebcffd7496 100644 --- a/configs/qwen3_5_9b.yaml +++ b/configs/qwen3_5_9b.yaml @@ -1,5 +1,8 @@ model: "qwen3_5_9b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. @@ -10,9 +13,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3_5_9b_tp2.yaml b/configs/qwen3_5_9b_tp2.yaml index b4bac9b89b..0949a72d89 100644 --- a/configs/qwen3_5_9b_tp2.yaml +++ b/configs/qwen3_5_9b_tp2.yaml @@ -1,5 +1,12 @@ model: "qwen3_5_9b" max_seq_len: 32768 +# GIL switch interval of the API server process (see docs/serving.rst): +# token streaming keeps three of its threads busy +api_switch_interval_s: 0.001 +# Speculate and pre-plan step N+1 of the TP LLM during step N, as single-GPU +# nodes do (see MSTAR_TP_ASYNC_SCHED); without it each step's schedule and +# plan run serially between the ranks' forwards. +tp_async_sched: true node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. @@ -14,9 +21,9 @@ resources: max_slots: 129 # state_dtype: bfloat16 # default: the checkpoint's mamba_ssm_dtype (float32) -# The vision encoder compiles on its first image shape and recompiles dynamic -# on the second (~1.5 s each); these two sizes (64 and 608 image tokens) take -# both compiles before the server binds. +# The vision encoder's embed and merger compile (dynamic) on the first image; +# these two sizes (64 and 608 image tokens) take that compile and warm two of +# the captured block loop's buckets before the server binds. warmup_requests: - text: What is in this image? files: [hf://nvidia/Cosmos3-Edge/assets/example_action_fd_umi_first_frame.png] diff --git a/configs/qwen3omni_pd_2gpu.yaml b/configs/qwen3omni_pd_2gpu.yaml new file mode 100644 index 0000000000..261b56e8d5 --- /dev/null +++ b/configs/qwen3omni_pd_2gpu.yaml @@ -0,0 +1,18 @@ +model: "qwen3_omni" +max_seq_len: 32768 +# qwen3omni_pd_disaggregated.yaml on two GPUs: the Talker shares rank 0 with the +# encoders, Code2Wav and the prefill Thinker; rank 1 decodes. +resources: + thinker_kv: + max_num_pages: 1024 +node_groups: + - node_names: [audio_encoder, vision_encoder, Code2Wav] + ranks: [0] + - node_names: [Thinker] + ranks: [0] + graph_walks: [prefill_text, prefill_audio, prefill_vision] + - node_names: [Thinker] + ranks: [1] + graph_walks: [thinker_decode] + - node_names: [Talker] + ranks: [0] diff --git a/configs/qwen3tts_split.yaml b/configs/qwen3tts_split.yaml new file mode 100644 index 0000000000..5376228ddf --- /dev/null +++ b/configs/qwen3tts_split.yaml @@ -0,0 +1,26 @@ +# Qwen3-TTS-12Hz-0.6B-CustomVoice with the codec on its own worker, as +# configs/qwen3tts_1p7b_split.yaml does for the 1.7B. Both workers share GPU 0 +# (rank_devices) as separate processes, so the Talker's steps no +# longer wait on the codec's host work. Start a CUDA MPS daemon for the GPU +# first so their kernels overlap. On an H200 (with MPS) it raised throughput +# ~11% at c=32/64 but cost ~14% at c=1 (RTF 0.072 vs 0.062) and first-audio +# latency at c=64 (~100 vs ~75 ms): a throughput-oriented alternative to +# configs/qwen3tts.yaml, not a replacement. +model: "qwen3_tts" +max_seq_len: 32768 +resources: + talker_attn: + flashinfer_backend: fa2 + # Two workers share the GPU: cap the KV pool at 1024 pages x 128 tokens + # (16 GB; 32 concurrent requests of a few hundred tokens use ~300 pages) + # so the Talker's weights, graphs and pool fit next to the codec worker's. + talker_kv: + max_num_pages: 1024 +node_groups: + - node_names: [Talker] + ranks: [0] + graph_walks: [talker_prefill, talker_decode] + - node_names: [Codec] + ranks: [1] + graph_walks: [codec_chunk] +rank_devices: {1: 0} diff --git a/configs/wan22.yaml b/configs/wan22.yaml index c705dc8cc5..6d4ab843e9 100644 --- a/configs/wan22.yaml +++ b/configs/wan22.yaml @@ -5,13 +5,13 @@ model: "wan22" max_seq_len: 512 # The untiled VAE decode runs the shared Wan VAE decoder, compiled (1.74x faster -# decode at 1280x704 x 121 frames on an H200). The compile is per output -# resolution, not frame count, about 20 s on an H200, so it runs only for the -# listed sizes (default: 1280x704); any other size decodes with stock vae.decode. -# Uncomment to list more sizes, or to restore stock vae.decode everywhere: -# model_kwargs: -# vae_decode_fast_sizes: ["1280x704", "832x480"] -# vae_decode_fast: false +# decode at 1280x704 x 121 frames on an H200; 1.59x at 832x480 x 81). The compile +# is per output resolution, not frame count, about 20 s on an H200, so it runs +# only for the listed sizes, the 720p and 480p tiers here; any other size +# decodes with stock vae.decode. vae_decode_fast: false restores stock +# vae.decode everywhere. +model_kwargs: + vae_decode_fast_sizes: ["1280x704", "832x480"] # Wan2.2-TI2V-5B single-GPU colocated: all four nodes on rank 0. node_groups: diff --git a/configs/wan22_dp2.yaml b/configs/wan22_dp2.yaml new file mode 100644 index 0000000000..16a040e96a --- /dev/null +++ b/configs/wan22_dp2.yaml @@ -0,0 +1,16 @@ +model: "wan22" + +# See wan22.yaml. Two GPUs, throughput-oriented: every node is replicated on +# both, and each request runs end to end on one of them. +max_seq_len: 512 + +model_kwargs: + vae_decode_fast_sizes: ["1280x704", "832x480"] + +node_groups: + - node_names: + - text_encoder + - vae_encoder + - vae_decoder + - dit + ranks: [0, 1] diff --git a/configs/wan22_sp2.yaml b/configs/wan22_sp2.yaml new file mode 100644 index 0000000000..d247cad0e5 --- /dev/null +++ b/configs/wan22_sp2.yaml @@ -0,0 +1,29 @@ +model: "wan22" + +# See wan22.yaml. Two GPUs, latency-oriented: the dit runs on both, each rank +# computing one classifier-free-guidance branch at batch 1, and the two noise +# predictions are all-gathered every step, so a guided denoise step costs one +# branch's forward. The text encoder and both VAE passes run on rank 0. +# For throughput under concurrency, wan22_dp2.yaml (one full replica per GPU) +# avoids rank 1 idling through rank 0's VAE decode. +max_seq_len: 512 + +model_kwargs: + vae_decode_fast_sizes: ["1280x704", "832x480"] + +node_groups: + - node_names: + - text_encoder + ranks: + - 0 + + - node_names: + - vae_encoder + - vae_decoder + ranks: + - 0 + + - node_names: + - dit + ranks: [0, 1] + sp_size: 2 diff --git a/configs/z_image_turbo.yaml b/configs/z_image_turbo.yaml index 1261693bd9..a8182c33eb 100644 --- a/configs/z_image_turbo.yaml +++ b/configs/z_image_turbo.yaml @@ -25,17 +25,14 @@ model_kwargs: # be padded up into a wider bucket without changing the image. A cap_len with no # captured graph falls back to the eager batched path. # - # The capture budget is len(capture_caption_lengths) x len(capture_batch_sizes), - # so all 16 lengths would be 128 graphs. Spend the budget on caption lengths - # instead of batch sizes: a captured bucket already serves SHORTER batches by - # padding them with dummy rows (CudaGraphRunner.DummyRowPool), so 1/2/4/8 covers - # every batch up to 8, where extra cap_lens cover nothing but themselves. - # 4 x 4 = 16 graphs, the same count as the [32, 64] x [1..8] this replaced, with - # twice the caption coverage. Widen if the measured startup and VRAM allow. + # The capture budget is len(capture_caption_lengths) x len(capture_batch_sizes) graphs. + # Every batch size up to max_batch_size is captured: with only [1, 2, 4, 8], a step of + # 3 / 5-7 rows replayed the next size with dummy rows, and under concurrency (requests at + # different steps, one or two outside the dit) that padding cost more than batching saved. + # H200, 1024^2, 8 steps, 2 rounds: c=4 1.05-1.08 -> 1.22 images/s, c=8 1.12-1.14 -> 1.21-1.22, + # c=16 1.17-1.19 -> 1.23, c=1 unchanged, same peak VRAM; ~100 s more warm startup (32 graphs). capture_caption_lengths: [32, 64, 96, 128] - # The scaffold default; 3/5/6/7 replay the next captured size with dummy rows. - # Costs some padded compute at those sizes, buys half the graphs. - capture_batch_sizes: [1, 2, 4, 8] + capture_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] max_batch_size: 8 # compiled VAE decode (max-autotune): 3x faster at 1024^2, ~56 dB from the eager decode; the autotuner may pick # other conv kernels in another server process (~64 dB apart). vae_compile: false for bit-exact, repeatable output diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index ef4c43662f..c8fb526207 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -514,11 +514,6 @@ Worker scheduling decoder) no longer pays the eager-step master write, a small synchronous host-to-device copy per step. ``0`` keeps every node on the per-request tensor path. - * - ``MSTAR_INGRAPH_DECODE_ROPE`` - - ``0`` - - A captured decode step builds its rotary tables inside the graph from - the position resource's planned positions instead of staging them. - Experiment knob, measured and gated on its own. * - ``MSTAR_SAMPLER_INGRAPH_SCATTER`` - ``0`` - The captured sampler gathers its RNG offsets from the slot masters diff --git a/docs/installation.rst b/docs/installation.rst index 09fd7d13ae..4f927bc53b 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -362,7 +362,9 @@ flash-attn-4 (Waypoint, FA3 sm90 kernel) ---------------------------------------- Waypoint's DiT attention runs torch flex-attention with the ``FLASH`` backend by -default. That backend needs the **flash-attn-4** package, which provides +default. LingBot-World's ``configs/lingbot_world_1_3b_fast.yaml`` (``reference_numerics: +false``) also calls it directly. Install it after ``flash-attn`` (FA2): the FA2 package +ships an older ``flash_attn/cute`` that overwrites flash-attn-4's. That backend needs the **flash-attn-4** package, which provides ``flash_attn.cute`` — the CuTe DSL rewrite of flash-attn; on H100 it runs the FA3-style sm90 kernel with TMA and warpgroup specialisation. It is **not on PyPI** as of 2026-09-17, and it is **not** pulled in by any extra. diff --git a/docs/models.rst b/docs/models.rst index 481177339f..4114265d2a 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -171,6 +171,35 @@ OmniVoice notes - The backbone is not autoregressive: it fills a fixed canvas of eight codebook rows over a few unmasking steps, so there is no KV cache and no per-token sampling loop. Serve it with ``mstar serve omnivoice --gpus 0``. +- The backbone's layer stack replays CUDA graphs per (requests, packed-token + bucket up to 4096 tokens) with attention through a ragged attention resource; + larger steps are GPU-bound and run eagerly. The layer loop is mstar's own over + the reference's fused weights: q/k norms read the fused QKV output in place + and each residual add folds into the next RMSNorm (flashinfer + ``fused_add_rmsnorm``), so logits match the reference to half-precision + rounding rather than bit for bit. ``model_kwargs: {backbone_graphs: false}`` + keeps the reference's eager packed path. +- ``configs/omnivoice_split.yaml`` runs code2wav (codec decode and the + host-side silence trim) on a second worker sharing GPU 0, so the unmask loop + no longer stops for each finished request. With a CUDA MPS daemon it raised + concurrency-32 throughput on an H200 from ~153 to 176-184 audio-s/s and cut + TTFA p95 from ~9 s to ~3 s; concurrency 1 is ~10% slower. + +Higgs Audio notes +~~~~~~~~~~~~~~~~~ + +- The audio tower runs outside ``torch.compile`` (a clip's chunk count changes + its shape, and every new shape recompiled on the serving path). Requests + whose 4 s chunks share a mel width are encoded together, and full-width + chunks replay a CUDA graph per chunk count; a clip under one chunk is padded + to the full width and masked (``HiggsAudioTower``'s ``valid_frames``) so it + replays a graph too. ``model_kwargs: {encoder_graphs: false}`` runs the + tower eagerly. +- One prefill step per request: the ChatML text around the audio and the audio + embeddings go into the KV cache as one span. +- Decode keeps each request's last token in the sampler's slot master and + sends the client's copy inline (``MSTAR_DEVICE_LOOPBACK``, on by default), + as Qwen3.5 does: a decode step stores and routes no tensors. Omnilingual ASR notes ~~~~~~~~~~~~~~~~~~~~~ @@ -205,7 +234,10 @@ Omnilingual ASR notes - ``dtype`` defaults to ``bfloat16``, the reference's inference dtype; ``float32`` also turns off TF32 for the node's matmuls. ``max_batch_size`` (16) caps the requests in one step; the step sorts them by length and splits them so that each forward's padded audio - (rows times the longest clip) stays within ``max_batch_seconds`` (320). In bf16 a + (rows times the longest clip) stays within ``max_batch_seconds`` (320). In bf16 the + encoder packs a step's clips into one unpadded forward (a captured graph up to 16384 + frames) when their total audio fits ``max_batch_seconds``; a step over it encodes padded, + per split. In bf16 a request batched with others can differ from the same request run alone where two characters nearly tie, as the reference's own batched ``transcribe`` does; ``max_batch_size: 1`` or ``dtype: float32`` give batch-invariant transcripts. @@ -213,6 +245,36 @@ Omnilingual ASR notes (``record_oracle.py`` runs the reference in its own venv, as it pins torch 2.8 and fairseq2). +BAGEL notes +~~~~~~~~~~~ + +- Text-to-image runs only the CFG branches a denoising step reads + (``cfg_prune``, on by default): two with image guidance off, and the main + branch alone outside ``cfg_interval``, about -37% latency at 1024² on an + H200. Outside the interval the step now returns the main branch's velocity + where it used to recombine all three and renormalize, which is equal only up + to rounding, so images move slightly (median 41 dB PSNR against the previous + output). ``model_kwargs: {cfg_prune: false}`` in the serving config restores + the always-three-branch math. + +Qwen3-Omni notes +~~~~~~~~~~~~~~~~ + +- The Talker -> Code2Wav edge sends a first chunk of 8 codec frames, then 25 + (``codec_chunk_schedule``), which halves time to first audio. The vocoder's + attention window is longer than the 25-frame left context, so streamed audio + was never the whole-utterance decode, and moving the chunk boundaries changes + the waveform (served WER measured unchanged). An empty + ``Code2WavConfig.codec_chunk_schedule`` restores fixed 25-frame chunks. +- Numerics: the code predictor's residual adds fold into the next RMSNorm + (FlashInfer's fused add + norm) and its attention is a Triton kernel + (``cp_attention``), neither bit-exact with the unfused ops. Its 15 depth + steps feed each sampled code into the next, so small differences can change + later codes at near-ties. +- The MoE layers route through the fused softmax-top-k kernel + (``TopKRouter(fused_softmax_topk=True)``), which gives equal probabilities to + the lower expert id where ``torch.topk`` may pick another. + Qwen3.5 / Qwen3.8 (``qwen3_5_*``, ``qwen3_8_27b``) -------------------------------------------------- @@ -267,6 +329,12 @@ not the 256-slot default. The state defaults to the checkpoint's ``gdn_state.state_dtype: bfloat16`` halves state traffic at some precision and needs FlashInfer's fused bf16 decode kernel (K = V = 128). +A captured decode step builds its rotary tables inside the graph from the +planned positions, rather than staging tables built eagerly before the replay +(that build sat on the per-step critical path: +13% output tok/s at c=32 on +0.8B, H200). Greedy output can differ from earlier releases at near-tie +tokens; prefill and eager steps are unchanged. + ``temperature``, ``top_p``, ``max_tokens`` and ``seed`` are the standard fields. ``repetition_penalty`` and ``enable_thinking`` (default true; the template opens a ```` block) are read by the model but are not OpenAI fields — pass them @@ -311,8 +379,12 @@ Kokoro notes graphs captured at start-up (fewer buckets on a smaller GPU), ``compile_decoder: false`` skips the ``torch.compile`` of the vocoder (start-up in seconds instead of minutes, about half the throughput on an H100) and ``decoder_dtype: bfloat16`` runs the vocoder trunk - in bf16 (about 13% more throughput at concurrency 32 in our runs, with the harmonic - source and the iSTFT kept in fp32; the parity test covers fp32 only). + in bf16 (about 13% more throughput at concurrency 32 on an H100 and 20-25% at + concurrency 32-64 on an H200, Whisper WER unchanged, with the harmonic source and the + iSTFT kept in fp32; the parity test covers fp32 only; ``configs/kokoro_bf16.yaml``). + ``fuse_lstm_directions`` (default on) runs each BiLSTM's two directions as one + block-diagonal recurrence, and ``decoder_cudnn_benchmark`` (default on) lets cuDNN + time the frame half's conv algorithms once per captured bucket. - Text is cut at sentence boundaries into chunks of at most 510 phonemes (the PL-BERT window); each chunk is emitted to the client as soon as it is synthesized, so ``stream=True`` on ``/v1/audio/speech`` returns audio sentence by @@ -364,25 +436,29 @@ Qwen3-TTS notes transformer stacks 8 such layers, so a window still differs slightly from a whole-utterance decode once an utterance outgrows the context; 25 frames of context differed far more. Each window size is a CUDA-graph bucket captured - for batch sizes 1 to 32; the stream buffer reports how many leading frames of - a window are repeated context, and the codec trims their audio. A voice - clone's reference tail (its last 72 frames) is the first window's context. + for batch sizes 1 to 32 (densely: the codec's cost is per row, so a padding + row costs as much as a real one); the stream buffer reports how many leading + frames of a window are repeated context, and the codec trims their audio. The + scheduler batches each window size apart, so a stream's small first windows + are not padded to a full window's width or queued behind a batch of them. A + voice clone's reference tail (its last 72 frames) is the first window's context. - The codec runs its transformer over the whole window but its conv stack (all but a few percent of the codec's time) only over the new frames plus the stack's causal receptive field, derived from the decoder's modules at load (10 frames for the 12 Hz decoder; the new frames' audio equals the whole-window decode). It also fuses the SnakeBeta activations into one kernel. -- Talker prefill replays a packed CUDA Graph for the smallest token bucket - (32 to 1024 tokens) that holds the batch; only the clone prefill, which also - pushes the reference clip's frames into the codec stream, runs eager. Decode - always uses the whole-walk CUDA Graph, with the 15-step +- Talker prefill, the clone prefill included, replays a packed CUDA Graph for + the smallest token bucket (32 to 1024 tokens) that holds the batch; a clone's + reference frames join the codec stream in ``postprocess``. Decode + always uses the whole-walk CUDA Graph (up to 64 rows a step), with the 15-step CodePredictor loop captured inside it; request-local EOS suppression is carried as a graph tensor input so replay does not consult capture-slot dummy request state. Residual ``subtalker_*`` sampling is per-request through the ``code_predictor`` aux sampler, so custom values neither block batching nor fall off the graph. On the 1.7B checkpoints the CodePredictor projects the Talker-width inputs through ``small_to_mtp_projection`` before its depth loop. -- ``configs/qwen3tts_1p7b_split.yaml`` runs the Codec on a second worker that +- ``configs/qwen3tts_1p7b_split.yaml`` (and ``configs/qwen3tts_split.yaml`` + for the 0.6B) runs the Codec on a second worker that shares GPU 0 (``rank_devices``) and caps the Talker KV pool so both fit. On its own it lowers first-audio latency at high concurrency; with a user-level CUDA MPS daemon (``nvidia-cuda-mps-control -d`` before @@ -479,9 +555,13 @@ Fish Audio S2-Pro notes probability carries the running total past top-p, while the reference drops it unless it is the most likely token. ``max_new_tokens`` (default 2048, at most 4096 frames of 2048 samples at 44.1 kHz) and ``ignore_eos`` pass through. -- Streamed audio matches the whole-utterance decode to within one int16 step: - each codec window carries the decoder's whole receptive field as context, and - the codec runs in fp32 with TF32 off. +- Each codec window carries the decoder's whole receptive field as context. + With ``model_kwargs: {codec_exact_fp32: true}`` the codec runs in fp32 with + TF32 off and the streamed audio matches the whole-utterance decode to within + one int16 step. The shipped ``configs/fish_s2_pro.yaml`` sets it to ``false``: + TF32 halves the codec's conv stack, most of the GPU time under concurrency + (c=32 throughput +35% on an H200), and the audio stays within ~60 dB SNR of + the exact decode. Nemotron VoiceChat (``nemotron_duplex``) notes ---------------------------------------------- @@ -639,6 +719,16 @@ Chatterbox notes window stays as close to the whole-utterance decode as the full history does (log-mel correlation 0.988 vs 0.985 on CPU). Requests whose chunks are ready together share one padded flow solve (up to 8 per step). +- ``configs/chatterbox_split.yaml`` (and ``chatterbox_turbo_split.yaml``, + ``chatterbox_multilingual_split.yaml``) runs S3Gen on a second worker that + shares GPU 0 (``rank_devices``), so T3's decode steps stop alternating with + S3Gen's flow solves on one worker loop. It only pays with a user-level CUDA + MPS daemon (``nvidia-cuda-mps-control -d`` before ``mstar serve``), which + lets the two workers' kernels overlap: on an H200 at concurrency 32 Chatterbox + went from 36-38 to 58-61 audio-s/s and Turbo from 69 to 94, at the cost of a + later first chunk (TTFA p50 1.6-1.8 s -> 1.9-2.1 s; Turbo 0.36 s -> 1.0 s); + concurrency 1 is unchanged, and WER held (1.22% -> 1.08%). Without MPS the + split measured the same as one worker. - S3Gen runs from CUDA graphs by default (``s3gen_graphs``): the flow solve (one graph per rows x frames x steps; the estimator's hundreds of tiny kernels per Euler step make the eager solve launch-bound, 170 ms vs 45 ms @@ -716,6 +806,19 @@ Cosmos3 environment requirements - The Wan-VAE decode dtype is gated on the cuDNN build: bf16 needs cuDNN >= 9.16 (fast Hopper bf16 conv3d); older cuDNN serves the decode in fp32/TF32 automatically. +- Whole-clip decodes (``vae_decode_fast``, default on) run the shared Wan VAE + decoder channels-last and compiled per latent, on the decoder node's own + decoder-only copy of the VAE: 17% off the decode on an H200 (Nano 832x480x189 + 2.9 -> 2.4 s, Edge 832x480x121 1.84 -> 1.54 s), within ~40 dB PSNR of the stock + decode. ``vae_decode_fast: false`` restores the compiled stock ``vae.decode``. +- Two GPUs: ``configs/cosmos3_nano_sp2.yaml`` halves one request's denoise + (Nano 832x480x189x35: 1.47 -> 0.78 s per step), while + ``configs/cosmos3_nano_dp2.yaml`` (a whole replica per GPU) serves more under + concurrency: 2.13 videos/min at 2 concurrent vs 1.63 for sp2 and 1.07 on one GPU. + For Edge (832x480x121, 20 steps): ``configs/cosmos3_edge_sp2.yaml`` takes one + request from 6.3 to 4.4 s (generation only: the reasoner has no SP path), and + ``configs/cosmos3_edge_dp2.yaml`` serves 15.8 videos/min at 4 concurrent vs 7.9 on + one GPU. Cosmos3-Edge reasoner and action loop ------------------------------------- @@ -942,12 +1045,18 @@ travels with the request rather than living in a scheduler object on one rank. Requests are therefore independent and the loop is resumable across ranks. **What is accelerated by default.** The DiT runs under ``torch.compile`` -(``compile_dit``) and the untiled VAE decode runs the shared Wan VAE decoder, +(``compile_dit``); its self-attention runs on the engine's ragged attention +resource (FlashInfer, FA3 on Hopper) as an opaque op inside the compiled graph, +6% off a 1280x704x81 denoise step on an H200 against cuDNN SDPA +(``model_kwargs: attention_backend: sdpa`` restores the reference kernel; the +cross-attention over the 512 text tokens stays on SDPA, which is faster at that +shape). The untiled VAE decode runs the shared Wan VAE decoder, channels-last and compiled (``vae_decode_fast``; set it to ``false`` under the yaml's ``model_kwargs`` for stock ``vae.decode``). Both compile on first use: the DiT once per resolution, frame count and CFG setting, so the first request at a new size is slower. The VAE decode compiles once per resolution, and only -for the sizes in ``vae_decode_fast_sizes`` (default ``1280x704``); any other +for the sizes in ``vae_decode_fast_sizes`` (model default ``1280x704``; the +shipped configs add ``832x480``, where it decodes 1.59x faster); any other size decodes with stock ``vae.decode``, frames identical to ``vae_decode_fast: false``. The compile uses Inductor's deterministic mode, so every server process decodes the same latents to the same frames; torch 2.9.1 @@ -958,6 +1067,17 @@ tiles only when free VRAM at decode time is short of the untiled estimate 32 GiB card. No CUDA-graph capture, no continuous batching (one request per step) and no component offload. +**Two GPUs.** ``configs/wan22_sp2.yaml`` runs the dit as a two-rank lockstep +instance (``sp_size: 2``) that splits the classifier-free guidance batch: each +rank runs one branch at batch 1 and the two noise predictions are all-gathered +every step, so a guided step costs one branch's forward (502 vs 1000 ms at +1280x704x81 on H200s; one request 13.4 s vs 23.3 s end to end at 20 steps). The +text encoder and VAE run on rank 0. Without guidance both ranks compute the same +step, correct but no faster. For throughput under concurrency, +``configs/wan22_dp2.yaml`` serves a whole replica per GPU instead: 5.2 videos/min +at 4 concurrent vs 4.7 for sp2 (sp2's rank 1 idles through rank 0's VAE +decode), while one request takes the 1-GPU time. + LTX-2.5 (``ltx2_5``) -------------------- diff --git a/docs/serving.rst b/docs/serving.rst index 40a4ecce7c..dc656b6543 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -261,8 +261,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key ``hf:////`` (a file in a model's HF repo); its modality comes from the extension, or give ``{path: ..., modality: ...}``. ``repeat: N`` sends ``text`` N times over, a long prompt for a path only - long prompts take (a long-context prefill). A malformed or failing warmup - is logged and skipped. + long prompts take (a long-context prefill). ``concurrency: N`` submits N + copies at once, so they batch, for a step that compiles per batch size. A + malformed or failing warmup is logged and skipped. * - ``dist_timeout_s`` - *(optional)* Timeout in seconds for the NCCL world group and its parallel subgroups during weight load and warmup; serving keeps PyTorch's default. @@ -284,6 +285,14 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key - *(optional)* Largest message, in KiB, the symmetric-memory path takes; larger ones go to NCCL. It is also the size of the buffer each group rendezvouses. Default 512. + * - ``api_switch_interval_s`` + - *(optional)* GIL switch interval of the API server process, in seconds + (``sys.setswitchinterval``); unset keeps Python's 5 ms. Under streaming load + that process keeps three threads busy (event loop, message thread, data + worker), and a shorter interval hands the GIL between them sooner. The + Qwen3.5 configs set ``0.001``: Qwen3.5-0.8B text, c=64, +4 to +8% tok/s. A + model whose preprocessing runs long on the data-worker thread can lose from + it (Kokoro's G2P: about -3.6% at c=32), so it is per deployment. * - ``tp_async_sched`` - *(optional)* Async scheduling for the parallel nodes: ``true``, ``false`` (default) or a list of node names (see ``MSTAR_TP_ASYNC_SCHED``, which diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 1dd6b41354..11f3a35e50 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -1325,12 +1325,13 @@ class WarmupRequest: files: list[WarmupMedia] = field(default_factory=list) output_modalities: list[str] = field(default_factory=lambda: ["text"]) model_kwargs: dict | None = None + concurrency: int = 1 @classmethod def from_config(cls, spec: Any) -> "WarmupRequest": if not isinstance(spec, dict): raise ValueError("not a mapping") - unknown = set(spec) - {"text", "files", "output_modalities", "model_kwargs", "repeat"} + unknown = set(spec) - {"text", "files", "output_modalities", "model_kwargs", "repeat", "concurrency"} if unknown: raise ValueError(f"unknown key(s) {sorted(unknown)}") text = spec.get("text") @@ -1341,6 +1342,9 @@ def from_config(cls, spec: Any) -> "WarmupRequest": raise ValueError(f"repeat must be a positive integer, got {repeat!r}") if repeat > 1 and not text: raise ValueError("repeat needs text") + concurrency = spec.get("concurrency", 1) + if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: + raise ValueError(f"concurrency must be a positive integer, got {concurrency!r}") out_mods = spec.get("output_modalities") or ["text"] if isinstance(out_mods, str): out_mods = [m.strip() for m in out_mods.split(",") if m.strip()] @@ -1355,6 +1359,7 @@ def from_config(cls, spec: Any) -> "WarmupRequest": files=[WarmupMedia.from_config(f) for f in files], output_modalities=list(out_mods), model_kwargs=model_kwargs, + concurrency=concurrency, ) def submit(self, server: APIServer, request_id: str) -> str: @@ -1380,13 +1385,29 @@ def submit(self, server: APIServer, request_id: str) -> str: ) -def _run_warmup_requests(server: APIServer, specs: list) -> None: +def _data_parallel_replicas(config: dict) -> int: + """The most replicas any node group serves (its ranks over its tp x sp + instance size): how many copies of a request it takes to reach them all.""" + replicas = 1 + for group in config.get("node_groups") or []: + instance = int(group.get("tp_size", 1)) * int(group.get("sp_size", 1)) + replicas = max(replicas, len(group.get("ranks") or []) // max(instance, 1)) + return replicas + + +def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> None: """Run the deployment's ``warmup_requests`` once the workers are ready and before the server binds, so the first client does not pay what a first request of a shape costs (torch.compile of a denoise step, a cold inductor cache). Each spec is parsed into a :class:`WarmupRequest`. A warmup that is malformed or fails is logged and skipped; it must not keep the server from coming up. + + Each spec runs as ``copies`` x its ``concurrency`` concurrent requests: + ``copies`` warms every replica of a data-parallel deployment, since the + conductor sends concurrent requests to different replicas (the + least-loaded one each), and ``concurrency`` lets them batch, for a step + that compiles per batch size. """ if not isinstance(specs, list): logger.warning("warmup_requests must be a list; ignoring %r", type(specs).__name__) @@ -1395,6 +1416,16 @@ def _run_warmup_requests(server: APIServer, specs: list) -> None: async def _one(request: WarmupRequest, request_id: str) -> None: await server.collect_results(request.submit(server, request_id)) + async def _all(request: WarmupRequest, index: int) -> None: + n = copies * request.concurrency + ids = [f"warmup-{index}"] + [f"warmup-{index}-{c}" for c in range(1, n)] + # every copy runs to its end before a failure is reported, so none is + # cancelled mid-request and still in the server when the next spec starts + results = await asyncio.gather(*(_one(request, rid) for rid in ids), return_exceptions=True) + errors = [r for r in results if isinstance(r, BaseException)] + if errors: + raise errors[0] + for index, spec in enumerate(specs): try: request = WarmupRequest.from_config(spec) @@ -1403,7 +1434,7 @@ async def _one(request: WarmupRequest, request_id: str) -> None: continue t0 = time.perf_counter() try: - asyncio.run(_one(request, f"warmup-{index}")) + asyncio.run(_all(request, index)) except Exception as exc: # noqa: BLE001 — a cold server is better than none logger.warning( "warmup request %d/%d failed after %.1f s: %s", @@ -1920,6 +1951,19 @@ def _set_preprocess_threads(model, config: dict) -> None: torch.set_num_threads(int(threads)) +def _set_switch_interval(config: dict) -> None: + """``api_switch_interval_s`` from the config: the GIL switch interval of + this process only (the conductor and workers are spawned fresh). Unset + keeps Python's 5 ms. Streaming many small chunks keeps three of this + process's threads busy (event loop, message thread, data worker); a + shorter interval hands the GIL between them sooner, which also preempts a + model's preprocessing on the data-worker thread more often.""" + interval = config.get("api_switch_interval_s") + if interval is not None: + sys.setswitchinterval(float(interval)) + logger.info("API server GIL switch interval: %s s", interval) + + def main(argv: list[str] | None = None): import argparse @@ -2030,6 +2074,7 @@ def main(argv: list[str] | None = None): **{**model_init_kwargs(model_name), **yaml_model_kwargs}, ) _set_preprocess_threads(model, config) + _set_switch_interval(config) global api_server log_stats = args.log_stats or args.log_stats_file is not None @@ -2080,8 +2125,12 @@ def main(argv: list[str] | None = None): # client never sees the first request of a shape. warmups = config.get("warmup_requests") or [] if warmups: - logger.info("Running %d warmup request(s) before binding", len(warmups)) - _run_warmup_requests(api_server, warmups) + copies = _data_parallel_replicas(config) + logger.info( + "Running %d warmup request(s) before binding, %d at a time (one per replica)", + len(warmups), copies, + ) + _run_warmup_requests(api_server, warmups, copies) if args.enable_metrics: api_server.enable_metrics() if args.rust_frontend: diff --git a/mstar/communication/tensor_store.py b/mstar/communication/tensor_store.py index e2699149e7..bd444dfff9 100644 --- a/mstar/communication/tensor_store.py +++ b/mstar/communication/tensor_store.py @@ -674,6 +674,15 @@ def _own( if cpu_tensor is not None: self._tensors_cpu[uuid] = cpu_tensor + def drop_cpu_tensors(self, uuids: list[int]) -> None: + """Forget the host copies stored for ``uuids``. They are views of the + producer's pinned buffers, which its next step overwrites, so a send + after that (a tensor held across steps) must read the device tensor.""" + if self._tensors_cpu: + pop = self._tensors_cpu.pop + for uuid in uuids: + pop(uuid, None) + def check_uuid_presence(self, uuid: int) -> bool: return uuid in self._tensors diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index e6bf14ba3e..59237008c0 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -1532,6 +1532,11 @@ def get_rx_info(self, request_id: Rid) -> list[RxInfo]: # ---- shared: TensorStore delegation ---- + def release_host_copies(self, uuids: list[int]) -> None: + """End of the step that stored ``uuids``: their host copies (views of + the worker's reused pinned buffers) may not be sent from after this.""" + self.tensor_store.drop_cpu_tensors(uuids) + def get_tensor(self, uuid: int) -> torch.Tensor: tensor = self.tensor_store.get_tensor(uuid) shard_dim = self.uuid_to_shard_dim.get(uuid) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 83584c424f..ba9d893806 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -816,6 +816,26 @@ def _assign_worker_graphs_to_workers(self) -> dict[int, list[str]]: """ # _group_id -> chosen DP-replica index within that group's ranks group_id_to_replica_idx: dict[int, int] = {} + # worker id -> requests it currently serves; built only when a group + # actually has more than one replica to choose from + load: dict[str, int] | None = None + + def pick(replicas: list[list[int]]) -> int: + """The least-loaded replica (by live requests on its workers), ties at + random. A random pick alone puts two of two concurrent requests on one + replica half the time, leaving the other idle.""" + nonlocal load + if len(replicas) == 1: + return 0 + if load is None: + load = defaultdict(int) + for request_data in self.requests.values(): + for worker_id in self._request_workers(request_data): + load[worker_id] += 1 + loads = [sum(load[f"worker_{r}"] for r in ranks) for ranks in replicas] + least = min(loads) + return int(np.random.choice([i for i, n in enumerate(loads) if n == least])) + result = {} for wg_id, wg in self.worker_graphs.items(): if wg._instance_ranks: @@ -823,16 +843,14 @@ def _assign_worker_graphs_to_workers(self) -> dict[int, list[str]]: # tp*sp instance runs the request together, so its TP all-reduce # and SP all-to-all collectives stay in sync. Picking a single TP # row here would desync the SP all-to-all across rows. - replica_idx = group_id_to_replica_idx.setdefault( - wg._group_id, np.random.randint(len(wg._instance_ranks)), - ) - ranks = wg._instance_ranks[replica_idx] + if wg._group_id not in group_id_to_replica_idx: + group_id_to_replica_idx[wg._group_id] = pick(wg._instance_ranks) + ranks = wg._instance_ranks[group_id_to_replica_idx[wg._group_id]] result[wg_id] = [f"worker_{r}" for r in ranks] else: - replica_idx = group_id_to_replica_idx.setdefault( - wg._group_id, np.random.randint(len(wg.ranks)), - ) - result[wg_id] = [f"worker_{wg.ranks[replica_idx]}"] + if wg._group_id not in group_id_to_replica_idx: + group_id_to_replica_idx[wg._group_id] = pick([[r] for r in wg.ranks]) + result[wg_id] = [f"worker_{wg.ranks[group_id_to_replica_idx[wg._group_id]]}"] return result def _build_request_sharding_config( diff --git a/mstar/engine/accelerator_graph_config.py b/mstar/engine/accelerator_graph_config.py index bbc33a55a3..82bd359941 100644 --- a/mstar/engine/accelerator_graph_config.py +++ b/mstar/engine/accelerator_graph_config.py @@ -299,6 +299,10 @@ class PiecewiseAcceleratorGraphConfig(ABC): # (`ModelInputsFromEngine.step` is the declaration that tells it so). # None: every batch that fits a bucket. can_lease: Callable[[list[NodeInputs]], bool] | None = None + # The region replays a whole step as ONE row, however many requests the + # step carries (a vision tower packing every request's images as segments + # of one ragged batch), so its slot is leased at batch size 1. + one_row_per_step: bool = False # static kwargs threaded into capture_fn (e.g. cond_tokens, is_causal) forward_kwargs: dict[str, Any] = field(default_factory=dict) # None => defer to the runner's default batch-size buckets diff --git a/mstar/engine/accelerator_graph_runner.py b/mstar/engine/accelerator_graph_runner.py index c3743f89d7..c4f15e6d86 100644 --- a/mstar/engine/accelerator_graph_runner.py +++ b/mstar/engine/accelerator_graph_runner.py @@ -2127,6 +2127,11 @@ def can_lease(self, inputs: list[NodeInputs]) -> bool: """The config's ``can_lease`` for this batch; True when it sets none.""" return self._config.can_lease is None or self._config.can_lease(inputs) + @property + def one_row_per_step(self) -> bool: + """Whether a replay carries the whole step as one row.""" + return self._config.one_row_per_step + def lease_slot( self, batch_size: int, total_tokens: int | None = None, ) -> SlotLease | None: diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 0ef8b9284e..ce62045d58 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1482,7 +1482,7 @@ def _maybe_lease_piecewise_regions( # a batch the region refuses runs it eagerly, with no slot if not runner.lease_before_step or not runner.can_lease(inputs): continue - lease = runner.lease_slot(bs, total_tokens) + lease = runner.lease_slot(1 if runner.one_row_per_step else bs, total_tokens) if lease is not None: leases[label] = lease if leases: diff --git a/mstar/engine/resources/attn/ragged/base.py b/mstar/engine/resources/attn/ragged/base.py index 02ebe9ce22..14bee191ba 100644 --- a/mstar/engine/resources/attn/ragged/base.py +++ b/mstar/engine/resources/attn/ragged/base.py @@ -27,7 +27,7 @@ def resolve_ragged_attn_backend( def _build_ragged(spec, info: EngineResourceInfo, manager_cls): config = spec.config - if info.joint_comm_group is not None: + if info.joint_comm_group is not None and config.shard_heads: config.shard(info.joint_comm_group.world_size) dtype = config.dtype if config.dtype is not None else info.kv_dtype if dtype is None: diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 027614b49d..2c597fc4c6 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -73,6 +73,13 @@ class RaggedAttentionConfig: backend: RaggedAttnBackend = RaggedAttnBackend.AUTO flashinfer_backend: str = "auto" + # Whether the node's lockstep instance splits the attention heads across its + # ranks (tensor parallelism, Ulysses sequence parallelism), so each rank's + # resource attends its slice. A node that parallelizes over something else -- + # each rank attending whole requests or guidance branches -- sets False and + # keeps every head on every rank. + shard_heads: bool = True + # Activation dtype of the q / k / v the node hands the kernel (and of its # output). A cacheless attention doesn't need to adhere to the dtype of its # KV cache, so it is free to declare a separate dtype here; when None, the diff --git a/mstar/engine/resources/convenience.py b/mstar/engine/resources/convenience.py index 9413ed29db..66e4f5934b 100644 --- a/mstar/engine/resources/convenience.py +++ b/mstar/engine/resources/convenience.py @@ -145,17 +145,54 @@ class RaggedAttentionCallable: until Dynamo gives up. """ - def __init__(self, attn: RaggedAttnManager | RaggedCrossAttnManager, label: str): + def __init__(self, attn: RaggedAttnManager | RaggedCrossAttnManager, label: str, opaque: bool = False): self.attn = attn self.label = label + # ``opaque``: call through the ``mstar::ragged_attention`` custom op, which + # a compiled region keeps in its graph as one node. Otherwise the call is a + # graph break per layer, which also stops Inductor fusing the pointwise + # work on either side of it: on Wan2.2's compiled, uncaptured DiT the + # breaks cost ~85 ms of a ~1 s step in extra elementwise kernels. + self._handle = _register_ragged(self) if opaque else None - @torch.compiler.disable def __call__( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: + if self._handle is not None: + return _ragged_attention_op(q, k, v, self._handle) + return self._run(q, k, v) + + @torch.compiler.disable + def _run(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return self.attn.run(q, k, v, label=self.label) +# Opaque callables by handle: a custom op takes tensors and scalars only, so it +# names its callable by index. One entry per callable built, i.e. per resource +# binding and label, which is a handful per process. +_RAGGED_CALLABLES: list[RaggedAttentionCallable] = [] + + +# Eager even when the callable is built inside a compiled region (a node that +# creates its callables lazily on first forward): traced, the append is not +# replayed but the handle Dynamo read is still baked into the graph. +@torch.compiler.disable +def _register_ragged(fn: RaggedAttentionCallable) -> int: + _RAGGED_CALLABLES.append(fn) + return len(_RAGGED_CALLABLES) - 1 + + +@torch.library.custom_op("mstar::ragged_attention", mutates_args=()) +def _ragged_attention_op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, handle: int) -> torch.Tensor: + # the wrapper returns a fresh tensor (or a clone of its graph buffer), never an alias + return _RAGGED_CALLABLES[handle]._run(q, k, v).to(q.dtype) + + +@_ragged_attention_op.register_fake +def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, handle: int) -> torch.Tensor: + return torch.empty_like(q) + + class Mamba2Callable: """``LinearAttnCallable`` for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor. diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index c533582451..807910da15 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -29,14 +29,12 @@ processed causally, then each image is processed bidirectionally. """ -import io import json import logging from pathlib import Path import torch from huggingface_hub import hf_hub_download, snapshot_download -from PIL import Image from torch import nn from mstar.communication.tensors import NameToTensorList @@ -75,6 +73,7 @@ from mstar.model.bagel.config import MIXED_TEXT, load_bagel_config from mstar.model.bagel.submodules import ( IMAGE_ADVANCE, + IMAGE_GEN_LOOP, IMAGE_SENTINELS, CombineCFGSubmodule, LLMSubmodule, @@ -85,6 +84,7 @@ vit_preprocess, ) from mstar.model.base import DECODE, ForwardPassArgs, Model, PrefixStream, ProcessPromptOutput, Span +from mstar.model.components.diffusion.image_io import encode_image from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -267,10 +267,13 @@ def __init__( model_path_hf: str, cache_dir: str | None = None, accelerator_graph: bool | None = None, + cfg_prune: bool | None = None, **kwargs ): if accelerator_graph is not None and not isinstance(accelerator_graph, bool): raise ValueError("accelerator_graph must be a boolean") + if cfg_prune is not None and not isinstance(cfg_prune, bool): + raise ValueError("cfg_prune must be a boolean") self.cache_dir = cache_dir config_path = hf_hub_download( @@ -281,6 +284,8 @@ def __init__( self.config = load_bagel_config(json.load(f)) if accelerator_graph is not None: self.config.accelerator_graph = accelerator_graph + if cfg_prune is not None: + self.config.cfg_prune = cfg_prune self.model_path_hf = model_path_hf @@ -776,11 +781,9 @@ def postprocess( logger.debug("OUTPUT TEXT %s", detok) return detok.encode("utf-8") if modality == "image": - output = output[0].permute(1, 2, 0) * 255 - img = Image.fromarray((output).to(torch.uint8).cpu().numpy()) - img_byte_arr = io.BytesIO() - img.save(img_byte_arr, format='PNG') - return img_byte_arr.getvalue() + # the same truncated pixels as before, through the shared writer: PIL's default PNG + # encode (adaptive filters, zlib 6) took ~0.4 s of every 1024^2 request + return encode_image((output[0] * 255).to(torch.uint8), request_kwargs) raise ValueError(f"Unsupported modality: {modality!r}") # CFG-parallel deployments add the two guidance branches. Ordinary @@ -1062,6 +1065,9 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: # there are N-1 intervals, so N-1 Euler steps are needed. image_gen = Sequential([ Loop( + # named: LLMSubmodule reads this loop's iteration count to + # pick the step's CFG branches on the host (cfg_branches) + name=IMAGE_GEN_LOOP, section=GraphNode( name="LLM", input_names=["latents", "time_index"], @@ -1098,6 +1104,8 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: # combine_cfg applies the CFG formula + Euler step after each iteration. image_gen_cfg = Sequential([ Loop( + # named, as image_gen's: every node of the step reads the counter (cfg_branches) + name=IMAGE_GEN_LOOP, section=Sequential([ Parallel([ GraphNode( diff --git a/mstar/model/bagel/components/vit_encoder.py b/mstar/model/bagel/components/vit_encoder.py index a774deeaae..9d4c858cc0 100644 --- a/mstar/model/bagel/components/vit_encoder.py +++ b/mstar/model/bagel/components/vit_encoder.py @@ -45,6 +45,14 @@ def _sdpa_varlen( # q/k/v: (total_seq_len, num_heads, head_dim). Build a block-diagonal mask # so packed images don't attend across boundaries, then run SDPA. total_len, num_heads, head_dim = q.shape + if cu_seqlens.numel() == 2 and not causal: + # One image (the unbatched ViT path): nothing to isolate, and without a mask SDPA can + # take its flash / cuDNN kernels instead of the masked memory-efficient one. + out = scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), scale=scale, + ) + return out.squeeze(0).transpose(0, 1).contiguous() seg_ids = torch.zeros(total_len, dtype=torch.int32, device=q.device) seg_ids[cu_seqlens[1:-1].long()] = 1 seg_ids = torch.cumsum(seg_ids, dim=0) diff --git a/mstar/model/bagel/config.py b/mstar/model/bagel/config.py index 297d9da0c8..426faa0d21 100644 --- a/mstar/model/bagel/config.py +++ b/mstar/model/bagel/config.py @@ -68,6 +68,9 @@ class BagelModelConfig: cfg_renorm_min: float = 0.0 think_mode: bool = False accelerator_graph: bool = True + # image_gen runs only the CFG branches a denoising step reads (see + # submodules.cfg_branches); False runs all three on every guidance-on step + cfg_prune: bool = True # Sampling defaults (per-request overridable via model_kwargs) temperature: float = 0.6 # 0 = greedy (argmax), >0 = sampling diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index e3cedd96e7..a216e37ce0 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -3,6 +3,7 @@ # --------------------------------------------------------------------------- +import functools import logging import os from collections.abc import Mapping @@ -125,6 +126,122 @@ def vit_patches( return (height // patch_size) * (width // patch_size) +# The image_gen denoise Loop. Named so its iteration count reaches the LLM +# node through ``fwd_info.dynamic_loop_iter_counts`` (an unnamed Loop gets a +# random uuid), which is how the step knows its timestep on the host. +IMAGE_GEN_LOOP = "image_gen_loop" + +# Rides on ``NodeInputs.kwargs`` from ``prepare_inputs`` to ``declare_step`` and +# the forward: how many CFG branches (1, 2 or 3) this image_gen step runs. +# Absent means the old behaviour, all three whenever guidance is on. +CFG_BRANCHES = "cfg_branches" + + +def _f32(x: float) -> float: + """``x`` rounded to float32, as a device compare against a Python scalar sees it.""" + return float(torch.tensor(float(x), dtype=torch.float32)) + + +@functools.lru_cache(maxsize=4096) +def image_gen_timestep(time_index: int, num_timesteps: int, shift: float) -> float: + """Host copy of the shifted timestep ``_forward_image_gen`` computes on the + device for iteration ``time_index``, with the same float32 ops in the same + order so the CFG-interval test agrees with the one the old code ran there. + Memoized: a request's schedule is a handful of distinct values.""" + t_uniform = 1.0 - torch.tensor([float(time_index)], dtype=torch.float32) / (num_timesteps - 1) + return float(LLMSubmodule._apply_timestep_shift(t=t_uniform, shift=shift)[0]) + + +@functools.lru_cache(maxsize=4096) +def cfg_branches( + t_value: float, cfg_interval: tuple[float, float], text_scale: float, img_scale: float, +) -> int: + """How many CFG branches an image_gen step has to run at timestep ``t_value``. + + 1 (main only): guidance off, or ``t_value`` outside ``(lo, hi]``. The old + combine then ran with both effective scales at exactly 1.0, which is + ``v_main`` up to float rounding (``a + 1 * (b - a)``, and a renorm ratio + ``|v_main| / (|v_main| + 1e-8)`` clamped to <= 1). + 2 (main, cfg_text): in the interval with ``img_scale <= 1``; the old combine + never read ``v_cfg_img`` then (``effective_img_scale > 1.0`` was False). + 3 (main, cfg_text, cfg_img): in the interval with image guidance on. + + The bounds are compared at float32, as the old device-side compare did. + ``cfg_interval`` must be hashable (a tuple); callers convert a list. + """ + if not (text_scale > 1.0 or img_scale > 1.0): + return 1 + lo, hi = cfg_interval + if not (_f32(lo) < t_value <= _f32(hi)): + return 1 + return 3 if img_scale > 1.0 else 2 + + +def image_gen_branches(inputs: list) -> int: + """The CFG branch count a guidance-on image_gen batch declares and runs: + what ``prepare_inputs`` stamped, or 3 (every branch) where it stamped + nothing — pruning off, or the old callers that predate it. The walk is + never batched; ``max`` keeps a hypothetical batch on the superset.""" + return max(inp.kwargs.get(CFG_BRANCHES, 3) for inp in inputs) + + +# The CFG branches in the order `cfg_branches` counts them: a step running n branches runs the first n. +CFG_BRANCH_LABELS = ("main", "cfg_text", "cfg_img") + + +def cfg_branch_skipped(label: str, n_branches: int | None) -> bool: + """Whether the parallel-CFG node owning ``label`` sits out a step that runs ``n_branches`` + (None: no count was stamped, so every branch runs).""" + return n_branches is not None and CFG_BRANCH_LABELS.index(label) >= n_branches + + +def combine_cfg_velocity( + v_main: torch.Tensor, + v_cfg_text: torch.Tensor, + v_cfg_img: torch.Tensor | None, + text_scale, + img_scale, + renorm_type: str, + renorm_min: float, +) -> torch.Tensor: + """BAGEL's two-node CFG velocity combine plus renormalization. + + ``v_cfg_img=None`` means the image branch did not run; that is only valid + where ``img_scale <= 1``, when the image term is skipped anyway. The scales + may be Python floats (pruned path, decided on the host) or 0-dim device + tensors (the old path, where ``img_scale > 1.0`` is a host sync). + """ + apply_img = v_cfg_img is not None and img_scale > 1.0 + if renorm_type == "text_channel": + # text_channel: renorm AFTER text CFG, THEN apply image CFG + v_text_guided = v_cfg_text + text_scale * (v_main - v_cfg_text) + norm_v = torch.norm(v_main, dim=-1, keepdim=True) + norm_v_text = torch.norm(v_text_guided, dim=-1, keepdim=True) + scale = (norm_v / (norm_v_text + 1e-8)).clamp(min=renorm_min, max=1.0) + v_text_renormed = v_text_guided * scale + if apply_img: + return v_cfg_img + img_scale * (v_text_renormed - v_cfg_img) + return v_text_renormed + + # global / channel: apply both text+image CFG, THEN renorm + v_text_guided = v_cfg_text + text_scale * (v_main - v_cfg_text) + if apply_img: + v_combined = v_cfg_img + img_scale * (v_text_guided - v_cfg_img) + else: + v_combined = v_text_guided + + if renorm_type == "channel": + renorm_scale = ( + v_main.norm(dim=-1, keepdim=True) / + (v_combined.norm(dim=-1, keepdim=True) + 1e-8) + ).clamp(min=renorm_min, max=1.0) + else: + renorm_scale = ( + v_main.norm() / (v_combined.norm() + 1e-8) + ).clamp(min=renorm_min, max=1.0) + return v_combined * renorm_scale + + class ViTEncoderSubmodule(NodeSubmodule): """SigLIP2 ViT + connector + vit_pos_embed: pixel patches -> ViT features. @@ -175,6 +292,7 @@ def __init__( # request, so the segment ceiling is the batch size. Raise only if that # changes — and raise the resource's `max_segments_per_request` with it. self._max_images_per_request = 1 + self._cu_seqlens: dict[tuple, torch.Tensor] = {} def prepare_inputs( self, @@ -218,9 +336,13 @@ def prepare_inputs( "not the count `vit_patches` gives a layout to key it by" ) - cu_seqlens = torch.tensor( - [0, num_tokens], dtype=torch.int32, device=device - ) + # built once per (length, device): `torch.tensor(..., device=)` syncs the stream, which + # held this request's prepare until the in-flight decode steps of others had finished + cu_seqlens = self._cu_seqlens.get((num_tokens, device)) + if cu_seqlens is None: + cu_seqlens = self._cu_seqlens[(num_tokens, device)] = torch.tensor( + [0, num_tokens], dtype=torch.int32, device=device + ) return NodeInputs( tensor_inputs={ @@ -736,6 +858,43 @@ def __init__( self.eoi_token_id = eoi_token_id self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id + # image_gen runs only the CFG branches a step needs (`cfg_branches`) + self._cfg_prune = config.cfg_prune + self._image_gen_shape_cache: dict[tuple, dict[str, torch.Tensor]] = {} + self._sentinel_ids: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} + self._warned_loop_counter = False + + def _image_gen_cfg_branches( + self, fwd_info: CurrentForwardPassInfo, first_iter: bool, + ) -> int | None: + """The branch count for this request's image_gen step, decided on the + host from the loop counter, or None for the old always-three behaviour. + + The counter is the step's ``time_index`` (0 on the iteration that + seeds the latents, +1 per Euler step). If it disagrees with whether + the latents were seeded here, the counter is not tracking this loop, + and guessing the timestep would change the output, so fall back. + """ + meta = fwd_info.step_metadata + if not (self._cfg_prune and meta.get("requires_cfg", False)): + return None + k = fwd_info.dynamic_loop_iter_counts.get(IMAGE_GEN_LOOP, 0) + if (k == 0) != first_iter: + if not self._warned_loop_counter: + self._warned_loop_counter = True + logger.warning( + "BAGEL image_gen: loop counter %r reads %d on a %s iteration; " + "running all CFG branches", IMAGE_GEN_LOOP, k, + "first" if first_iter else "later", + ) + return None + t = image_gen_timestep(k, self.config.num_timesteps, self.config.timestep_shift) + return cfg_branches( + t, + tuple(meta.get("cfg_interval", self.config.cfg_interval)), + meta.get("cfg_text_scale", self.config.cfg_text_scale), + meta.get("cfg_img_scale", self.config.cfg_img_scale), + ) def _get_image_pos_ids( self, labels: list[str], @@ -891,16 +1050,12 @@ def prepare_inputs( node_inputs.tensor_inputs = self._get_text_vae_idxs(seq_len, device) if graph_walk in ("image_gen", "image_gen_cfg"): - tensor_inputs = {} H = fwd_info.step_metadata.get("height", 1024) W = fwd_info.step_metadata.get("width", 1024) - - tensor_inputs["vae_position_ids"] = get_flattened_position_ids_extrapolate( - H, W, - self.config.latent_downsample, - max_num_patches_per_side=self.config.max_latent_size - ) - if "latents" not in inputs or len(inputs["latents"]) == 0: + shape_inputs = self._image_gen_shape_inputs(H, W, device) + tensor_inputs = {"vae_position_ids": shape_inputs["vae_position_ids"]} + first_iter = "latents" not in inputs or len(inputs["latents"]) == 0 + if first_iter: node_inputs.input_embeds, tensor_inputs["time_index"] = _init_latents_and_time_index( self.config, device, seed=fwd_info.random_seed, H=H, W=W ) @@ -908,6 +1063,11 @@ def prepare_inputs( node_inputs.input_embeds = inputs["latents"][0] tensor_inputs["time_index"] = inputs["time_index"][0] + if graph_walk in ("image_gen", "image_gen_cfg"): + n_branches = self._image_gen_cfg_branches(fwd_info, first_iter) + if n_branches is not None: + node_inputs.kwargs[CFG_BRANCHES] = n_branches + tensor_inputs["empty_combined_emb"] = self._wrap_with_boi_eoi( torch.empty( (node_inputs.input_embeds.shape[0], self.config.hidden_size), @@ -923,11 +1083,32 @@ def prepare_inputs( ) node_inputs.tensor_inputs = { **tensor_inputs, - **self._get_text_vae_idxs(seq_len, device) + "text_indexes": shape_inputs["text_indexes"], + "vae_token_indexes": shape_inputs["vae_token_indexes"], + "text_mask": shape_inputs["text_mask"], } return node_inputs + def _image_gen_shape_inputs(self, H: int, W: int, device) -> dict[str, torch.Tensor]: + """The image_gen inputs that depend only on the output size, on the device, built once. + + Building them per step copied host data to the device (``torch.tensor(..., device=)`` and + a CPU ``vae_position_ids`` indexing a device table), each a stream sync that held the next + step's ``prepare_inputs`` until the current step had finished. Read-only downstream. + """ + key = (H, W, str(device)) + cached = self._image_gen_shape_cache.get(key) + if cached is None: + pos = get_flattened_position_ids_extrapolate( + H, W, self.config.latent_downsample, max_num_patches_per_side=self.config.max_latent_size, + ).to(device) + cached = self._image_gen_shape_cache[key] = { + "vae_position_ids": pos, + **self._get_text_vae_idxs(pos.shape[0] + IMAGE_SENTINELS, device), + } + return cached + def declare_step( self, graph_walk: str, @@ -946,8 +1127,19 @@ def declare_step( ``NodeInputs.resource_step_info`` from ``prepare_inputs`` rather than being read off the batch here. """ + if graph_walk == "image_gen_cfg" and all( + cfg_branch_skipped(NODE_TO_CFG_LABEL.get(self.node_name, "main"), inp.kwargs.get(CFG_BRANCHES)) + for inp in inputs + ): + # this branch sits the step out (`cfg_branches`): no plan, no attention + return None requires_cfg = requires_cfg_for_inputs(inputs) labels = self._get_active_labels(graph_walk, requires_cfg) + if graph_walk == "image_gen" and requires_cfg: + # Only the branches this step runs (`cfg_branches`): labels are + # ordered main, cfg_text, cfg_img, so the first n. One branch is + # the plain "main" step below; the forward reads the same count. + labels = labels[:image_gen_branches(inputs)] spans = tuple(inp.input_seq_len for inp in inputs) # Label-major, matching how a combined plan concatenates its sources; @@ -964,9 +1156,9 @@ def declare_step( ] for label in labels } - if graph_walk == "image_gen" and requires_cfg: - # Batched CFG: one combined plan across all three labels so - # image_gen runs one forward instead of three. The frozen caches + if graph_walk == "image_gen" and len(labels) > 1: + # Batched CFG: one combined plan across the active branch labels + # so image_gen runs one forward instead of n. The frozen caches # never grow during flow matching, so nothing commits. pos_ids = [ tensor for label in labels for tensor in per_label_pos_ids[label] @@ -1093,6 +1285,13 @@ def preprocess( # and the step it was declared under cannot disagree during capture, which # passes in dummy metadata result["requires_cfg"] = requires_cfg_for_inputs(inputs) + # collate hands NodeInputs.kwargs over as per-row lists (and the dict + # itself under "kwargs"); the branch count goes in as the one int + # `declare_step` planned for, or not at all (the old 3-branch step). + result.pop("kwargs", None) + result.pop(CFG_BRANCHES, None) + if graph_walk in ("image_gen", "image_gen_cfg") and CFG_BRANCHES in inputs[0].kwargs: + result[CFG_BRANCHES] = image_gen_branches(inputs) result["sample_token"] = any([ info.step_metadata.get("sample_prefill_token", True) \ for info in engine_inputs.per_request_info.values() @@ -1289,17 +1488,19 @@ def _forward_image_gen( text_mask: torch.Tensor, time_index: torch.Tensor, requires_cfg: bool = True, + cfg_branches: int | None = None, **kwargs, ) -> NameToTensorList: - """Flow matching Euler step with optional 3-pass CFG. + """Flow matching Euler step with optional batched CFG. - Runs against the 3 frozen KV caches (main, cfg_text, cfg_img) - through one combined plan; the declared step writes no store and - commits nothing since the caches are frozen during flow - matching. + Runs against the frozen KV caches (main, cfg_text, cfg_img) through + one combined plan; the declared step writes no store and commits + nothing since the caches are frozen during flow matching. - When requires_cfg is False, runs a single forward pass (main only) - without CFG, saving 2/3 of the compute. + ``cfg_branches`` (from ``prepare_inputs`` via ``declare_step``) is how + many of those branches this step runs: 3, 2 (no image guidance) or 1 + (outside the CFG interval: main only, like requires_cfg=False). None + is the old behaviour, all three whenever requires_cfg. Renormalization modes (cfg_renorm_type in config): - "global": single scalar renorm over all dimensions (default) @@ -1340,7 +1541,10 @@ def _forward_image_gen( empty_combined_emb[1:-1] = latents_ logger.debug("packed_seq = %s", empty_combined_emb) # tensor repr: keep lazy - if requires_cfg: + if requires_cfg and cfg_branches is None: + # cfg_prune: false (or a step that never stamped a branch + # count): every guidance-on step runs all three branches, and the + # interval acts through the effective scales. cfg_text_scale = kwargs.pop("cfg_text_scale", self.config.cfg_text_scale) cfg_img_scale = kwargs.pop("cfg_img_scale", self.config.cfg_img_scale) renorm_type = kwargs.pop("cfg_renorm_type", self.config.cfg_renorm_type) @@ -1355,69 +1559,41 @@ def _forward_image_gen( effective_text_scale = cfg_text_scale * in_cfg_interval + 1.0 * (1 - in_cfg_interval) effective_img_scale = cfg_img_scale * in_cfg_interval + 1.0 * (1 - in_cfg_interval) - # Batched CFG: run all 3 branches in a single LLM forward. - # plan_attention_batched_cfg created a single FlashInfer plan - # treating (main, cfg_text, cfg_img) as 3 "virtual requests". - S = empty_combined_emb.shape[0] - batched_emb = empty_combined_emb.repeat(3, 1) # [3S, hidden] - - # Expand MoT indexes with absolute offsets for each copy - batched_text_indexes = torch.cat([ - text_indexes + i * S for i in range(3) - ]) - batched_vae_indexes = torch.cat([ - vae_token_indexes + i * S for i in range(3) - ]) - batched_text_mask = text_mask.repeat(3) - - hidden = self.language_model( - batched_emb, mode="gen", - label=self.CFG_BATCHED_LABEL, - vae_token_indexes=batched_vae_indexes, - text_indexes=batched_text_indexes, - text_mask=batched_text_mask, - **kwargs, + v_main, v_cfg_text, v_cfg_img = self._image_gen_branches_forward( + 3, empty_combined_emb, text_indexes, vae_token_indexes, text_mask, kwargs, ) - # Split output back to 3 branches and project to VAE space - v_main = self.llm2vae(hidden[:S])[1:-1] - v_cfg_text = self.llm2vae(hidden[S:2*S])[1:-1] - v_cfg_img = self.llm2vae(hidden[2*S:])[1:-1] - # Two-node CFG velocity combination + renormalization cfg_renorm_min = kwargs.pop("cfg_renorm_min", self.config.cfg_renorm_min) + # the scales are device tensors here, so the `img_scale > 1.0` + # test inside is a host sync (the pruned path decides it on the host) + v_final = combine_cfg_velocity( + v_main, v_cfg_text, v_cfg_img, + effective_text_scale, effective_img_scale, + renorm_type, cfg_renorm_min, + ) + elif requires_cfg and cfg_branches > 1: + # In the CFG interval (`cfg_branches`, decided on the host): the + # scales apply as given. Two branches when image guidance is off + # (img scale <= 1), where the combine never reads v_cfg_img. + cfg_text_scale = kwargs.pop("cfg_text_scale", self.config.cfg_text_scale) + cfg_img_scale = kwargs.pop("cfg_img_scale", self.config.cfg_img_scale) + renorm_type = kwargs.pop("cfg_renorm_type", self.config.cfg_renorm_type) + kwargs.pop("cfg_interval", None) + cfg_renorm_min = kwargs.pop("cfg_renorm_min", self.config.cfg_renorm_min) - if renorm_type == "text_channel": - # text_channel: renorm AFTER text CFG, THEN apply image CFG - v_text_guided = v_cfg_text + effective_text_scale * (v_main - v_cfg_text) - norm_v = torch.norm(v_main, dim=-1, keepdim=True) - norm_v_text = torch.norm(v_text_guided, dim=-1, keepdim=True) - scale = (norm_v / (norm_v_text + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) - v_text_renormed = v_text_guided * scale - if effective_img_scale > 1.0: - v_final = v_cfg_img + effective_img_scale * (v_text_renormed - v_cfg_img) - else: - v_final = v_text_renormed - else: - # global / channel: apply both text+image CFG, THEN renorm - v_text_guided = v_cfg_text + effective_text_scale * (v_main - v_cfg_text) - if effective_img_scale > 1.0: - v_combined = v_cfg_img + effective_img_scale * (v_text_guided - v_cfg_img) - else: - v_combined = v_text_guided - - if renorm_type == "channel": - renorm_scale = ( - v_main.norm(dim=-1, keepdim=True) / - (v_combined.norm(dim=-1, keepdim=True) + 1e-8) - ).clamp(min=cfg_renorm_min, max=1.0) - else: - renorm_scale = ( - v_main.norm() / (v_combined.norm() + 1e-8) - ).clamp(min=cfg_renorm_min, max=1.0) - v_final = v_combined * renorm_scale + v_main, v_cfg_text, v_cfg_img = self._image_gen_branches_forward( + cfg_branches, empty_combined_emb, text_indexes, vae_token_indexes, text_mask, kwargs, + ) + v_final = combine_cfg_velocity( + v_main, v_cfg_text, v_cfg_img, + cfg_text_scale, cfg_img_scale, + renorm_type, cfg_renorm_min, + ) else: - # No CFG: single forward pass over the declared plan + # No CFG, or a guidance-on step outside the CFG interval + # (cfg_branches == 1), where the combine reduces to v_main: a + # single forward pass over the declared "main" plan. hidden = self.language_model( empty_combined_emb, mode="gen", label="main", @@ -1437,6 +1613,47 @@ def _forward_image_gen( "time_index": [time_index + 1] } + def _image_gen_branches_forward( + self, + n: int, + empty_combined_emb: torch.Tensor, + text_indexes: torch.Tensor, + vae_token_indexes: torch.Tensor, + text_mask: torch.Tensor, + kwargs: dict, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """The first ``n`` CFG branches (main, cfg_text[, cfg_img]) in one LLM + forward over the combined plan `declare_step` made of the same ``n`` + labels: the plan treats each branch as a "virtual request" over its + own cache. Returns their VAE-space velocities, cfg_img None if n < 3. + """ + S = empty_combined_emb.shape[0] + batched_emb = empty_combined_emb.repeat(n, 1) # [nS, hidden] + + # Expand MoT indexes with absolute offsets for each copy + batched_text_indexes = torch.cat([ + text_indexes + i * S for i in range(n) + ]) + batched_vae_indexes = torch.cat([ + vae_token_indexes + i * S for i in range(n) + ]) + batched_text_mask = text_mask.repeat(n) + + hidden = self.language_model( + batched_emb, mode="gen", + label=self.CFG_BATCHED_LABEL, + vae_token_indexes=batched_vae_indexes, + text_indexes=batched_text_indexes, + text_mask=batched_text_mask, + **kwargs, + ) + + # Split output back per branch and project to VAE space + v_main = self.llm2vae(hidden[:S])[1:-1] + v_cfg_text = self.llm2vae(hidden[S:2*S])[1:-1] + v_cfg_img = self.llm2vae(hidden[2*S:3*S])[1:-1] if n >= 3 else None + return v_main, v_cfg_text, v_cfg_img + def _forward_image_gen_single_branch( self, input_embeds: torch.Tensor, @@ -1463,8 +1680,18 @@ def _forward_image_gen_single_branch( kwargs.pop("cfg_renorm_type", None) kwargs.pop("cfg_interval", None) kwargs.pop("cfg_renorm_min", None) + n_branches = kwargs.pop(CFG_BRANCHES, None) latents = input_embeds + output_name = { + "LLM": "v_main", + "LLM_cfg_text": "v_cfg_text", + "LLM_cfg_img": "v_cfg_img", + }.get(self.node_name, "v_main") + if cfg_branch_skipped(self._NODE_TO_CFG_LABEL.get(self.node_name, "main"), n_branches): + # sits this step out: combine_cfg never reads the velocity (see `cfg_branches`), + # but the edge still has to carry a tensor of the right shape + return {output_name: [torch.zeros_like(latents)]} # See _forward_image_gen: tie the separately-passed sequence dims so # torch.compile's dynamic shapes can prove the adds below line up. @@ -1498,12 +1725,6 @@ def _forward_image_gen_single_branch( # Output velocity. Main LLM also passes through latents/time_index # for the combine_cfg node. - output_name = { - "LLM": "v_main", - "LLM_cfg_text": "v_cfg_text", - "LLM_cfg_img": "v_cfg_img", - }.get(self.node_name, "v_main") - result: NameToTensorList = {output_name: [projected]} if self.node_name == "LLM": result["time_index"] = [time_index] @@ -1600,8 +1821,16 @@ def _wrap_with_boi_eoi(self, emb: torch.Tensor) -> torch.Tensor: assert self.boi_token_id is not None and self.eoi_token_id is not None device = emb.device - boi_ids = torch.tensor([self.boi_token_id], device=device) - eoi_ids = torch.tensor([self.eoi_token_id], device=device) + # built once per device: `torch.tensor(..., device=)` is a host copy that syncs the + # stream, which held every image_gen step's prepare_inputs until the step before it + # had finished (no speculation across denoising steps) + ids = self._sentinel_ids.get(device) + if ids is None: + ids = self._sentinel_ids[device] = ( + torch.tensor([self.boi_token_id], device=device), + torch.tensor([self.eoi_token_id], device=device), + ) + boi_ids, eoi_ids = ids with torch.no_grad(): boi_emb = self.embed_tokens(boi_ids).to(emb.dtype) eoi_emb = self.embed_tokens(eoi_ids).to(emb.dtype) @@ -1765,9 +1994,28 @@ def prepare_inputs( "time_index": inputs["time_index"][0], **result, } - return NodeInputs( - tensor_inputs=result, - kwargs=fwd_info.step_metadata + kwargs = dict(fwd_info.step_metadata) + n_branches = self._cfg_branches(fwd_info, first_iter="latents" not in inputs or not inputs["latents"]) + if n_branches is not None: + kwargs[CFG_BRANCHES] = n_branches + return NodeInputs(tensor_inputs=result, kwargs=kwargs) + + def _cfg_branches(self, fwd_info: CurrentForwardPassInfo, first_iter: bool) -> int | None: + """The step's CFG branch count (`cfg_branches`), as the LLM nodes stamp it, or None to + combine all three branches as before (pruning off, or a counter that is not tracking + this loop).""" + meta = fwd_info.step_metadata + if not (self.config.cfg_prune and meta.get("requires_cfg", False)): + return None + k = fwd_info.dynamic_loop_iter_counts.get(IMAGE_GEN_LOOP, 0) + if (k == 0) != first_iter: + return None + t = image_gen_timestep(k, self.config.num_timesteps, self.config.timestep_shift) + return cfg_branches( + t, + tuple(meta.get("cfg_interval", self.config.cfg_interval)), + meta.get("cfg_text_scale", self.config.cfg_text_scale), + meta.get("cfg_img_scale", self.config.cfg_img_scale), ) def forward( @@ -1784,6 +2032,7 @@ def forward( cfg_renorm_type: str = None, cfg_renorm_min: float = None, cfg_interval: tuple = None, + cfg_branches: int | None = None, **kwargs, ) -> NameToTensorList: if cfg_text_scale is None: @@ -1822,6 +2071,21 @@ def forward( v_ct = v_cfg_text v_ci = v_cfg_img + if cfg_branches is not None: + # decided on the host (`cfg_branches`): outside the interval the combine is v_main, + # and the image branch is read only with image guidance on; no device-side test + if cfg_branches == 1: + v_final = v_m + else: + v_final = combine_cfg_velocity( + v_m, v_ct, v_ci if cfg_branches == 3 else None, + cfg_text_scale, cfg_img_scale, cfg_renorm_type, cfg_renorm_min, + ) + new_latents = latents - v_final * dt + if torch.is_autocast_enabled(): + new_latents = new_latents.to(torch.get_autocast_gpu_dtype()) + return {"latents": [new_latents], "time_index": [time_index + 1]} + # CFG interval cfg_lo, cfg_hi = cfg_interval t_val = timestep[0] diff --git a/mstar/model/base.py b/mstar/model/base.py index 08584c7686..5e829aea65 100644 --- a/mstar/model/base.py +++ b/mstar/model/base.py @@ -27,7 +27,7 @@ TensorPointerInfo, top_level_entities, ) -from mstar.model.media_limits import VideoBudget, read_image +from mstar.model.media_limits import VIDEO_DECODE_THREADS, VideoBudget, read_image from mstar.model.multimodal import PromptPart DECODE = "decode" @@ -785,7 +785,7 @@ def load_audio(self, filepath: str, device: str) -> TensorAndMetadata: def load_video(self, filepath: str, device: str): from torchcodec.decoders import VideoDecoder - decoder = VideoDecoder(filepath, device=device) + decoder = VideoDecoder(filepath, device=device, num_ffmpeg_threads=VIDEO_DECODE_THREADS) budget = VideoBudget() frames = [] for frame in decoder: diff --git a/mstar/model/chatterbox/components/t3.py b/mstar/model/chatterbox/components/t3.py index ca25c61720..fd205a03f4 100644 --- a/mstar/model/chatterbox/components/t3.py +++ b/mstar/model/chatterbox/components/t3.py @@ -307,10 +307,12 @@ def embed_text(self, ids: torch.Tensor, text_positions: torch.Tensor) -> torch.T emb = emb + self.text_pos_emb(text_positions) return emb - def embed_speech(self, ids: torch.Tensor, speech_positions: torch.Tensor) -> torch.Tensor: + def embed_speech(self, ids: torch.Tensor, speech_position: int) -> torch.Tensor: + """Speech tokens that all sit at one speech position. The position row + is a view of the table, so no position index is staged on the device.""" emb = self.speech_emb(ids) if self.speech_pos_emb is not None: - emb = emb + self.speech_pos_emb(speech_positions) + emb = emb + self.speech_pos_emb.weight[speech_position] return emb def build_prefill_embeds( @@ -331,7 +333,7 @@ def build_prefill_embeds( if self.text_pos_emb is not None: text = text + self.text_pos_emb(torch.arange(text_ids.shape[0], device=device)) bos = torch.full((1,), self.config.start_speech_token, dtype=torch.long, device=device) - speech = self.embed_speech(bos, torch.zeros(1, dtype=torch.long, device=device)) + speech = self.embed_speech(bos, 0) parts = [cond_emb.to(text.dtype), text, speech] if self.config.duplicate_bos_in_prefill: parts.append(speech) diff --git a/mstar/model/chatterbox/submodules.py b/mstar/model/chatterbox/submodules.py index bf19e96a1f..543ac319e8 100644 --- a/mstar/model/chatterbox/submodules.py +++ b/mstar/model/chatterbox/submodules.py @@ -278,7 +278,9 @@ def prepare_inputs( tensor_inputs["uncond_embeds"] = self.model.build_prefill_embeds( cond, text_ids, uncond=True, ) - state.add_all(generated=0, requires_cfg=requires_cfg) + # decode reuses this device tensor rather than staging the + # request's constant guidance weight every step + state.add_all(generated=0, requires_cfg=requires_cfg, cfg_weight=tensor_inputs["cfg_weight"]) elif graph_walk == "decode": requires_cfg = bool(state["requires_cfg"]) # The input token's speech position is how many tokens were @@ -290,8 +292,8 @@ def prepare_inputs( f"speech position {step} is past the T3 table ({self.t3.speech_pos_table_size})" ) token = inputs[PREV_TOKEN][0].to(device, torch.long).reshape(-1) - embeds = self.model.embed_speech(token, torch.tensor([step], device=device)) - tensor_inputs = self._scalar_inputs(fwd_info, requires_cfg) + embeds = self.model.embed_speech(token, step) + tensor_inputs = {"cfg_weight": state["cfg_weight"]} else: raise ValueError(f"Unknown T3 walk {graph_walk!r}") return ARNodeInputs( diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index 377ffb582c..d85944114f 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -58,10 +58,21 @@ from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule +from mstar.utils.device import DeviceCapability logger = logging.getLogger(__name__) +def _to_device_no_sync(tensor: torch.Tensor, device: torch.device) -> torch.Tensor: + """A small host tensor onto ``device`` without a stream sync. A pageable ``.to(device)`` + waits for everything already queued on the stream (other requests' denoise steps), which + stalls the GPU thread under concurrency; a pinned non-blocking copy does not.""" + device = torch.device(device) + if tensor.device.type != "cpu" or not DeviceCapability.of(device).is_accelerator: + return tensor.to(device) + return tensor.pin_memory().to(device, non_blocking=True) + + class StepScalar(NamedTuple): """One per-step per-request scalar, sliced out of the request's schedule. @@ -371,7 +382,7 @@ def _step_scalar_sources( f"entries for {schedule.num_steps} steps, but " f"{scalar.source} has {tensor.shape[0]}." ) - staged.setdefault(scalar.source, tensor.to(device)) + staged.setdefault(scalar.source, _to_device_no_sync(tensor, device)) return staged def _carried_inputs( diff --git a/mstar/model/components/diffusion/qwen3/encoder.py b/mstar/model/components/diffusion/qwen3/encoder.py index 06fbcd03c2..47a1a9c1f4 100644 --- a/mstar/model/components/diffusion/qwen3/encoder.py +++ b/mstar/model/components/diffusion/qwen3/encoder.py @@ -209,6 +209,7 @@ def __init__( raise ValueError("tap layers are 1-indexed decoder layers (0 would be the raw embeddings)") self.tap_layers = tuple(tap_layers) self.head_dim, self.rope_theta = head_dim, rope_theta + self._rotary_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} self.embed_tokens = nn.Embedding(vocab_size, hidden_size) self.layers = nn.ModuleList( DecoderLayer( @@ -226,18 +227,32 @@ def __init__( def dtype(self) -> torch.dtype: return self.embed_tokens.weight.dtype - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + def forward( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor, padded: bool | None = None, + ) -> torch.Tensor: """``input_ids, attention_mask [B, S]`` (right padded) -> ``[B, S, len(taps) * hidden]``. Follows HF's two attention paths exactly: a batch without padding attends mask-free (``is_causal`` + native GQA), a padded one through a boolean causal-and-key-padding mask; the two pick different SDPA kernels, so the choice is part of parity. + + ``padded`` is whether any row has padding, when the caller knows it on the host; + None reads it off ``attention_mask``, a host sync (and not graph-capturable). """ seq = input_ids.shape[1] - rotary = qwen3_rotary_tables(torch.arange(seq, device=input_ids.device), self.head_dim, self.rope_theta) + # cached per (length, device): building them copies inv_freq from the host, which a + # captured forward cannot do (and the tables never change) + key = (seq, input_ids.device) + rotary = self._rotary_cache.get(key) + if rotary is None: + rotary = self._rotary_cache[key] = qwen3_rotary_tables( + torch.arange(seq, device=input_ids.device), self.head_dim, self.rope_theta, + ) # No padding anywhere in the batch: HF's create_causal_mask returns None and attention # runs its is_causal path; mirror that rather than passing an all-True mask. - mask = None if bool(attention_mask.all()) else padded_causal_mask(attention_mask) + if padded is None: + padded = not bool(attention_mask.all()) + mask = padded_causal_mask(attention_mask) if padded else None hidden = self.embed_tokens(input_ids) taps = [] for layer_idx, layer in enumerate(self.layers, start=1): diff --git a/mstar/model/components/moe.py b/mstar/model/components/moe.py index f48b755330..17d864c63b 100644 --- a/mstar/model/components/moe.py +++ b/mstar/model/components/moe.py @@ -96,6 +96,11 @@ class TopKRouter(nn.Module): to (top-k). norm_topk_prob: if True, renormalize the top-k probabilities so they sum to 1. + fused_softmax_topk: on CUDA, route through one Triton kernel + (``mstar::moe_softmax_topk``) instead of softmax, ``torch.topk`` + and the renormalize. Equal probabilities go to the lower expert + id, where ``torch.topk`` may pick another, so a model opts in + once it has checked its outputs. """ def __init__( @@ -104,11 +109,13 @@ def __init__( num_experts: int, num_experts_per_tok: int, norm_topk_prob: bool = True, + fused_softmax_topk: bool = False, ) -> None: super().__init__() self.top_k = num_experts_per_tok self.num_experts = num_experts self.norm_topk_prob = norm_topk_prob + self.fused_softmax_topk = fused_softmax_topk self.weight = nn.Parameter(torch.zeros(num_experts, hidden_size)) def forward( @@ -130,6 +137,13 @@ def forward( """ hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) router_logits = F.linear(hidden_states, self.weight) + if self.fused_softmax_topk and router_logits.is_cuda: + from mstar.utils.fused_moe.router import softmax_topk + + routing_weights, router_indices = softmax_topk( + router_logits, self.top_k, self.norm_topk_prob, + ) + return routing_weights, router_indices, None router_logits = F.softmax(router_logits, dtype=torch.float, dim=-1) router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) @@ -221,6 +235,7 @@ def __init__( moe_intermediate_size: int, norm_topk_prob: bool = True, router: nn.Module | None = None, + fused_softmax_topk: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -233,6 +248,7 @@ def __init__( num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, norm_topk_prob=norm_topk_prob, + fused_softmax_topk=fused_softmax_topk, ) self.experts = nn.Module() self.experts.gate_up_proj = nn.Parameter( @@ -282,6 +298,7 @@ def __init__( shared_expert: nn.Module, norm_topk_prob: bool = False, router: nn.Module | None = None, + fused_softmax_topk: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -294,6 +311,7 @@ def __init__( num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, norm_topk_prob=norm_topk_prob, + fused_softmax_topk=fused_softmax_topk, ) self.experts = nn.Module() self.experts.gate_up_proj = nn.Parameter( @@ -394,6 +412,7 @@ def __init__( norm_topk_prob: bool = True, router: nn.Module | None = None, comm_group: CommGroup | None = None, + fused_softmax_topk: bool = False, ) -> None: super().__init__() if comm_group is None: @@ -412,6 +431,7 @@ def __init__( num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, norm_topk_prob=norm_topk_prob, + fused_softmax_topk=fused_softmax_topk, ) shard_inter = divide(moe_intermediate_size, tp_size) @@ -499,6 +519,7 @@ def __init__( norm_topk_prob: bool = False, router: nn.Module | None = None, comm_group: CommGroup | None = None, + fused_softmax_topk: bool = False, ) -> None: super().__init__() if comm_group is None: @@ -517,6 +538,7 @@ def __init__( num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, norm_topk_prob=norm_topk_prob, + fused_softmax_topk=fused_softmax_topk, ) shard_inter = divide(moe_intermediate_size, tp_size) diff --git a/mstar/model/cosmos3/config.py b/mstar/model/cosmos3/config.py index 5bb0262468..d6c1c4c782 100644 --- a/mstar/model/cosmos3/config.py +++ b/mstar/model/cosmos3/config.py @@ -375,6 +375,12 @@ class Cosmos3Config: # sharing the whole-clip decoder's, which keeps the default layout: a fixed # ~1.03 GiB of bf16 weights (Wan2.2 VAE) in a process hosting both nodes. vae_decode_channels_last: bool = True + # Decode a whole clip (the non-windowed video and image walks) with the shared + # Wan VAE decoder: the same per-latent decoder calls as ``vae.decode`` from an + # empty cache, channels-last and compiled per call in Inductor's deterministic + # mode, on a decoder-only copy of the VAE (~1 GiB bf16). False decodes with the + # compiled stock ``vae.decode`` on the shared instance. + vae_decode_fast: bool = True # Let cuDNN time its conv algorithms for each new streaming-decode shape # (once, then cached) instead of its heuristic pick. Same math, but the # timed pick can differ between server launches, and so can the frames @@ -526,7 +532,8 @@ def check_streaming_decode(self) -> None: reinterpret silently: a non-boolean switch (a yaml ``"false"`` is truthy), a negative warm-stream cap (read as no cap), or more latents per call than the default window holds (read as one call per window).""" - for name in ("windowed_decode_warm_cache", "vae_decode_channels_last", "vae_decode_cudnn_benchmark"): + for name in ("windowed_decode_warm_cache", "vae_decode_channels_last", "vae_decode_cudnn_benchmark", + "vae_decode_fast"): value = getattr(self, name) if not isinstance(value, bool): raise ValueError(f"Cosmos3 {name} must be true or false, got {value!r}.") diff --git a/mstar/model/cosmos3/cosmos3_model.py b/mstar/model/cosmos3/cosmos3_model.py index f552a67b47..e952399278 100644 --- a/mstar/model/cosmos3/cosmos3_model.py +++ b/mstar/model/cosmos3/cosmos3_model.py @@ -1757,9 +1757,14 @@ def _create_submodule(self, node_name: str, device: str, tp_group=None, sp_group vae=self._build_encode_vae(device), config=self.config ) if node_name == VAE_DECODER_NODE: - return Cosmos3VAEDecoderSubmodule( - vae=self._build_vae(device), config=self.config + # The fast decode converts its VAE to channels-last in place, so it + # takes a decoder-only copy rather than the shared instance. + vae = ( + self._build_decoder_only_vae(device) + if self.config.vae_decode_fast + else self._build_vae(device) ) + return Cosmos3VAEDecoderSubmodule(vae=vae, config=self.config) if node_name == VAE_DECODER_AR_NODE: # Shares the decoder VAE weights unless it decodes channels-last, # which converts the VAE in place; only the streaming chunk state diff --git a/mstar/model/cosmos3/submodules.py b/mstar/model/cosmos3/submodules.py index e899794de2..c29005222f 100644 --- a/mstar/model/cosmos3/submodules.py +++ b/mstar/model/cosmos3/submodules.py @@ -298,6 +298,11 @@ def __init__(self, transformer, config, scheduler=None): self.transformer.denoise_step_batched_cfg = torch.compile( self.transformer.denoise_step_batched_cfg, fullgraph=False, dynamic=False, ) + # One graph per batch size: dynamo unrolls the per-request loop, so + # the first batch of each size compiles (seconds with a warm cache). + self.transformer.denoise_step_action_batched = torch.compile( + self.transformer.denoise_step_action_batched, fullgraph=False, dynamic=False, + ) logger.info("Cosmos3 denoise compute torch.compile enabled") def to(self, *args, **kwargs): @@ -2678,6 +2683,9 @@ class Cosmos3VAEDecoderSubmodule(NodeSubmodule): # One-shot decode per request; CUDA-graph capture (not torch.compile) is the # speedup path. disable_torch_compile = True + # Whole-clip decodes through the shared Wan VAE decoder (config.vae_decode_fast); + # the streaming subclass sets up its own decoder step. + _whole_clip_fast = True def __init__(self, vae, config): super().__init__() @@ -2685,6 +2693,7 @@ def __init__(self, vae, config): self.config = config self._decode_dtype_cached = None self._channels_last = False + self._fast_step = None # The decode is 3D-conv bound, one-shot per request at request-specific # shapes, and not graph-captured; torch.compile fuses its pointwise # epilogues (~20-30% off the 189-frame video decode, neutral for @@ -2694,15 +2703,26 @@ def __init__(self, vae, config): self._decode = vae.decode if vae is not None else None self._decode_compiled = False self._warmed_decode_shapes: set[tuple[int, ...]] = set() + fast = vae is not None and self._whole_clip_fast and config.vae_decode_fast if vae is not None and os.environ.get("COSMOS3_COMPILE_VAE", "1").lower() not in ( "0", "false", "no", "off", ): - self._decode = _wan_vae().compile_decode(vae.decode) self._decode_compiled = True + # The fast path compiles its own DecoderStep; the stock decode is unused there. + if not fast: + self._decode = _wan_vae().compile_decode(vae.decode) logger.info("Cosmos3 VAE decode torch.compile enabled") # Resolve + log the decode dtype now (a cheap cuDNN-version read) so the # choice is fixed at startup, not on the first request. self._decode_dtype() + if fast: + # Cast before the layout switch and the compile trace, so both see the + # weights the decode runs with (this instance is the node's own). + self.vae = vae = vae.to(self._decode_dtype()) + _wan_vae().to_channels_last(vae) + self._channels_last = True + self._fast_step = _wan_vae().DecoderStep(vae, compile=self._decode_compiled, deterministic=True) + logger.info("Cosmos3 VAE decode: whole clips on the shared Wan VAE decoder (channels-last)") def prepare_inputs(self, graph_walk, fwd_info, inputs, **kwargs) -> NodeInputs: return NodeInputs(tensor_inputs={"latents": inputs["latents"][0]}) @@ -2757,6 +2777,13 @@ def _decode_pixels(self, latents: torch.Tensor) -> torch.Tensor: """Latents -> uint8 pixel frames ``[1, 3, T, H, W]``.""" vae_dtype = self._decode_dtype() z = self._vae_input(latents) + if self._fast_step is not None: + wan_vae = _wan_vae() + with torch.autocast(device_type="cuda", enabled=False): + decoded = wan_vae.decode_causal( + self.vae, z, wan_vae.new_feat_cache(self.vae), self._fast_step, clip_start=True, + ) + return self._to_uint8(decoded) with torch.autocast(device_type="cuda", enabled=False): if ( self._decode_compiled @@ -2816,6 +2843,8 @@ class Cosmos3VAEDecoderARSubmodule(Cosmos3VAEDecoderSubmodule): done, so it is not vetoed. """ + _whole_clip_fast = False + def __init__(self, vae, config): super().__init__(vae, config) # Per-session decode context (the latents behind a session's last diff --git a/mstar/model/cosmos3/tests/test_serving.py b/mstar/model/cosmos3/tests/test_serving.py index 1970fbe825..e7708074aa 100644 --- a/mstar/model/cosmos3/tests/test_serving.py +++ b/mstar/model/cosmos3/tests/test_serving.py @@ -409,7 +409,8 @@ def decode(self, z): monkeypatch.setenv("COSMOS3_COMPILE_VAE", "0") # keep _decode the raw stub vae = _StubVAE() - sub = Cosmos3VAEDecoderSubmodule(vae, Cosmos3Config()) + # the stock compiled-decode path (vae_decode_fast routes around it) + sub = Cosmos3VAEDecoderSubmodule(vae, Cosmos3Config(vae_decode_fast=False)) sub._decode_compiled = True sub._decode_dtype_cached = torch.bfloat16 @@ -749,3 +750,44 @@ def test_video_postprocess_uses_request_fps(tmp_path) -> None: if NANO_DIR.exists(): test_process_prompt_emits_cond_and_uncond() print("PASS") + + +def test_vae_decode_fast_routes_whole_clips_through_the_shared_decoder(monkeypatch) -> None: + """With vae_decode_fast the whole-clip node decodes through + wan_vae.decode_causal from an empty cache, channels-last; the streaming + subclass keeps its own decoder step.""" + from types import SimpleNamespace + + import torch + + from mstar.model.components import wan_vae + from mstar.model.cosmos3.config import Cosmos3Config + from mstar.model.cosmos3.submodules import Cosmos3VAEDecoderSubmodule + + class _StubVAE(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + self.decoder = torch.nn.Identity() + self.config = SimpleNamespace(latents_mean=[0.0] * 4, latents_std=[1.0] * 4) + + def decode(self, z): # the stock path, never taken here + raise AssertionError("stock vae.decode called") + + calls = [] + monkeypatch.setenv("COSMOS3_COMPILE_VAE", "0") + monkeypatch.setattr(wan_vae, "to_channels_last", lambda vae: calls.append("channels_last")) + monkeypatch.setattr(wan_vae, "new_feat_cache", lambda vae: []) + + def fake_decode_causal(vae, z, feat_cache, step, *, clip_start, latents_per_call=1): + calls.append(("decode_causal", clip_start, feat_cache == [])) + t = 1 + (z.shape[2] - 1) * 4 + return torch.zeros(1, 3, t, 8, 8, dtype=torch.bfloat16) + + monkeypatch.setattr(wan_vae, "decode_causal", fake_decode_causal) + sub = Cosmos3VAEDecoderSubmodule(_StubVAE(), Cosmos3Config()) + sub._decode_dtype_cached = torch.bfloat16 + out = sub.forward("video_gen", None, torch.zeros(1, 4, 3, 8, 8, dtype=torch.bfloat16)) + assert calls == ["channels_last", ("decode_causal", True, True)] + assert out["video_output"][0].shape == (1, 3, 9, 8, 8) + assert out["video_output"][0].dtype == torch.uint8 diff --git a/mstar/model/cosmos3/tests/test_windowed_decode.py b/mstar/model/cosmos3/tests/test_windowed_decode.py index 607e99d34d..bc8eb0811e 100644 --- a/mstar/model/cosmos3/tests/test_windowed_decode.py +++ b/mstar/model/cosmos3/tests/test_windowed_decode.py @@ -347,6 +347,7 @@ def test_whole_clip_decoder_keeps_its_default_layout_and_cudnn_flag(monkeypatch) monkeypatch.setenv("COSMOS3_COMPILE_VAE", "0") cfg = Cosmos3Config() cfg.vae_decode_cudnn_benchmark = True + cfg.vae_decode_fast = False # the stock whole-clip path; the fast one is tested below assert cfg.vae_decode_channels_last vae = _tiny_vae() sub = Cosmos3VAEDecoderSubmodule(vae, cfg) @@ -361,6 +362,27 @@ def test_whole_clip_decoder_keeps_its_default_layout_and_cudnn_flag(monkeypatch) assert torch.equal(video, (ref / 2 + 0.5).clamp(0, 1).mul(255).to(torch.uint8)) +def test_fast_whole_clip_decode_matches_stock_decode(monkeypatch) -> None: + """vae_decode_fast decodes a whole clip through the shared decoder, + channels-last, on the node's own VAE: the frames of the stock decode.""" + monkeypatch.setenv("COSMOS3_COMPILE_VAE", "0") + ref_vae = _tiny_vae() + stream = _stream(5) + stock = Cosmos3VAEDecoderSubmodule(ref_vae, Cosmos3Config(vae_decode_fast=False)) + stock._decode_dtype_cached = torch.float32 + want = stock._decode_pixels(stream) + fast_vae = _tiny_vae() + fast_vae.load_state_dict(ref_vae.state_dict()) + cfg = Cosmos3Config() + assert cfg.vae_decode_fast + monkeypatch.setenv("COSMOS3_VAE_DECODE_FP32", "1") + fast = Cosmos3VAEDecoderSubmodule(fast_vae, cfg) + assert is_channels_last(fast_vae) + got = fast._decode_pixels(stream) + assert got.shape == want.shape + assert (got.int() - want.int()).abs().max() <= 1 + + @pytest.mark.parametrize("enabled", [False, True]) @pytest.mark.parametrize("warm", [True, False]) def test_streaming_cudnn_benchmark_is_opt_in_and_scoped(monkeypatch, vae, enabled, warm) -> None: diff --git a/mstar/model/fish_s2_pro/components/dac_layers.py b/mstar/model/fish_s2_pro/components/dac_layers.py index cc083ed3e7..1907fe550a 100644 --- a/mstar/model/fish_s2_pro/components/dac_layers.py +++ b/mstar/model/fish_s2_pro/components/dac_layers.py @@ -45,9 +45,38 @@ import torch.nn.functional as F from torch import nn +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def _snake_kernel(x_ptr, out_ptr, alpha_ptr, inv_alpha_ptr, channels, length, BLOCK: tl.constexpr): + row = tl.program_id(0) # one (batch, channel) row of a contiguous [B, C, T] + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = cols < length + channel = row % channels + alpha = tl.load(alpha_ptr + channel) + inv_alpha = tl.load(inv_alpha_ptr + channel) + offsets = row.to(tl.int64) * length + cols + x = tl.load(x_ptr + offsets, mask=mask) + s = tl.sin(alpha * x) + tl.store(out_ptr + offsets, x + inv_alpha * (s * s), mask=mask) + class Snake1d(nn.Module): - """``x + sin(alpha * x)^2 / alpha``, per channel.""" + """``x + sin(alpha * x)^2 / alpha``, per channel. + + On fp32 CUDA inputs one Triton pass replaces the five elementwise kernels + the expression launches over the decoder's audio-rate activations; every + intermediate stays fp32, as in the eager form. + """ + + BLOCK = 1024 def __init__(self, channels: int): super().__init__() @@ -56,8 +85,15 @@ def __init__(self, channels: int): def forward(self, x: torch.Tensor) -> torch.Tensor: shape = x.shape x = x.reshape(shape[0], shape[1], -1) - x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) - return x.reshape(shape) + inv_alpha = (self.alpha + 1e-9).reciprocal() + if triton is None or not x.is_cuda or x.dtype != torch.float32 or self.alpha.dtype != torch.float32: + return (x + inv_alpha * torch.sin(self.alpha * x).pow(2)).reshape(shape) + x = x.contiguous() + out = torch.empty_like(x) + batch, channels, length = x.shape + grid = (batch * channels, triton.cdiv(length, self.BLOCK)) + _snake_kernel[grid](x, out, self.alpha, inv_alpha, channels, length, BLOCK=self.BLOCK) + return out.reshape(shape) class VectorQuantize(nn.Module): diff --git a/mstar/model/fish_s2_pro/submodules.py b/mstar/model/fish_s2_pro/submodules.py index bd16497842..c3302779a6 100644 --- a/mstar/model/fish_s2_pro/submodules.py +++ b/mstar/model/fish_s2_pro/submodules.py @@ -33,7 +33,11 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo -from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, BatchedAcceleratorGraphConfig +from mstar.engine.accelerator_graph_config import ( + AcceleratorGraphConfig, + BatchedAcceleratorGraphConfig, + PackedAcceleratorGraphConfig, +) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep from mstar.engine.resources.sampler.resource import SamplerResource @@ -94,8 +98,15 @@ class ARSubmodule(ARNodeSubmodule): """ disable_torch_compile = True - MAX_BATCH_SIZE = 32 - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # The top decode bucket is also the cap on rows per AR step; a step + # costs about the same at 64 rows as at 32, so 64 streams do not run as + # two alternating batches. + DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 24, 32, 48, 64] + MAX_BATCH_SIZE = DECODE_CAPTURE_BATCH_SIZES[-1] + # Plain prompts are short (the chat frame around one text). Eager, a + # prefill step took ~60 ms of mostly launch-bound stream time on an H200. + PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] + PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8] def __init__(self, slow_ar: FishSlowAR, fast_ar: FishFastAR, config: FishS2ProConfig) -> None: super().__init__() @@ -354,10 +365,29 @@ def max_batch_size(self, graph_walk: str) -> int: def get_accelerator_graph_configs( self, device: torch.device, tp_world_size: int = 1 ) -> list[AcceleratorGraphConfig]: - """Capture the whole one-token decode step (both ARs and sampling) per batch bucket.""" + """Capture the whole one-token decode step (both ARs and sampling) per + batch bucket, and prefills packed into token buckets. A voice clone's + prefill is the same forward over a prompt that embeds the reference + frames, so it replays the same captures when it fits the buckets.""" del tp_world_size dtype = self.slow_ar.semantic_head.dtype + + def prefill_input(num_tokens: int) -> ARNodeInputs: + return ARNodeInputs( + input_embeds=torch.zeros(num_tokens, self.config.slow_ar.dim, dtype=dtype, device=device), + input_seq_len=num_tokens, + kwargs={"suppress_eos": True}, + ) + return [ + PackedAcceleratorGraphConfig( + capture_graph_walk="ar_prefill", + replay_graph_walks=["ar_prefill", "ar_prefill_clone"], + capture_token_lengths=self.PREFILL_TOKEN_BUCKETS, + make_node_input=prefill_input, + capture_batch_sizes=self.PREFILL_CAPTURE_BATCH_SIZES, + compile=False, + ), BatchedAcceleratorGraphConfig( capture_graph_walk="ar_decode", single_request_inputs=ARNodeInputs( @@ -366,12 +396,14 @@ def get_accelerator_graph_configs( kwargs={"suppress_eos": False}, ), capture_batch_sizes=self.DECODE_CAPTURE_BATCH_SIZES, - compile=False, + compile=True, ) ] def can_use_accelerator_graphs(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: - if batch.graph_walk != "ar_decode" or not self.can_batch(batch, model_inputs): + if batch.graph_walk not in ("ar_decode", "ar_prefill", "ar_prefill_clone") or not self.can_batch( + batch, model_inputs, + ): return False return super().can_use_accelerator_graphs(batch, model_inputs) diff --git a/mstar/model/flux2_klein/flux2_klein_model.py b/mstar/model/flux2_klein/flux2_klein_model.py index cb50f94bcc..d8c5166c60 100644 --- a/mstar/model/flux2_klein/flux2_klein_model.py +++ b/mstar/model/flux2_klein/flux2_klein_model.py @@ -43,7 +43,7 @@ from mstar.engine.resources import NodeResourceSpec, RaggedAttentionConfig, RaggedAttentionSpec from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata +from mstar.model.base import ForwardPassArgs, Model, ProcessPromptOutput, TensorAndMetadata from mstar.model.components.diffusion.decode_utils import VAE_DECODE_BATCH_SIZES from mstar.model.components.diffusion.image_io import encode_image from mstar.model.components.diffusion.lora import LoraSpec @@ -61,6 +61,7 @@ TEXT_EMBEDS, TEXT_INPUTS, TEXT_MASK, + TEXT_PADDED, KleinDenoiseSubmodule, KleinShape, KleinTextEncoderSubmodule, @@ -84,6 +85,9 @@ class Flux2KleinModel(Model): """FLUX.2 [klein] (4B by default; the 9B checkpoint loads through the same class).""" + # settled by process_prompt (tokenization); a client cannot set it + INTERNAL_KWARGS = (TEXT_PADDED,) + def __init__( self, model_path_hf: str = FLUX2_KLEIN_4B, @@ -334,7 +338,11 @@ def process_prompt( if float(kwargs.get("guidance_scale", 1.0)) != 1.0: logger.warning("guidance_scale is ignored: FLUX.2 klein is step-distilled and runs without CFG") ids, mask = self.tokenize(prompt) - return {TEXT_INPUTS: [ids], TEXT_MASK: [mask]} + # right padded: the last position is padding unless the prompt fills the context + return ProcessPromptOutput( + new_input_tensors={TEXT_INPUTS: [ids], TEXT_MASK: [mask]}, + metadata={TEXT_PADDED: not bool(mask[-1])}, + ) # --------------------------------------------------------- state machine def _step_metadata(self, metadata: CurrentForwardConductorMetadata) -> dict: @@ -344,6 +352,7 @@ def _step_metadata(self, metadata: CurrentForwardConductorMetadata) -> dict: "height": kw["height"], "width": kw["width"], "num_inference_steps": kw["num_inference_steps"], "ref_grids": kw["ref_grids"], + TEXT_PADDED: kw.get(TEXT_PADDED), } def get_initial_forward_pass_args( @@ -365,6 +374,7 @@ def get_initial_forward_pass_args( "width": width, "num_inference_steps": self._resolve_steps(model_kwargs), "ref_grids": [list(self.config.latent_grid(h, w)) for h, w in ref_dims], + TEXT_PADDED: model_kwargs.get(TEXT_PADDED), } logger.info( "FLUX.2 klein request: %dx%d steps=%d refs=%s seed=%s", width, height, kwargs["num_inference_steps"], @@ -470,7 +480,11 @@ def _create_submodule(self, node_name: str, device) -> NodeSubmodule | None: if node_name == "text_encoder": encoder = build_text_encoder(self.config.text_encoder, self.snapshot, device) - return KleinTextEncoderSubmodule(encoder, self.config, max_batch_size=self.max_batch_size) + return KleinTextEncoderSubmodule( + encoder, self.config, max_batch_size=self.max_batch_size, + capture_walk=ENCODE_TEXT_WALK if self.accelerator_graph else None, + capture_batch_sizes=self.capture_batch_sizes, + ) if node_name == "vae_encoder": return KleinVaeEncoderSubmodule(self._vae_module(device), self.config) if node_name == "vae_decoder": diff --git a/mstar/model/flux2_klein/submodules.py b/mstar/model/flux2_klein/submodules.py index c5304ec783..a672b8bb13 100644 --- a/mstar/model/flux2_klein/submodules.py +++ b/mstar/model/flux2_klein/submodules.py @@ -49,6 +49,8 @@ # Edge names shared with the model file. TEXT_INPUTS, TEXT_MASK, TEXT_EMBEDS = "text_inputs", "text_mask", "text_embeds" +# step metadata: whether the tokenized prompt is padded, settled on the host at tokenization +TEXT_PADDED = "text_padded" IMAGE_INPUTS, REF_LATENTS, IMAGE_OUTPUT = "image_inputs", "ref_latents", "image_output" @@ -93,22 +95,35 @@ class KleinTextEncoderSubmodule(BatchedRows, NodeSubmodule): ``text_inputs`` / ``text_mask`` rows stack into ``[B, 512]`` and each request gets its ``[1, 512, joint_dim]`` slice of the output. All 512 positions (pad tokens included) are emitted, as the reference feeds them all to the DiT. + + With ``capture_batch_sizes`` the encoder is captured per batch size. Eager it is + launch-bound (~1.6k small kernels, ~10 ms of GPU work in a ~50 ms forward at B=1). + Only the padded-prompt path is captured: whether a prompt has padding is settled + on the host at tokenization (``TEXT_PADDED`` in the step metadata), and a batch + holding a prompt that fills all 512 tokens runs the encoder's unmasked path eagerly. """ disable_torch_compile = True output_keys = (TEXT_EMBEDS,) keep_batch_dim = True - def __init__(self, encoder: nn.Module, config: Flux2KleinConfig, max_batch_size: int = 16): + def __init__( + self, encoder: nn.Module, config: Flux2KleinConfig, max_batch_size: int = 16, + capture_walk: str | None = None, capture_batch_sizes: Sequence[int] = (), + ): super().__init__() self.encoder = encoder self.config = config self._max_batch_size = max_batch_size + self._capture_walk = capture_walk + self._capture_batch_sizes = [b for b in capture_batch_sizes if b <= max_batch_size] def prepare_inputs(self, graph_walk, fwd_info, inputs: NameToTensorList, **kwargs) -> NodeInputs: return NodeInputs( tensor_inputs={TEXT_INPUTS: inputs[TEXT_INPUTS][0], TEXT_MASK: inputs[TEXT_MASK][0]}, input_seq_len=self.config.text_encoder.max_sequence_length, + # None (not settled at tokenization) keeps the encoder's own device-side check + resource_step_info=fwd_info.step_metadata.get(TEXT_PADDED), ) def can_batch(self, batch, model_inputs) -> bool: @@ -117,17 +132,53 @@ def can_batch(self, batch, model_inputs) -> bool: def max_batch_size(self, graph_walk: str): return self._max_batch_size + def cg_key_info(self, graph_walk: str, per_request_info, per_request_input_metadata=None, **kwargs): + """True when every real row is a padded prompt: the only captured path. Padding rows + of a captured bucket carry no step metadata.""" + del graph_walk, per_request_input_metadata, kwargs + flags = [info.step_metadata.get(TEXT_PADDED) for info in per_request_info.values() if info.step_metadata] + return bool(flags) and all(flag is True for flag in flags) + def preprocess(self, graph_walk, engine_inputs, inputs: list[NodeInputs]) -> dict: + flags = [inp.resource_step_info for inp in inputs] return { TEXT_INPUTS: torch.stack([inp.tensor_inputs[TEXT_INPUTS] for inp in inputs]), TEXT_MASK: torch.stack([inp.tensor_inputs[TEXT_MASK] for inp in inputs]), + # any padded row puts the batch on the masked path, as the reference batches it + "padded": None if None in flags else any(flags), } - def run_batch(self, text_inputs: torch.Tensor, text_mask: torch.Tensor, **kwargs): + def run_batch(self, text_inputs: torch.Tensor, text_mask: torch.Tensor, padded: bool | None = None, **kwargs): device = self.get_device() - embeds = self.encoder(text_inputs.to(device=device, dtype=torch.long), text_mask.to(device=device)) + embeds = self.encoder( + text_inputs.to(device=device, dtype=torch.long), text_mask.to(device=device), padded=padded, + ) return {TEXT_EMBEDS: embeds.to(self.encoder.dtype)} + def get_accelerator_graph_configs(self, device: torch.device, tp_world_size: int = 1): + if self._capture_walk is None or not self._capture_batch_sizes: + return [] + from mstar.engine.accelerator_graph_config import BatchedAcceleratorGraphConfig + + seq = self.config.text_encoder.max_sequence_length + # a padding row's prompt: one real token, so every query row attends to something + mask = torch.zeros(seq, dtype=torch.long, device=device) + mask[0] = 1 + single = NodeInputs( + tensor_inputs={TEXT_INPUTS: torch.zeros(seq, dtype=torch.long, device=device), TEXT_MASK: mask}, + input_seq_len=seq, resource_step_info=True, + ) + return [BatchedAcceleratorGraphConfig( + capture_graph_walk=self._capture_walk, + single_request_inputs=single, + additional_key_info=True, + compile=False, + capture_batch_sizes=list(self._capture_batch_sizes), + # batches past the largest capture (or holding an unpadded prompt) run eager + caps_eager_batch_size=False, + input_seq_dims={TEXT_INPUTS: 0, TEXT_MASK: 0}, + )] + # --------------------------------------------------------------------------- # vae_encoder (reference images for editing) diff --git a/mstar/model/higgs_audio/components/audio_tower.py b/mstar/model/higgs_audio/components/audio_tower.py index 9c6bdc6a19..9322286b48 100644 --- a/mstar/model/higgs_audio/components/audio_tower.py +++ b/mstar/model/higgs_audio/components/audio_tower.py @@ -59,17 +59,43 @@ def __init__(self, config: HiggsAudioModelConfig): self.avg_pooler = nn.AvgPool1d(2, stride=2) self.layer_norm = nn.LayerNorm(embed_dim) - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - """(batch, num_mel_bins, T_mel) -> (batch, ~T_mel/4, d_model).""" + def forward(self, input_features: torch.Tensor, valid_frames: torch.Tensor | None = None) -> torch.Tensor: + """(batch, num_mel_bins, T_mel) -> (batch, ~T_mel/4, d_model). + + ``valid_frames`` ``[batch]``: rows are mel prefixes right-padded with + zeros, and each row comes out as its prefix alone would: the padded + tail is zeroed where a convolution would read it and kept out of + attention as keys. Rows past their valid outputs are left unspecified. + """ hidden = F.gelu(self.conv1(input_features)) + if valid_frames is not None: + # conv2 reads one frame past an odd-length prefix: it must be the zero + # the unpadded conv pads with + hidden = hidden * _prefix_mask(valid_frames, hidden.shape[-1])[:, None, :] hidden = F.gelu(self.conv2(hidden)) hidden = hidden.permute(0, 2, 1) # (B, T, D) hidden = hidden + self.embed_positions.weight[: hidden.shape[1]] + attn_mask = None + if valid_frames is not None: + keys = _prefix_mask((valid_frames - 1) // 2 + 1, hidden.shape[1]) + attn_mask = torch.zeros(keys.shape, dtype=hidden.dtype, device=hidden.device).masked_fill( + ~keys, torch.finfo(hidden.dtype).min + )[:, None, None, :] for layer in self.layers: - out = layer(hidden, None, None) if _LAYER_TAKES_HEAD_MASK else layer(hidden, None) + out = layer(hidden, attn_mask, None) if _LAYER_TAKES_HEAD_MASK else layer(hidden, attn_mask) hidden = out[0] if isinstance(out, tuple) else out # tf4 returns (hidden, attn) hidden = self.avg_pooler(hidden.permute(0, 2, 1)).permute(0, 2, 1) - return self.layer_norm(hidden) + hidden = self.layer_norm(hidden) + if valid_frames is not None: + # the projector's stride-2 conv reads one step past an odd-length row + pooled = ((valid_frames - 1) // 2 + 1) // 2 + hidden = hidden * _prefix_mask(pooled, hidden.shape[1])[:, :, None] + return hidden + + +def _prefix_mask(lengths: torch.Tensor, size: int) -> torch.Tensor: + """``[B, size]`` bool, True where ``t < lengths[b]``.""" + return torch.arange(size, device=lengths.device)[None, :] < lengths[:, None] class HiggsAudioFeatureProjector(nn.Module): diff --git a/mstar/model/higgs_audio/higgs_audio_model.py b/mstar/model/higgs_audio/higgs_audio_model.py index 167ed92c1c..dd73aae368 100644 --- a/mstar/model/higgs_audio/higgs_audio_model.py +++ b/mstar/model/higgs_audio/higgs_audio_model.py @@ -14,10 +14,10 @@ audio_encoder (enc_dec) — audio_tower + projector (checkpoint code) LLM (ar) — dense Qwen3, paged KV cache -Graph walks (sequential prefill schedule, qwen3_omni-style): - prefill_text — text span -> LLM KV cache (the second/post-audio - span is the last prefill: samples the first token) - prefill_audio — audio_encoder -> LLM KV cache +Graph walks: + prefill_audio — audio_encoder -> LLM: the ChatML text around the audio + and the audio embeddings, prefilled as one span that + samples the first token decode — LLM loop Not implemented from the reference pipeline: silero-VAD chunk @@ -86,6 +86,8 @@ def __init__( ): self.cache_dir = cache_dir self.model_path_hf = model_path_hf + # ``model_kwargs: {encoder_graphs: false}`` runs the audio tower eagerly + self.encoder_graphs = bool(kwargs.get("encoder_graphs", True)) self.local_dir = _resolve_local_hf_snapshot(model_path_hf, cache_dir=cache_dir) self.config = HiggsAudioModelConfig.from_pretrained(self.local_dir) @@ -181,12 +183,6 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: ), ] - prefill_text = GraphNode( - name="LLM", - input_names=["text_inputs"], - outputs=llm_prefill_outputs, - ) - prefill_audio = Sequential([ GraphNode( name="audio_encoder", @@ -195,7 +191,7 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: ), GraphNode( name="LLM", - input_names=["audio_embeds"], + input_names=["text_pre", "audio_embeds", "text_post"], outputs=[edge.clone() for edge in llm_prefill_outputs], ), ]) @@ -205,15 +201,17 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: section=GraphNode( name="LLM", input_names=["text_inputs"], + # one signal to the client and back into the loop, so a step + # can keep it on the device and send the client's copy inline outputs=[ GraphEdge( - next_node=EMIT_TO_CLIENT, - name="new_token", - output_modality="text", + next_node="LLM", + name="text_inputs", ), GraphEdge( - next_node="LLM", + next_node=EMIT_TO_CLIENT, name="text_inputs", + output_modality="text", ), ], ), @@ -222,7 +220,6 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: ) return dict( - prefill_text=prefill_text, prefill_audio=prefill_audio, decode=decode, ) @@ -235,7 +232,8 @@ def _build_prefill_schedule( self, input_signals: dict[str, list[TensorPointerInfo]], ) -> list[tuple[str, dict[str, TensorPointerInfo]]]: - """[text-before-audio] + [audio] + [text-after-audio]. + """One ``prefill_audio`` step: the ChatML text before the audio, the + audio embeddings and the text after it, prefilled as one span. ``process_prompt`` emits two ``text_inputs`` tensors (the ChatML prompt split at the audio splice point) plus the per-chunk mel @@ -244,29 +242,28 @@ def _build_prefill_schedule( texts = input_signals.get("text_inputs", []) audio_features = input_signals.get("audio_features", []) audio_feature_lens = input_signals.get("audio_feature_lens", []) - - schedule: list[tuple[str, dict[str, TensorPointerInfo]]] = [] - if texts: - schedule.append(("prefill_text", {"text_inputs": texts[0]})) - if audio_features: - entry = {"audio_features": audio_features[0]} - if audio_feature_lens: - entry["audio_feature_lens"] = audio_feature_lens[0] - schedule.append(("prefill_audio", entry)) - for text in texts[1:]: - schedule.append(("prefill_text", {"text_inputs": text})) - return schedule + if len(texts) != 2 or not audio_features or not audio_feature_lens: + raise ValueError( + "Higgs-Audio expects the prompt split around one audio input; got " + f"{len(texts)} text spans and {len(audio_features)} audio inputs." + ) + return [("prefill_audio", { + "text_pre": texts[0], + "audio_features": audio_features[0], + "audio_feature_lens": audio_feature_lens[0], + "text_post": texts[1], + })] def _prefill_inputs( self, metadata: CurrentForwardConductorMetadata, ) -> list[GraphEdge]: schedule = metadata.kwargs["prefill_schedule"] step = metadata.kwargs["prefill_step"] - walk_name, tensor_dict = schedule[step] - target_node = "audio_encoder" if walk_name == "prefill_audio" else "LLM" + _, tensor_dict = schedule[step] edges: list[GraphEdge] = [] for input_name, tensor_info in tensor_dict.items(): + target_node = "audio_encoder" if input_name.startswith("audio_") else "LLM" edge = GraphEdge(next_node=target_node, name=input_name) edge.tensor_info = [tensor_info] edges.append(edge) @@ -501,6 +498,7 @@ def _create_encoder_submodule(self, device: str) -> NodeSubmodule: from mstar.model.higgs_audio.submodules import HiggsAudioEncoderSubmodule return HiggsAudioEncoderSubmodule( audio_tower=audio_tower, projector=projector, config=self.config, + graphs=self.encoder_graphs, ) @staticmethod diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index b2a287d117..1505f3f3b9 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -15,6 +15,7 @@ from typing import Any import torch +import torch.nn.functional as F from torch import nn from mstar.communication.tensors import NameToTensorList @@ -23,6 +24,9 @@ AcceleratorGraphConfig, BatchedAcceleratorGraphConfig, PackedAcceleratorGraphConfig, + PiecewiseAcceleratorGraphConfig, + PiecewiseBatchedConfig, + PiecewiseCallInputs, ) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep @@ -32,9 +36,12 @@ from mstar.model.submodule_base import ( ARNodeInputs, ARNodeSubmodule, + BatchedModelOutput, + HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + device_loopback_enabled, ) logger = logging.getLogger(__name__) @@ -55,19 +62,40 @@ class HiggsAudioEncoderSubmodule(NodeSubmodule): Mirrors the reference pipeline: batch-encode padded chunks, run the projector on the padded batch, then slice each chunk to its valid - (downsampled) length and concatenate in order. + (downsampled) length and concatenate in order. Chunks attend only + within themselves, so requests whose chunks share a width are encoded + as one batch. + + The chunk count and, for a clip under one chunk, the mel width change + per request; compiled, every new shape recompiled on the serving path + (tens of seconds with the worker blocked). So the tower runs eagerly, + and full-width chunks replay a CUDA graph by chunk count. A clip under + one chunk is padded to the full width and replays a second graph whose + tower masks the padding out (``HiggsAudioTower.forward``'s + ``valid_frames``), so short clips batch together too. """ + disable_torch_compile = True + # requests per encoder step, and full-width chunks per captured replay + MAX_BATCH_SIZE = 16 + CAPTURE_CHUNK_COUNTS = [1, 2, 4, 8, 16, 32, 64] + TOWER = "tower" + MASKED_TOWER = "tower_masked" + def __init__( self, audio_tower: nn.Module, projector: nn.Module, config: HiggsAudioModelConfig, + graphs: bool = True, ): super().__init__() self.audio_tower = audio_tower self.projector = projector self.config = config + self.graphs = graphs + # mel frames of a full chunk: the feature extractor's 10 ms hop + self.chunk_frames = int(config.chunk_size_seconds * 100) def prepare_inputs( self, @@ -76,37 +104,138 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs, ) -> NodeInputs: + features = inputs["audio_features"][0] + if features.dim() == 2: + features = features.unsqueeze(0) # (1, num_mel_bins, T) + # one host read per request: the valid embeddings of each chunk + valid = [ + self.config.encoder_output_length(int(n)) + for n in inputs["audio_feature_lens"][0].reshape(-1).tolist() + ] return NodeInputs( - tensor_inputs={ - "audio_features": inputs["audio_features"][0], - "audio_feature_lens": inputs["audio_feature_lens"][0], - } + tensor_inputs={"audio_features": features}, + kwargs={"valid_lengths": valid}, + input_seq_len=features.shape[0], ) - def forward( + def _padded(self, width: int) -> bool: + """A clip under one chunk: run at the full width, its tail masked.""" + return self.graphs and width < self.chunk_frames + + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: + # chunks are padded to their own request's longest, so only requests of + # one mel width share a batch -- or, padded and masked, short clips + widths = {inp.tensor_inputs["audio_features"].shape[-1] for inp in model_inputs} + chunks = sum(inp.input_seq_len for inp in model_inputs) + same_kind = len(widths) == 1 or all(self._padded(w) for w in widths) + return same_kind and chunks <= self.CAPTURE_CHUNK_COUNTS[-1] + + def max_batch_size(self, graph_walk: str) -> int: + return self.MAX_BATCH_SIZE + + def preprocess( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - audio_features: torch.Tensor, - audio_feature_lens: torch.Tensor, - **kwargs, - ) -> NameToTensorList: - device = self.get_device() + inputs: list[NodeInputs], + ) -> dict[str, Any]: dtype = next(self.audio_tower.parameters()).dtype - feats = audio_features.to(device=device, dtype=dtype) - if feats.dim() == 2: - feats = feats.unsqueeze(0) # (1, num_mel_bins, T) - - encoded = self.audio_tower(feats) # (num_chunks, T_out, 1280) - projected = self.projector(encoded) # (num_chunks, T', hidden) + features = [inp.tensor_inputs["audio_features"] for inp in inputs] + widths = [f.shape[-1] for f in features] + padded = self._padded(widths[0]) + if padded: + features = [F.pad(f, (0, self.chunk_frames - f.shape[-1])) for f in features] + out = { + "audio_features": torch.cat(features, dim=0).to(device=self.get_device(), dtype=dtype), + "valid_lengths": [inp.kwargs["valid_lengths"] for inp in inputs], + } + if padded: + frames = torch.tensor(widths, dtype=torch.long) + if out["audio_features"].is_cuda: + frames = frames.pin_memory() + out["valid_frames"] = frames.to(self.get_device(), non_blocking=True) + return out + + def _encode(self, features: torch.Tensor, valid_frames: torch.Tensor | None = None) -> torch.Tensor: + return self.projector(self.audio_tower(features, valid_frames)) + + def _capture_tower(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + return {"projected": self._encode(inp.static_inputs["audio_features"])} + + def _capture_masked_tower(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + s = inp.static_inputs + # padding rows replay zeroed: give them one valid frame + return {"projected": self._encode(s["audio_features"], s["valid_frames"].clamp(min=1))} + + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, **kwargs: Any, + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + """Tower + projector over full-width chunks, by chunk count.""" + del autocast_dtype, tp_world_size, kwargs + if not self.graphs: + return {} + dtype = next(self.audio_tower.parameters()).dtype + mels, frames = self.config.audio_num_mel_bins, self.chunk_frames + return { + self.TOWER: PiecewiseBatchedConfig( + capture_fn=self._capture_tower, + make_static_inputs=lambda shape: { + "audio_features": torch.zeros(shape.bs, mels, frames, dtype=dtype, device=device), + }, + seq_len=frames, + capture_batch_sizes=list(self.CAPTURE_CHUNK_COUNTS), + ), + # clips under one chunk, one per row, padded to the full width + self.MASKED_TOWER: PiecewiseBatchedConfig( + capture_fn=self._capture_masked_tower, + make_static_inputs=lambda shape: { + "audio_features": torch.zeros(shape.bs, mels, frames, dtype=dtype, device=device), + "valid_frames": torch.full((shape.bs,), frames, dtype=torch.long, device=device), + }, + seq_len=frames, + capture_batch_sizes=[bs for bs in self.CAPTURE_CHUNK_COUNTS if bs <= self.MAX_BATCH_SIZE], + ), + } - chunks = [] - for i, mel_len in enumerate(audio_feature_lens.tolist()): - valid = self.config.encoder_output_length(int(mel_len)) - chunks.append(projected[i, :valid]) - audio_embeds = torch.cat(chunks, dim=0) # (total_audio_tokens, hidden) + def forward_batched( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + audio_features: torch.Tensor, + valid_lengths: list[list[int]], + valid_frames: torch.Tensor | None = None, + **kwargs, + ) -> dict[str, NameToTensorList]: + num_chunks = audio_features.shape[0] + static_inputs = {"audio_features": audio_features} + if valid_frames is not None: + static_inputs["valid_frames"] = valid_frames + runner = engine_inputs.piecewise_runners.get(self.TOWER if valid_frames is None else self.MASKED_TOWER) + if ( + runner is not None + and audio_features.shape[-1] == self.chunk_frames + and runner.can_run(num_chunks) + ): + projected = runner.run(static_inputs=static_inputs, real_bs=num_chunks).get_view("projected") + else: + projected = self._encode(audio_features, valid_frames) + + outputs: dict[str, NameToTensorList] = {} + chunk = 0 + for rid, valid in zip(engine_inputs.request_ids, valid_lengths, strict=True): + embeds = torch.cat([projected[chunk + i, :n] for i, n in enumerate(valid)], dim=0) + chunk += len(valid) + outputs[rid] = {"audio_embeds": [embeds]} + return outputs - return {"audio_embeds": [audio_embeds]} + def forward( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + **kwargs, + ) -> NameToTensorList: + rid = engine_inputs.request_ids[0] + return self.forward_batched(graph_walk, engine_inputs, **kwargs)[rid] # =================================================================== @@ -118,18 +247,20 @@ class HiggsAudioLLMSubmodule(ARNodeSubmodule): """Dense Qwen3 LLM. Dispatches on graph_walk: - - prefill_text: embed a text span, extend the KV cache. - - prefill_audio: splice the encoder's audio embeddings, extend the - KV cache. + - prefill_audio: the ChatML text around the encoder's audio + embeddings, one span into the KV cache. - decode: embed the previous token, single-step decode. - Every walk samples a token; ``postprocess`` drops it for the prefill - steps that are not the last one, so the captured graph's output keys - stay the same whatever the step metadata says. The two prefill walks - share one packed capture — both hand the forward ``input_embeds`` — - while decode carries bare ``input_ids`` and embeds inside the graph. + Every walk samples a token; ``postprocess`` drops it for a prefill step + that is not the last one, so the captured graph's output keys stay the + same whatever the step metadata says. Prefill replays a packed capture + keyed on ``input_embeds``, while decode carries bare ``input_ids`` and + embeds inside the graph. """ + # the decode walk reads its input ids off the sampler's slot masters + reads_device_loopback = True + PREFILL_TOKEN_BUCKETS = [64, 128, 256, 512, 1024, 2048] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8] DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] @@ -152,8 +283,7 @@ def get_accelerator_graph_configs( capture_batch_sizes=self.DECODE_CAPTURE_BATCH_SIZES, ), PackedAcceleratorGraphConfig( - capture_graph_walk="prefill_text", - replay_graph_walks=["prefill_text", "prefill_audio"], + capture_graph_walk="prefill_audio", make_node_input=lambda n: ARNodeInputs( input_seq_len=n, input_embeds=torch.zeros( @@ -171,30 +301,49 @@ def prepare_inputs( graph_walk: str, fwd_info: CurrentForwardPassInfo, inputs: NameToTensorList, + has_loopback: bool = False, **kwargs, ) -> ARNodeInputs: device = self.get_device() + if graph_walk == "decode" and has_loopback and self._device_loopback(): + # the token is read off the sampler's slot master in `preprocess` + return ARNodeInputs(input_seq_len=1) + if graph_walk == "prefill_audio": - audio_embeds = inputs["audio_embeds"][0].to(device) + # the audio between the ChatML text around it, prefilled as one span + audio = inputs["audio_embeds"][0].to(device) + pre, post = ( + self.model.embed_tokens(inputs[name][0].to(device).reshape(-1)) + if inputs.get(name) else None + for name in ("text_pre", "text_post") + ) + dtype = (pre if pre is not None else audio).dtype + embeds = torch.cat([t.to(dtype) for t in (pre, audio, post) if t is not None], dim=0) return ARNodeInputs( - input_seq_len=audio_embeds.shape[0], - input_embeds=audio_embeds, + input_seq_len=embeds.shape[0], + input_embeds=embeds, ) + # decode: embedding happens inside the captured graph; a decode row + # is one id, so staging it costs 8 bytes instead of a hidden vector token_ids = inputs["text_inputs"][0].to(device).reshape(-1) - if graph_walk == "decode": - # embedding happens inside the captured graph; a decode row is - # one id, so staging it costs 8 bytes instead of a hidden vector - return ARNodeInputs(input_seq_len=1, input_ids=token_ids) - - # prefill_text shares prefill_audio's capture, which is keyed on - # embeddings, so the text span is embedded here - return ARNodeInputs( - input_seq_len=token_ids.shape[0], - input_ids=token_ids, - input_embeds=self.model.embed_tokens(token_ids), - ) + return ARNodeInputs(input_seq_len=1, input_ids=token_ids) + + def uniform_row_inputs( + self, graph_walk: str, has_loopback: bool = False, + ) -> ARNodeInputs | None: + if graph_walk == "decode" and has_loopback and self._device_loopback(): + return ARNodeInputs(input_seq_len=1) + return None + + def _device_loopback(self) -> bool: + return device_loopback_enabled() and self.node_resources[SAMPLER].has_slot_masters + + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: + if graph_walk == "decode" and self._device_loopback(): + return frozenset({"text_inputs"}) + return frozenset() def declare_step( self, @@ -237,9 +386,19 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: if graph_walk == "decode": - return { - "input_ids": torch.cat([inp.input_ids for inp in inputs]), - } + if all(inp.input_ids is not None for inp in inputs): + return {"input_ids": torch.cat([inp.input_ids for inp in inputs])} + # device loop-back: each row's last sampled token, read by slot; + # rows whose previous step was the prefill brought theirs + sampler: SamplerResource = engine_inputs.resources[SAMPLER] + ids = sampler.loopback_tokens(engine_inputs.request_ids) + routed = [i for i, inp in enumerate(inputs) if inp.input_ids is not None] + if routed: + ids.index_copy_( + 0, torch.tensor(routed, device=ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + return {"input_ids": ids} return { "input_embeds": torch.cat( [inp.input_embeds for inp in inputs], dim=0 @@ -257,6 +416,8 @@ def _forward( sampler: SamplerResource = engine_inputs.resources[SAMPLER] attn: AttentionManager = engine_inputs.resources[ATTN] + if input_ids is None and input_embeds is None: + input_ids = sampler.loopback_tokens(engine_inputs.request_ids) if input_embeds is None: input_embeds = self.model.embed_tokens(input_ids) hidden = self.model(input_embeds=input_embeds, label="main") @@ -298,17 +459,18 @@ def forward_batched( input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, - ) -> dict[str, NameToTensorList]: + ) -> BatchedModelOutput: new_tokens = self._forward( graph_walk=graph_walk, engine_inputs=engine_inputs, input_ids=input_ids, input_embeds=input_embeds, ) - return { - rid: {"new_token": [new_tokens[i:i + 1]]} - for i, rid in enumerate(engine_inputs.request_ids) - } + # row i is request i: one clone and one device-to-host copy per step + return BatchedModelOutput( + row_outputs={"new_token": new_tokens}, + check_stop_buffers={"new_token": new_tokens}, + ) def postprocess( self, @@ -342,3 +504,43 @@ def check_stop( decoded_tokens >= request_info.max_tokens: return {"decode_loop"} return set() + + def step_is_input_free(self, graph_walk: str) -> bool: + # a decode row is one token, like the capture template's + return graph_walk == "decode" + + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: + # the decode loop's client edge carries the token the stop check + # already has on the host; prefill's client edge is persisted + return {"text_inputs": "new_token"} if graph_walk == "decode" else {} + + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for the whole batch off one ``tolist``.""" + if step_walk is not None and step_walk != "decode": + return {} + tokens = host_rows.buffers.get("new_token") + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + stop_ids = self.config.stop_token_ids + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + if step_walk is None and info.graph_walk != "decode": + continue + if ( + values[i] in stop_ids + and not info.resource_configs[SAMPLER].ignore_eos + ) or info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens: + stops[rid] = {"decode_loop"} + return stops diff --git a/mstar/model/kokoro/components/lstm.py b/mstar/model/kokoro/components/lstm.py index ee788ff6a8..4b6f74c338 100644 --- a/mstar/model/kokoro/components/lstm.py +++ b/mstar/model/kokoro/components/lstm.py @@ -25,11 +25,47 @@ class MaskedBiLSTM(nn.Module): def __init__(self, input_size: int, hidden_size: int): super().__init__() + self.hidden_size = hidden_size self.fwd = nn.LSTM(input_size, hidden_size, batch_first=True) self.bwd = nn.LSTM(input_size, hidden_size, batch_first=True) + # Set by ``fuse``; kept out of the state dict, which stays the checkpoint's. + self._fused: list[nn.LSTM] = [] + + @torch.no_grad() + def fuse(self) -> None: + """Run both directions as one LSTM of twice the width. + + Its weights are block-diagonal: each gate's first ``hidden`` rows read + the forward input and state, the second the reversed input and the + backward state, so the two directions stay independent while their + time steps advance together. The recurrence is a chain of tiny, + latency-bound kernels per step, so this halves the steps' cost. + """ + h, i = self.hidden_size, self.fwd.input_size + ref = self.fwd.weight_ih_l0 + fused = nn.LSTM(2 * i, 2 * h, batch_first=True).to(device=ref.device, dtype=ref.dtype) + for name in ("weight_ih_l0", "weight_hh_l0", "bias_ih_l0", "bias_hh_l0"): + dst = getattr(fused, name) + dst.zero_() + for direction, (rows, cols) in enumerate(((0, 0), (h, i if "ih" in name else h))): + src = getattr(self.bwd if direction else self.fwd, name) + for gate in range(4): + block = src[gate * h : (gate + 1) * h] + r = gate * 2 * h + rows + if block.dim() == 1: + dst[r : r + h] = block + else: + dst[r : r + h, cols : cols + block.shape[1]] = block + fused.flatten_parameters() + self._fused = [fused] def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: - forward_out, _ = self.fwd(x) - backward_out, _ = self.bwd(reverse_padded(x, lengths)) + reversed_x = reverse_padded(x, lengths) + if self._fused: + out, _ = self._fused[0](torch.cat([x, reversed_x], dim=-1)) + forward_out, backward_out = out.split(self.hidden_size, dim=-1) + else: + forward_out, _ = self.fwd(x) + backward_out, _ = self.bwd(reversed_x) out = torch.cat([forward_out, reverse_padded(backward_out, lengths)], dim=-1) return mask_features(out, length_mask(lengths, x.shape[1])) diff --git a/mstar/model/kokoro/components/synthesizer.py b/mstar/model/kokoro/components/synthesizer.py index d9610493b6..eff51b5071 100644 --- a/mstar/model/kokoro/components/synthesizer.py +++ b/mstar/model/kokoro/components/synthesizer.py @@ -18,6 +18,7 @@ from mstar.model.kokoro.components.albert import PLBert from mstar.model.kokoro.components.decoder import Decoder +from mstar.model.kokoro.components.lstm import MaskedBiLSTM from mstar.model.kokoro.components.masking import length_mask, mask_channels, mask_features from mstar.model.kokoro.components.prosody import ProsodyPredictor, TextEncoder from mstar.model.kokoro.config import KokoroModelConfig @@ -33,6 +34,11 @@ def __init__(self, config: KokoroModelConfig): self.text_encoder = TextEncoder(config) self.decoder = Decoder(config) + def fuse_lstm_directions(self) -> None: + for module in self.modules(): + if isinstance(module, MaskedBiLSTM): + module.fuse() + def _split_style(self, style: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """A pack row is [decoder style, predictor style].""" return style[:, : self.config.style_dim], style[:, self.config.style_dim :] diff --git a/mstar/model/kokoro/config.py b/mstar/model/kokoro/config.py index 52509dd065..b0bbd928f9 100644 --- a/mstar/model/kokoro/config.py +++ b/mstar/model/kokoro/config.py @@ -152,6 +152,13 @@ class KokoroModelConfig: # Snake) and the fused kernels are what the CUDA graphs capture. One # compile serves every bucket. Off on CPU. compile_decoder: bool = True + # Run each bidirectional LSTM's two directions as one block-diagonal LSTM + # (see ``MaskedBiLSTM.fuse``): half the sequential steps. + fuse_lstm_directions: bool = True + # Let cuDNN time its conv algorithms for the frame half's captured buckets + # (once each, at startup) instead of taking its heuristic pick; the eager + # fallback keeps the heuristic. + decoder_cudnn_benchmark: bool = True # --------------------------------------------------------------------- @property diff --git a/mstar/model/kokoro/kokoro_model.py b/mstar/model/kokoro/kokoro_model.py index 95a6db6742..8340b1419c 100644 --- a/mstar/model/kokoro/kokoro_model.py +++ b/mstar/model/kokoro/kokoro_model.py @@ -56,7 +56,7 @@ SERVING_OVERRIDES = { "default_voice", "default_speed", "first_chunk_target_phonemes", "chunk_target_phonemes", "max_chunks", "text_buckets", "frame_buckets", "capture_batch_sizes", "max_batch_frames", "frame_grouping", "decoder_dtype", - "compile_decoder", + "compile_decoder", "fuse_lstm_directions", "decoder_cudnn_benchmark", } @@ -284,5 +284,7 @@ def _create_synth_submodule(self, device: str) -> NodeSubmodule: model = KokoroTTS(self.config) load_kokoro_weights(model, Path(local_dir) / self.config.weights_file) model = model.to(device).eval() + if self.config.fuse_lstm_directions: + model.fuse_lstm_directions() logger.info("Loaded Kokoro (%d voices) on %s", len(self.voices.names), device) return KokoroSynthSubmodule(model, self.config) diff --git a/mstar/model/kokoro/submodules.py b/mstar/model/kokoro/submodules.py index ac8c36088b..45f6e84296 100644 --- a/mstar/model/kokoro/submodules.py +++ b/mstar/model/kokoro/submodules.py @@ -22,6 +22,7 @@ from __future__ import annotations +import contextlib import logging import os import time @@ -65,6 +66,18 @@ PCM16_SCALE = 32767 +@contextlib.contextmanager +def cudnn_benchmark(enabled: bool): + """cuDNN times its conv algorithms per new shape inside the block (the frame + half runs on the worker's one GPU thread); the process flag is restored.""" + previous = torch.backends.cudnn.benchmark + torch.backends.cudnn.benchmark = enabled or previous + try: + yield + finally: + torch.backends.cudnn.benchmark = previous + + def text_region(bucket: int) -> str: return f"text_T{bucket}" @@ -256,6 +269,8 @@ def _synthesize( real_bs=len(rows), )["audio"] else: + # cuDNN's heuristic pick: autotuning here would search again + # for every new frame count the eager fallback sees num_frames = max(sizes[r] for r in rows) audio, _ = self.model.synthesize_frames( d[index], t_en[index], pred_dur[index], style[index], num_frames @@ -356,7 +371,9 @@ def _text_capture(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: def _frames_capture(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: si = inp.static_inputs - return {"audio": self.model.decode_frames(si["en"], si["asr"], si["frame_lengths"].clamp(min=1), si["style"])} + with cudnn_benchmark(self.config.decoder_cudnn_benchmark): + audio = self.model.decode_frames(si["en"], si["asr"], si["frame_lengths"].clamp(min=1), si["style"]) + return {"audio": audio} def _text_static_inputs(self, shape: PiecewiseCaptureShape, bucket: int, device) -> dict[str, torch.Tensor]: return { diff --git a/mstar/model/lingbot_world/components/dit.py b/mstar/model/lingbot_world/components/dit.py index 606927d6c3..fff2cb0159 100644 --- a/mstar/model/lingbot_world/components/dit.py +++ b/mstar/model/lingbot_world/components/dit.py @@ -59,9 +59,17 @@ def require_flash_attn() -> None: ) from e -def _flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: +def _flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, fa4: bool = False) -> torch.Tensor: """The reference ``flash_attention`` for one unpadded sequence: q/k/v - ``[1, L, H, D]`` bf16 through FA2 varlen; returns ``[1, Lq, H, D]``.""" + ``[1, L, H, D]`` bf16 through FA2 varlen; returns ``[1, Lq, H, D]``. + ``fa4`` runs flash-attn-4's sm90 kernel instead (not the reference's + numerics; see ``LingBotWorldConfig.reference_numerics``).""" + if fa4: + from flash_attn.cute import flash_attn_func + + out = flash_attn_func(q.to(v.dtype), k.to(v.dtype), v) + out = out[0] if isinstance(out, tuple) else out + return out.type(q.dtype) import flash_attn b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype @@ -116,7 +124,8 @@ def causal_rope_apply( freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), ], dim=-1).reshape(seq_len, 1, -1) x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[0, seq_len:]]) + if seq_len < x.size(1): + x_i = torch.cat([x_i, x[0, seq_len:]]) return x_i.unsqueeze(0).type_as(x) @@ -160,6 +169,9 @@ class RollingKVCache: local_end: int = 0 cross_k: list[torch.Tensor] = field(default_factory=list) cross_v: list[torch.Tensor] = field(default_factory=list) + # Per-block camera (scale, shift) for the chunk whose rays are ``cam_rays``. + cam_rays: torch.Tensor | None = None + cam_mod: list[tuple[torch.Tensor, torch.Tensor]] = field(default_factory=list) @classmethod def allocate( @@ -224,22 +236,25 @@ def __init__(self, dim: int, num_heads: int, eps: float): self.o = nn.Linear(dim, dim) self.norm_q = WanRMSNorm(dim, eps=eps) self.norm_k = WanRMSNorm(dim, eps=eps) + self.fa4 = False def forward( self, x: torch.Tensor, grid: tuple[int, int, int], freqs: torch.Tensor, start_frame: int, cache: RollingKVCache, layer: int, plan: _WritePlan, ) -> torch.Tensor: b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - q = self.norm_q(self.q(x)).view(b, s, n, d) - k = self.norm_k(self.k(x)).view(b, s, n, d) - v = self.v(x).view(b, s, n, d) + q, k, v = self._qkv(x) + q, k, v = q.view(b, s, n, d), k.view(b, s, n, d), v.view(b, s, n, d) roped_query = causal_rope_apply(q, grid, freqs, start_frame).type_as(v) roped_key = causal_rope_apply(k, grid, freqs, start_frame).type_as(v) cache.write(layer, roped_key, v, plan) window = slice(max(0, plan.local_end - cache.capacity), plan.local_end) - x = _flash_attention(roped_query, cache.k[layer][:, window], cache.v[layer][:, window]) + x = _flash_attention(roped_query, cache.k[layer][:, window], cache.v[layer][:, window], self.fa4) return self.o(x.flatten(2)) + def _qkv(self, x): + return self.norm_q(self.q(x)), self.norm_k(self.k(x)), self.v(x) + class CrossAttention(nn.Module): """Cross-attention to the 512-token text stream. The text K/V are computed @@ -255,6 +270,7 @@ def __init__(self, dim: int, num_heads: int, eps: float): self.o = nn.Linear(dim, dim) self.norm_q = WanRMSNorm(dim, eps=eps) self.norm_k = WanRMSNorm(dim, eps=eps) + self.fa4 = False def text_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: b, n, d = context.size(0), self.num_heads, self.head_dim @@ -262,8 +278,11 @@ def text_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: def forward(self, x: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: b, n, d = x.size(0), self.num_heads, self.head_dim - q = self.norm_q(self.q(x)).view(b, -1, n, d) - return self.o(_flash_attention(q, k, v).flatten(2)) + q = self._q(x).view(b, -1, n, d) + return self.o(_flash_attention(q, k, v, self.fa4).flatten(2)) + + def _q(self, x): + return self.norm_q(self.q(x)) class CausalBlock(nn.Module): @@ -281,25 +300,39 @@ def __init__(self, dim: int, ffn_dim: int, num_heads: int, eps: float): self.cam_scale_layer = nn.Linear(dim, dim) self.cam_shift_layer = nn.Linear(dim, dim) + def cam_modulation(self, cam: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The block's camera scale and shift; a function of the chunk's Plucker + rays alone, so the same for all of a chunk's forwards.""" + cam_hidden = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(cam))) + cam_hidden = cam_hidden + cam + return self.cam_scale_layer(cam_hidden), self.cam_shift_layer(cam_hidden) + def forward( - self, x: torch.Tensor, e: torch.Tensor, cam: torch.Tensor, cross_k: torch.Tensor, + self, x: torch.Tensor, e: torch.Tensor, cam_mod: tuple[torch.Tensor, torch.Tensor], cross_k: torch.Tensor, cross_v: torch.Tensor, attn_args: dict, ) -> torch.Tensor: with torch.amp.autocast("cuda", dtype=torch.float32): e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2) - y = self.self_attn(self.norm1(x).float() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), **attn_args) + y = self.self_attn(self._pre(x, e[0], e[1]), **attn_args) + x = self._post_self(x, y, e[2], *cam_mod) + x = x + self.cross_attn(self._norm3(x), cross_k, cross_v) + return self._ffn(x, e[3], e[4], e[5]) + + def _pre(self, x, e0, e1): + return self.norm1(x).float() * (1 + e1.squeeze(2)) + e0.squeeze(2) + + def _post_self(self, x, y, e2, cam_scale, cam_shift): with torch.amp.autocast("cuda", dtype=torch.float32): - x = x + y * e[2].squeeze(2) + x = x + y * e2.squeeze(2) + return (1.0 + cam_scale) * x + cam_shift - cam_hidden = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(cam))) - cam_hidden = cam_hidden + cam - x = (1.0 + self.cam_scale_layer(cam_hidden)) * x + self.cam_shift_layer(cam_hidden) + def _norm3(self, x): + return self.norm3(x) - x = x + self.cross_attn(self.norm3(x), cross_k, cross_v) - y = self.ffn(self.norm2(x).float() * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + def _ffn(self, x, e3, e4, e5): + y = self.ffn(self.norm2(x).float() * (1 + e4.squeeze(2)) + e3.squeeze(2)) with torch.amp.autocast("cuda", dtype=torch.float32): - x = x + y * e[5].squeeze(2) - return x + return x + y * e5.squeeze(2) class CausalHead(nn.Module): @@ -341,6 +374,26 @@ def __init__(self, config: LingBotWorldConfig): rope_params(n, 2 * (d // 6)), ], dim=1) self._freqs: dict[torch.device, torch.Tensor] = {} + if not config.reference_numerics: + self._fast_numerics() + + def _fast_numerics(self) -> None: + """Trade bit-parity with the reference for speed: flash-attn-4 + attention, and each block's pointwise chains (modulation, norms, FFN + activation, gated residuals) compiled around the GEMMs. Attention, RoPE + and the cache writes stay eager: they take the chunk's Python-int + cache indices, which would recompile per chunk.""" + from mstar.model.components.diffusion.compile_utils import with_inductor_rounding + + def compiled(fn): + return with_inductor_rounding(torch.compile(fn, dynamic=False), eager_rounding=True) + + for block in self.blocks: + block.self_attn.fa4 = block.cross_attn.fa4 = True + for name in ("_pre", "_post_self", "_norm3", "_ffn"): + setattr(block, name, compiled(getattr(block, name))) + block.self_attn._qkv = compiled(block.self_attn._qkv) + block.cross_attn._q = compiled(block.cross_attn._q) def freqs(self, device: torch.device) -> torch.Tensor: if device not in self._freqs: @@ -378,6 +431,7 @@ def _forward(self, x, t, context, y, plucker, cache, current_start, frame_seqlen x = x.flatten(2).transpose(1, 2) seq_len = x.size(1) + per_token_t = None if t.dim() == 1 else t if t.dim() == 1: t = t.expand(t.size(0), seq_len) with torch.amp.autocast("cuda", dtype=torch.float32): @@ -385,6 +439,12 @@ def _forward(self, x, t, context, y, plucker, cache, current_start, frame_seqlen e = self.time_embedding( sinusoidal_embedding_1d(cfg.freq_dim, t.flatten()).unflatten(0, (bt, seq_len)).float()) e0 = self.time_projection(e).unflatten(2, (6, cfg.dim)) + if per_token_t is None: + # One timestep for the whole chunk: every token's row of e/e0 is + # the same bits, so keep one and let the modulation broadcast. The + # per-token [L, 6, dim] fp32 tensors otherwise cost ~10% of a + # forward in elementwise traffic, in every block. + e, e0 = e[:, :1], e0[:, :1] if not cache.cross_k: padded = torch.cat([context, context.new_zeros(cfg.text_len - context.size(0), context.size(1))]) @@ -394,12 +454,17 @@ def _forward(self, x, t, context, y, plucker, cache, current_start, frame_seqlen cache.cross_k.append(k) cache.cross_v.append(v) - cam = rearrange( - plucker, "1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)", - c1=cfg.patch_size[0], c2=cfg.patch_size[1], c3=cfg.patch_size[2], - ).contiguous() - cam = self.patch_embedding_wancamctrl(cam) - cam = cam + self.c2ws_hidden_states_layer2(F.silu(self.c2ws_hidden_states_layer1(cam))) + # A chunk's denoise and commit forwards share its Plucker rays, so the + # camera path (four GEMMs a block) runs once per chunk. + if cache.cam_rays is not plucker: + cam = rearrange( + plucker, "1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)", + c1=cfg.patch_size[0], c2=cfg.patch_size[1], c3=cfg.patch_size[2], + ).contiguous() + cam = self.patch_embedding_wancamctrl(cam) + cam = cam + self.c2ws_hidden_states_layer2(F.silu(self.c2ws_hidden_states_layer1(cam))) + cache.cam_mod = [block.cam_modulation(cam) for block in self.blocks] + cache.cam_rays = plucker plan = cache.plan(current_start, seq_len) start_frame = current_start // frame_seqlen @@ -408,7 +473,7 @@ def _forward(self, x, t, context, y, plucker, cache, current_start, frame_seqlen f"latent frame {start_frame + grid[0]} is past the RoPE table ({freqs.shape[0]} frames)") for i, block in enumerate(self.blocks): attn_args = dict(grid=grid, freqs=freqs, start_frame=start_frame, cache=cache, layer=i, plan=plan) - x = block(x, e0, cam, cache.cross_k[i], cache.cross_v[i], attn_args) + x = block(x, e0, cache.cam_mod[i], cache.cross_k[i], cache.cross_v[i], attn_args) cache.global_end = plan.current_end cache.local_end = plan.local_end diff --git a/mstar/model/lingbot_world/config.py b/mstar/model/lingbot_world/config.py index b057e08ca4..1c40b1ddd3 100644 --- a/mstar/model/lingbot_world/config.py +++ b/mstar/model/lingbot_world/config.py @@ -74,7 +74,23 @@ class LingBotWorldConfig: # The reference decodes in fp32. bf16 is faster on cuDNN >= 9.16 but is not # the reference's numerics. vae_decode_fp32: bool = True - vae_decode_channels_last: bool = False + # Channels-last, compiled fp32 decode in Inductor's deterministic mode: + # 547 -> 347 ms per 16-frame chunk at 480x832 on an H200, frames within + # 73 dB of the contiguous eager decode (accumulation order only), the same + # in every process. cuDNN autotuning takes it to 230 ms but picks conv + # algorithms per process, so the model turns it on only with + # reference_numerics False. + vae_decode_channels_last: bool = True + vae_decode_cudnn_benchmark: bool = False + vae_decode_compile: bool = True + + # False trades bit-parity with the reference for speed: flash-attn-4 + # attention (needs flash-attn-4, see docs/installation.rst), compiled block + # pointwise ops, and a cuDNN-autotuned condition encoder and bf16 decode: + # 12.6 -> 18.6 fps at 480x832 on an H200. A free-running rollout then drifts + # from the reference's about as far as it does under any other attention + # kernel (24.6 dB over 93 frames; FA4 instead of FA2 alone gives 26.5 dB). + reference_numerics: bool = True @property def head_dim(self) -> int: diff --git a/mstar/model/lingbot_world/lingbot_world_model.py b/mstar/model/lingbot_world/lingbot_world_model.py index 014790c4f9..8383385e60 100644 --- a/mstar/model/lingbot_world/lingbot_world_model.py +++ b/mstar/model/lingbot_world/lingbot_world_model.py @@ -61,11 +61,16 @@ def __init__( cache_dir: str | None = None, variant: str = LINGBOT_VARIANT_1_3B, skip_weight_loading: bool = False, + reference_numerics: bool = True, **kwargs, ): if variant != LINGBOT_VARIANT_1_3B: raise NotImplementedError(f"LingBot-World variant {variant!r} is not implemented yet.") self.config = LINGBOT_CONFIGS[variant]() + self.config.reference_numerics = bool(reference_numerics) + if not reference_numerics: + self.config.vae_decode_fp32 = False + self.config.vae_decode_cudnn_benchmark = True self.model_path_hf = model_path_hf or LINGBOT_HF_REPOS[variant] self.cache_dir = cache_dir self.skip_weight_loading = skip_weight_loading @@ -378,6 +383,8 @@ def _create_submodule(self, node_name: str, device: str) -> NodeSubmodule | None from mstar.model.lingbot_world.weight_loader import build_lingbot_dit require_flash_attn() + if not self.config.reference_numerics: + import flash_attn.cute # noqa: F401 (flash-attn-4; docs/installation.rst) dit = build_lingbot_dit(self.config, self.model_path_hf, device=device, cache_dir=self.cache_dir) return LingBotDitSubmodule( dit, self._load_vae().to(device), self.config, encoder_vae=self._load_vae().to(device), diff --git a/mstar/model/lingbot_world/submodules.py b/mstar/model/lingbot_world/submodules.py index 217244cc02..f53a646d08 100644 --- a/mstar/model/lingbot_world/submodules.py +++ b/mstar/model/lingbot_world/submodules.py @@ -169,9 +169,15 @@ def __init__(self, dit: LingBotCausalDiT, vae: nn.Module, config: LingBotWorldCo self.encoder_vae = encoder_vae self.config = config self.schedule = make_fast_schedule(config.sample_shift, config.timesteps_index, config.num_train_timesteps) + if vae is not None and not config.vae_decode_fp32: + # The latents are cast to bf16 for the decode, so the weights must be too. + vae.to(torch.bfloat16) if vae is not None and config.vae_decode_channels_last: wan_vae.to_channels_last(vae) - self._decode_step = wan_vae.DecoderStep(vae, compile=False) if vae is not None else None + self._decode_step = ( + wan_vae.DecoderStep(vae, compile=config.vae_decode_compile, deterministic=True) + if vae is not None else None + ) # Debug hook for the parity harness: keep the last chunk's latents and # float frames in the request state. self.keep_chunk_outputs = False @@ -228,7 +234,10 @@ def forward( cs = cfg.chunk_size frames = slice(chunk * cs, (chunk + 1) * cs) latent = state["noise"][:, frames] - y = encode_condition_chunk(self.encoder_vae, image, state["encoder_stream"], chunk * cs, cs, cfg) + # Autotuned convs pick other algorithms than the reference's, which a + # rollout amplifies; only off reference numerics. + with wan_vae.cudnn_benchmark(not cfg.reference_numerics): + y = encode_condition_chunk(self.encoder_vae, image, state["encoder_stream"], chunk * cs, cs, cfg) _, height, width = image.shape plucker = plucker_embedding( camera_poses[frames].float().to(device), intrinsics.float(), height, width, @@ -271,7 +280,7 @@ def _decode(self, latents: torch.Tensor, state, clip_start: bool) -> torch.Tenso mean = torch.tensor(self.vae.config.latents_mean).view(1, cfg.vae_z_dim, 1, 1, 1).to(z) inv_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, cfg.vae_z_dim, 1, 1, 1).to(z) z = z / inv_std + mean - with torch.amp.autocast("cuda", enabled=False): + with torch.amp.autocast("cuda", enabled=False), wan_vae.cudnn_benchmark(cfg.vae_decode_cudnn_benchmark): return wan_vae.decode_causal( self.vae, z, state["feat_cache"], self._decode_step, clip_start=clip_start, latents_per_call=1, ) diff --git a/mstar/model/media_limits.py b/mstar/model/media_limits.py index f22bb75161..b4e2fa2e3b 100644 --- a/mstar/model/media_limits.py +++ b/mstar/model/media_limits.py @@ -15,6 +15,11 @@ import torch +# FFmpeg threads per video decode. torchcodec decodes single-threaded by +# default, which dominates a long clip's request latency; frame threading is +# bit-identical. Capped so concurrent requests don't oversubscribe the host. +VIDEO_DECODE_THREADS = min(8, os.cpu_count() or 1) + # PIL's DecompressionBombWarning threshold (about 9459 x 9459 RGB) DEFAULT_MAX_IMAGE_PIXELS = 89_478_485 # about 39 s of 720p30, or 2.6 min of 640x360 at 30 fps, decoded at once: a diff --git a/mstar/model/omnilingual_asr/components/wav2vec2.py b/mstar/model/omnilingual_asr/components/wav2vec2.py index 55033db803..00d0f33d71 100644 --- a/mstar/model/omnilingual_asr/components/wav2vec2.py +++ b/mstar/model/omnilingual_asr/components/wav2vec2.py @@ -106,11 +106,50 @@ def __init__(self, model_dim: int, kernel_size: int, num_groups: int): bias=True, ) self.remove_pad = kernel_size % 2 == 0 + # Set by a server on an accelerator: run the conv as per-group GEMMs + # (see `_conv_unfolded`) instead of through cuDNN + self.unfolded = False + + # elements of the unfolded input one GEMM chunk may materialize (1 GiB in bf16) + UNFOLD_CHUNK_ELEMS = 1 << 29 + + def _conv_unfolded(self, x: torch.Tensor) -> torch.Tensor: + """``conv`` on channels-last ``(N, T, C)`` as one batched GEMM per chunk + of frames: the windows of each group, ``(N * T, C/G * K)``, times that + group's weight. The same products as the conv, summed in another order. + + cuDNN's grouped kernels for this shape (wide groups, a 128-tap kernel) + are erratic: 0.8 ms or 35 ms for a single clip at 2048 channels, + depending on its length.""" + conv = self.conv + n, t, c = x.shape + groups, (kernel,), (pad,) = conv.groups, conv.kernel_size, conv.padding + per_group = c // groups + # group-major first (a copy of x), so the windows below take one copy + xg = F.pad(x, (0, 0, pad, pad)).view(n, t + 2 * pad, groups, per_group) + xg = xg.permute(2, 0, 1, 3).contiguous() # (G, N, T + 2 * pad, C/G) + t_out = t + 2 * pad - kernel + 1 + weight = conv.weight.reshape(groups, per_group, per_group * kernel).transpose(1, 2) + chunk = max(1, self.UNFOLD_CHUNK_ELEMS // (n * c * kernel)) + outs = [] + for start in range(0, t_out, chunk): + stop = min(start + chunk, t_out) + # (G, N, T', C/G, K) -> (G, N * T', C/G * K), k fastest like the weight + win = xg[:, :, start:stop + kernel - 1].unfold(2, kernel, 1) + y = torch.bmm(win.reshape(groups, n * (stop - start), per_group * kernel), weight) + outs.append(y.view(groups, n, stop - start, per_group).permute(1, 2, 0, 3).reshape(n, stop - start, c)) + y = outs[0] if len(outs) == 1 else torch.cat(outs, dim=1) + return y + conv.bias if conv.bias is not None else y def forward(self, x: torch.Tensor, frame_mask: torch.Tensor | None) -> torch.Tensor: if frame_mask is not None: # padded frames must be zero, or they leak into the valid ones' encodings x = x * frame_mask.unsqueeze(-1).to(x.dtype) + if self.unfolded: + enc = self._conv_unfolded(x) + if self.remove_pad: + enc = enc[:, :-1] + return x + F.gelu(enc) enc = self.conv(x.transpose(1, 2)) if self.remove_pad: enc = enc[:, :, :-1] @@ -135,7 +174,14 @@ def __init__(self, model_dim: int, num_heads: int): self.v_proj = nn.Linear(model_dim, model_dim) self.output_proj = nn.Linear(model_dim, model_dim) - def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor: + def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None, attn_fn=None) -> torch.Tensor: + if attn_fn is not None: + # packed rows: (T, M) -> (T, H, K), attending within each row's span + q, k, v = ( + proj(x).unflatten(-1, (self.num_heads, -1)) + for proj in (self.q_proj, self.k_proj, self.v_proj) + ) + return self.output_proj(attn_fn(q, k, v).flatten(-2)) # (N, S, M) -> (N, S, H, K) -> (N, H, S, K) q = self.q_proj(x).unflatten(-1, (self.num_heads, -1)).transpose(1, 2) k = self.k_proj(x).unflatten(-1, (self.num_heads, -1)).transpose(1, 2) @@ -164,8 +210,8 @@ def __init__(self, cfg): self.ffn_layer_norm = nn.LayerNorm(cfg.model_dim, eps=cfg.layer_norm_eps) self.ffn = FeedForward(cfg.model_dim, cfg.ffn_inner_dim) - def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor: - x = self.self_attn(self.self_attn_layer_norm(x), attn_mask) + x + def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None, attn_fn=None) -> torch.Tensor: + x = self.self_attn(self.self_attn_layer_norm(x), attn_mask, attn_fn) + x return self.ffn(self.ffn_layer_norm(x)) + x @@ -175,9 +221,9 @@ def __init__(self, cfg): self.layers = nn.ModuleList(EncoderLayer(cfg) for _ in range(cfg.num_layers)) self.layer_norm = nn.LayerNorm(cfg.model_dim, eps=cfg.layer_norm_eps) - def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor: + def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None, attn_fn=None) -> torch.Tensor: for layer in self.layers: - x = layer(x, attn_mask) + x = layer(x, attn_mask, attn_fn) return self.layer_norm(x) @@ -198,6 +244,33 @@ def __init__(self, cfg): self.encoder = Encoder(cfg) self.final_proj = nn.Linear(cfg.model_dim, cfg.vocab_size) + def frontend( + self, waveforms: torch.Tensor, seq_lens: list[int] | None = None, features=None, + ) -> tuple[torch.Tensor, list[int], torch.Tensor | None]: + """``(N, S)`` waveforms -> ``(N, T, model_dim)`` encoder inputs (padded + frames zeroed), each row's valid frame count, and the ``(N, T)`` frame + mask (None when no row is padded). ``features`` stands in for the + feature extractor (e.g. a compiled copy of it).""" + front = self.encoder_frontend + x = (front.feature_extractor if features is None else features)(waveforms) + n, t = x.shape[:2] + frame_lens = [t] * n if seq_lens is None else self.frame_lengths(seq_lens) + frame_mask = None + if any(f < t for f in frame_lens): + lens = torch.tensor(frame_lens) + if x.is_cuda: + lens = lens.pin_memory() + lens = lens.to(x.device, non_blocking=True) + frame_mask = torch.arange(t, device=x.device)[None, :] < lens[:, None] + x = front.model_dim_proj(front.post_extract_layer_norm(x)) + return front.pos_encoder(x, frame_mask), frame_lens, frame_mask + + def encode_packed(self, x: torch.Tensor, attn_fn) -> torch.Tensor: + """``(T, model_dim)``, every row's valid frames end to end -> ``(T, vocab)`` + logits. ``attn_fn(q, k, v)`` attends within each row's span (a varlen + kernel): no padding reaches the encoder.""" + return self.final_proj(self.encoder(x, None, attn_fn)) + def frame_lengths(self, seq_lens: list[int]) -> list[int]: return [self.encoder_frontend.feature_extractor.output_length(n) for n in seq_lens] @@ -205,25 +278,14 @@ def forward( self, waveforms: torch.Tensor, seq_lens: list[int] | None = None, + features=None, ) -> tuple[torch.Tensor, list[int]]: """``(N, S)`` waveforms (and each row's sample count, for a padded batch) -> ``(N, T, vocab)`` logits and each row's valid frame count. The lengths are host ints, so a padded batch is masked without reading anything back from the device.""" - front = self.encoder_frontend - x = front.feature_extractor(waveforms) - n, t = x.shape[:2] - frame_lens = [t] * n if seq_lens is None else self.frame_lengths(seq_lens) - frame_mask = None - if any(f < t for f in frame_lens): - lens = torch.tensor(frame_lens) - if x.is_cuda: - lens = lens.pin_memory() - lens = lens.to(x.device, non_blocking=True) - frame_mask = torch.arange(t, device=x.device)[None, :] < lens[:, None] - x = front.model_dim_proj(front.post_extract_layer_norm(x)) - x = front.pos_encoder(x, frame_mask) + x, frame_lens, frame_mask = self.frontend(waveforms, seq_lens, features) # cuDNN attention builds a plan per new sequence length, and nearly # every upload has its own length with sdpa_kernel(_ATTENTION_BACKENDS): diff --git a/mstar/model/omnilingual_asr/config.py b/mstar/model/omnilingual_asr/config.py index 40206ee7bf..bd66424283 100644 --- a/mstar/model/omnilingual_asr/config.py +++ b/mstar/model/omnilingual_asr/config.py @@ -35,6 +35,8 @@ TRANSCRIBE_WALK = "transcribe" WAVEFORM = "waveform" TOKEN_IDS = "token_ids" +# the encoder's varlen self-attention (a ragged attention resource) +ENCODER_ATTN = "encoder_attn" SIZES = ("300m", "1b", "3b", "7b") VERSIONS = ("v1", "v2") diff --git a/mstar/model/omnilingual_asr/omnilingual_asr_model.py b/mstar/model/omnilingual_asr/omnilingual_asr_model.py index 9056f66e44..4beea6424d 100644 --- a/mstar/model/omnilingual_asr/omnilingual_asr_model.py +++ b/mstar/model/omnilingual_asr/omnilingual_asr_model.py @@ -38,6 +38,7 @@ from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata from mstar.model.omnilingual_asr.audio import load_waveform, preprocess from mstar.model.omnilingual_asr.config import ( + ENCODER_ATTN, ENCODER_NODE, MAX_AUDIO_SECONDS, SAMPLE_RATE, @@ -120,8 +121,27 @@ def __init__( # -- graph ------------------------------------------------------------- def get_node_resources(self) -> list[NodeResourceSpec]: - """Not autoregressive: no KV cache, no sampler.""" - return [] + """Not autoregressive: no KV cache, no sampler. A half-precision encoder + attends through a varlen kernel over the step's packed rows; fp32 (the + reference's numerics) keeps the padded SDPA path.""" + if self.dtype not in (torch.bfloat16, torch.float16): + return [] + from mstar.engine.resources.attn.ragged.config import RaggedAttentionConfig, RaggedAttentionSpec + + return [ + RaggedAttentionSpec( + resource_key=ENCODER_ATTN, nodes={ENCODER_NODE}, + config=RaggedAttentionConfig( + num_qo_heads=self.config.num_heads, + num_kv_heads=self.config.num_heads, + head_dim=self.config.model_dim // self.config.num_heads, + dtype=self.dtype, + # one segment per request, and a captured step packs a + # whole batch into one row + max_segments_per_request=self.max_batch_size, + ), + ), + ] def get_graph_walk_graphs(self) -> dict[str, GraphSection]: return { @@ -238,6 +258,12 @@ def _create_encoder_submodule(self, device: str) -> NodeSubmodule: model = Wav2Vec2CTC(self.config) load_into(model, convert_state_dict(read_checkpoint(path), self.dtype, device)) model.eval() + pos = model.encoder_frontend.pos_encoder + # cuDNN's grouped kernels for the positional conv are erratic once its + # groups are wide (1b and up): 0.8-40 ms for one 2048-channel clip by + # length, against ~1.3 ms per 1000 frames as per-group GEMMs. At 300m's + # 64-channel groups cuDNN is the faster (0.4-1.5 ms a step). + pos.unfolded = torch.device(device).type != "cpu" and pos.conv.in_channels // pos.conv.groups >= 80 logger.info( "Loaded Omnilingual ASR CTC %s (%s, %s) from %s on %s", self.size, @@ -246,4 +272,6 @@ def _create_encoder_submodule(self, device: str) -> NodeSubmodule: Path(path).name, device, ) - return OmniASRCTCSubmodule(model, self.max_batch_size, self.max_batch_seconds) + submodule = OmniASRCTCSubmodule(model, self.max_batch_size, self.max_batch_seconds) + submodule.warm_compile(device) + return submodule diff --git a/mstar/model/omnilingual_asr/submodules.py b/mstar/model/omnilingual_asr/submodules.py index 99cb43f570..fae5b48cad 100644 --- a/mstar/model/omnilingual_asr/submodules.py +++ b/mstar/model/omnilingual_asr/submodules.py @@ -19,6 +19,7 @@ from __future__ import annotations +from collections.abc import Mapping from contextlib import contextmanager from typing import Any @@ -26,10 +27,20 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.accelerator_graph_config import ( + PiecewiseAcceleratorGraphConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, + PiecewisePackedConfig, +) from mstar.engine.engine import ExecutingBatch +from mstar.engine.resources.attn.config import AttentionStep +from mstar.engine.resources.convenience import RaggedAttentionCallable +from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep from mstar.model.omnilingual_asr.components.wav2vec2 import Wav2Vec2CTC -from mstar.model.omnilingual_asr.config import SAMPLE_RATE, TOKEN_IDS, WAVEFORM +from mstar.model.omnilingual_asr.config import ENCODER_ATTN, SAMPLE_RATE, TOKEN_IDS, WAVEFORM from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule +from mstar.platforms.compile import with_inductor_overrides @contextmanager @@ -48,20 +59,49 @@ def _full_fp32_matmuls(enabled: bool): torch.set_float32_matmul_precision(previous) +# The packed encoder (layers, final norm, CTC head, argmax) as a piecewise +# capture region +ENCODER_LOOP = "encoder_loop" + + class OmniASRCTCSubmodule(NodeSubmodule): - # Every request has its own length, so a compiled forward would recompile - # per shape; the encoder runs eagerly. + # Every request has its own length, so the engine's compile of the whole + # forward would recompile per shape; the encoder runs eagerly. The conv + # feature extractor alone is compiled (below), dynamic from the first call. disable_torch_compile = True # The model runs in its own dtype end to end, as the reference does (only # the conv stack's LayerNorms upcast to fp32), not under autocast. disable_autocast = True + # Frames of one packed encoder replay (50 a second); a step past the top + # bucket runs the encoder eagerly. 16384 frames is 320 s, the default + # `max_batch_seconds`. + ENCODER_FRAME_BUCKETS = [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384] + def __init__(self, model: Wav2Vec2CTC, max_batch_size: int, max_batch_seconds: float): super().__init__() self.model = model self.dtype = next(model.parameters()).dtype self._max_batch_size = max_batch_size self._max_batch_samples = int(max_batch_seconds * SAMPLE_RATE) + self._max_batch_frames = model.frame_lengths([self._max_batch_samples])[0] + # Eager, each conv layer's channel LayerNorm (in fp32, over a transposed + # copy) costs more than the conv: about half the GPU time of a batched + # step. Inductor fuses the transposes, the norm and the GELU, under the + # platform's inductor settings as the engine's compiles are. + self._compiled_features = with_inductor_overrides( + torch.compile(model.encoder_frontend.feature_extractor, dynamic=True) + ) + self._attn_fns: dict[int, RaggedAttentionCallable] = {} + + def warm_compile(self, device: str) -> None: + """Compile the feature extractor at load rather than on a live request: + one graph for a lone clip (size-1 dims specialize) and the dynamic one.""" + if torch.device(device).type == "cpu": + return + with torch.no_grad(): + for rows in (1, 2): + self._compiled_features(torch.zeros(rows, SAMPLE_RATE, dtype=self.dtype, device=device)) def prepare_inputs( self, @@ -72,7 +112,80 @@ def prepare_inputs( ) -> NodeInputs: del graph_walk, fwd_info, kwargs waveform = inputs[WAVEFORM][0] - return NodeInputs(tensor_inputs={WAVEFORM: waveform}, input_seq_len=int(waveform.shape[-1])) + # in encoder frames, what the captured encoder is leased by + frames = self.model.frame_lengths([int(waveform.shape[-1])])[0] + return NodeInputs(tensor_inputs={WAVEFORM: waveform}, input_seq_len=frames) + + def declare_step( + self, + graph_walk: str, + request_ids: list[str], + inputs: list[NodeInputs], + piecewise_leases: Mapping[str, SlotLease] | None = None, + **kwargs: Any, + ) -> SubmoduleStep | None: + """One attention segment per request, its valid frames, when the + encoder attends packed (see `get_node_resources`). Leased, the captured + encoder plans its own.""" + if ( + ENCODER_ATTN not in self.node_resources + or (piecewise_leases or {}).get(ENCODER_LOOP) + or not self._packs([inp.input_seq_len for inp in inputs]) + ): + return None + return SubmoduleStep( + segments=[ + Segment(request_id=rid, label="main", span=inp.input_seq_len) + for rid, inp in zip(request_ids, inputs, strict=True) + ], + steps={ENCODER_ATTN: AttentionStep(causal=False)}, + ) + + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, **kwargs: Any, + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + """The packed encoder, one replay a step: every request's frames are a + segment of one row (`one_row_per_step`). Eager, the encoder is a few + hundred small launches, which bounds a lone clip's latency.""" + if ENCODER_ATTN not in self.node_resources: + return {} + model_dim = self.model.config.model_dim + + def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return {"hidden": torch.zeros(shape.total_tokens, model_dim, dtype=self.dtype, device=device)} + + def declare_step(request_ids: list[str], seq_lens: list[int]) -> SubmoduleStep: + (rid,) = request_ids + return SubmoduleStep( + segments=[Segment(request_id=rid, label="main", span=n) for n in seq_lens], + steps={ENCODER_ATTN: AttentionStep(causal=False)}, + ) + + return { + ENCODER_LOOP: PiecewisePackedConfig( + capture_fn=self._capture_encoder, + make_static_inputs=make_static_inputs, + declare_step=declare_step, + lease_before_step=True, + can_lease=lambda inputs: self._packs([inp.input_seq_len for inp in inputs]), + one_row_per_step=True, + total_tokens=self.ENCODER_FRAME_BUCKETS, + capture_batch_sizes=[1], + ) + } + + def _capture_encoder(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + """The captured region: reads the runner's buffer, never rebinds.""" + attn = self._attn_fn(inp.resources[ENCODER_ATTN]) + return {"ids": self.model.encode_packed(inp.static_inputs["hidden"], attn).argmax(dim=-1)} + + def _attn_fn(self, resource) -> RaggedAttentionCallable: + """The resource's attention over the step's segments, one callable per + bound resource.""" + fn = self._attn_fns.get(id(resource)) + if fn is None: + fn = self._attn_fns[id(resource)] = RaggedAttentionCallable(resource, "main") + return fn def max_batch_size(self, graph_walk: str) -> int: del graph_walk @@ -104,9 +217,61 @@ def sub_batches(self, lengths: list[int]) -> list[list[int]]: groups.append([row]) return groups - def _transcribe(self, waveforms: list[torch.Tensor]) -> list[torch.Tensor]: + def _packs(self, frames: list[int]) -> bool: + """Whether a step's rows encode packed, in one call: their frames fit + the audio budget. A step over it encodes padded per sub-batch instead, + as the fp32 path does; packing each sub-batch would need an attention + plan per call.""" + return sum(frames) <= self._max_batch_frames + + def _features_fn(self, batch: torch.Tensor): + # CPU (tests) stays eager + return self._compiled_features if batch.device.type != "cpu" else None + + def _transcribe_packed( + self, waveforms: list[torch.Tensor], attn, engine_inputs: ModelInputsFromEngine, + ) -> list[torch.Tensor]: + """Argmax ids per row: the conv frontend runs padded per sub-batch, then + every row's valid frames go through the encoder packed, in request order + (the order `declare_step` laid the segments out in).""" + lengths = [w.shape[0] for w in waveforms] + rows: list[torch.Tensor | None] = [None] * len(waveforms) + for group in self.sub_batches(lengths): + group_lens = [lengths[row] for row in group] + longest = max(group_lens) + batch = waveforms[group[0]].new_zeros(len(group), longest) + for i, row in enumerate(group): + batch[i, : lengths[row]] = waveforms[row] + padded = any(n < longest for n in group_lens) + x, frame_lens, _ = self.model.frontend( + batch, group_lens if padded else None, features=self._features_fn(batch), + ) + for i, (row, n) in enumerate(zip(group, frame_lens, strict=True)): + rows[row] = x[i, :n] + packed = rows[0] if len(rows) == 1 else torch.cat(rows) + frames = [r.shape[0] for r in rows] + runner = engine_inputs.piecewise_runners.get(ENCODER_LOOP) + if runner is not None and runner.can_run(1, packed.shape[0]): + # all of the step's segments under one row; a copy, since the + # output buffer is overwritten by the next step's replay + ids = runner.run( + static_inputs={"hidden": packed}, + request_ids=[engine_inputs.request_ids[0]], + seq_lens=frames, + real_bs=1, + ).get("ids") + else: + ids = self.model.encode_packed(packed, self._attn_fn(attn)).argmax(dim=-1) + return list(ids.split(frames)) + + def _transcribe( + self, waveforms: list[torch.Tensor], engine_inputs: ModelInputsFromEngine | None = None, + ) -> list[torch.Tensor]: """Argmax ids per row, over the row's valid frames.""" + attn = getattr(engine_inputs, "resources", {}).get(ENCODER_ATTN) lengths = [w.shape[0] for w in waveforms] + if attn is not None and self._packs(self.model.frame_lengths(lengths)): + return self._transcribe_packed(waveforms, attn, engine_inputs) out: list[torch.Tensor | None] = [None] * len(waveforms) for group in self.sub_batches(lengths): group_lens = [lengths[row] for row in group] @@ -116,7 +281,9 @@ def _transcribe(self, waveforms: list[torch.Tensor]) -> list[torch.Tensor]: batch[i, : lengths[row]] = waveforms[row] padded = any(n < longest for n in group_lens) with _full_fp32_matmuls(self.dtype == torch.float32 and batch.is_cuda): - logits, frame_lens = self.model(batch, group_lens if padded else None) + logits, frame_lens = self.model( + batch, group_lens if padded else None, features=self._features_fn(batch), + ) ids = logits.argmax(dim=-1) for i, (row, n) in enumerate(zip(group, frame_lens, strict=True)): out[row] = ids[i, :n] @@ -129,8 +296,8 @@ def forward( waveforms: list[torch.Tensor], **kwargs: Any, ) -> NameToTensorList: - del graph_walk, engine_inputs, kwargs - return {TOKEN_IDS: [self._transcribe(waveforms)[0]]} + del graph_walk, kwargs + return {TOKEN_IDS: [self._transcribe(waveforms, engine_inputs)[0]]} def forward_batched( self, @@ -140,5 +307,5 @@ def forward_batched( **kwargs: Any, ) -> dict[str, NameToTensorList]: del graph_walk, kwargs - rows = self._transcribe(waveforms) + rows = self._transcribe(waveforms, engine_inputs) return {rid: {TOKEN_IDS: [ids]} for rid, ids in zip(engine_inputs.request_ids, rows, strict=True)} diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index ac26f79cc2..db80b0be2c 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -6,11 +6,13 @@ Nothing is padded, so a batch mixing a 2-second and a 25-second utterance wastes nothing, and the head GEMM runs only at the positions whose logits are read. -The per-step kernel path is the reference's own fastest mode, reused rather than -re-derived: ``apply_flashinfer`` fuses QKV, swaps in flashinfer RoPE, RMSNorm -and ``silu_and_mul``, and keeps the whole layer in ``(S, H, D)`` so no transpose -copies happen. Re-implementing that here would mean a second copy of tuned -kernels whose only job is to end up numerically identical. +The per-step kernels are the reference's own fastest mode: ``apply_flashinfer`` +fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE, +RMSNorm and ``silu_and_mul``. ``OmniVoiceBackbone.encode`` drives those fused +weights with its own layer loop rather than through the patched ``llm``: it +norms q and k straight off the fused QKV output (the reference copies each out +first) and folds every residual add into the next RMSNorm. At 16 requests a +step those copies and adds were a tenth of the step's GPU time. What M* adds on top is what the reference cannot do from inside one ``generate()`` call: the documents in a step come from *different concurrent @@ -18,8 +20,10 @@ own server has no batching at all (k2-fsa/OmniVoice#251). The coupling to ``omnivoice.models.omnivoice_flashinfer`` is deliberate but -narrow — four module-private names, checked at load by ``assert_flashinfer_api`` -so an upstream rename fails at boot rather than at the first request. +narrow — four module-private names and the four fused-weight attributes its +``apply_flashinfer`` sets on each layer, checked at load by +``assert_flashinfer_api`` and ``assert_flashinfer_patched`` so an upstream +rename fails at boot rather than at the first request. """ import logging @@ -27,12 +31,20 @@ from dataclasses import dataclass, field import torch +import torch.nn.functional as F from torch import nn logger = logging.getLogger(__name__) -# The private surface this integration rides on. Checked once at load. +# The private surface this integration rides on. Checked once at load: the +# module's names before ``apply_flashinfer`` runs, and the fused weights it +# leaves on every layer's attention and MLP (read by ``OmniVoiceBackbone.encode``) +# after. _REQUIRED_FI_NAMES = ("_CTX", "PackedAttnRunner", "_forward_logits", "apply_flashinfer") +_REQUIRED_FI_LAYER_ATTRS = { + "self_attn": ("_fi_w_qkv", "_fi_qkv_split", "_fi_rope_theta"), + "mlp": ("_fi_w_gate_up",), +} def assert_flashinfer_api(): @@ -64,6 +76,24 @@ def assert_flashinfer_api(): return fi +def assert_flashinfer_patched(model: nn.Module) -> None: + """Fail at boot if ``apply_flashinfer`` no longer leaves the fused weights + ``OmniVoiceBackbone.encode`` reads, rather than at capture.""" + missing = sorted({ + f"{part}.{name}" + for layer in model.llm.layers + for part, names in _REQUIRED_FI_LAYER_ATTRS.items() + for name in names + if not hasattr(getattr(layer, part), name) + }) + if missing: + raise RuntimeError( + f"omnivoice's apply_flashinfer no longer sets {missing} on the layers; " + "mstar's OmniVoice layer loop reads them. Pin the omnivoice version, " + "or port the fused weights into mstar." + ) + + @dataclass class CanvasItem: """One request's contribution to a step's packed batch. @@ -159,78 +189,60 @@ def build_packed_canvas( if not items: raise ValueError("build_packed_canvas called with no items") - num_codebook = items[0].prefix_ids.shape[0] - + # One concatenation per device tensor, and the index tensors built on the + # host and copied once: a step packs up to 16 requests, and ten small + # kernels per request were a fifth of the step's host time. + id_parts: list[torch.Tensor] = [] + mask_parts: list[torch.Tensor] = [] + positions: list[torch.Tensor] = [] doc_lens: list[int] = [] - offsets: list[tuple[int, int]] = [] # (cond_offset, uncond_offset or -1) - cursor = 0 - for item in items: - c_off = cursor - cursor += item.cond_len - doc_lens.append(item.cond_len) - if item.does_cfg: - u_off = cursor - cursor += item.target_len - doc_lens.append(item.target_len) - else: - u_off = -1 - offsets.append((c_off, u_off)) - total = cursor - - packed_ids = torch.full( - (1, num_codebook, total), audio_mask_id, dtype=torch.long, device=device - ) - audio_mask = torch.zeros((1, total), dtype=torch.bool, device=device) - position_ids = torch.zeros((1, total), dtype=torch.long, device=device) + cond_rows: list[torch.Tensor] = [] + uncond_rows: list[torch.Tensor] = [] + longest = max(item.target_len for item in items) + revealed = torch.ones(longest, dtype=torch.bool, device=device) - cond_ranges: list[torch.Tensor] = [] - uncond_ranges: list[torch.Tensor] = [] - flat_cursor = 0 - flat_u_cursor = 0 - - for item, (c_off, u_off) in zip(items, offsets, strict=True): + cursor = flat_cursor = flat_u_cursor = 0 + for item in items: c_len, t_len = item.cond_len, item.target_len - - prefix_ids = item.prefix_ids.to(device) tokens = item.tokens[0].to(device) + target_mask = revealed[:t_len] # Conditional document: [style | text | ref | target]. - packed_ids[0, :, c_off : c_off + c_len] = torch.cat([prefix_ids, tokens], dim=-1) - audio_mask[0, c_off : c_off + c_len] = torch.cat( - [ - item.prefix_audio_mask.to(device), - torch.ones(t_len, dtype=torch.bool, device=device), - ], - dim=-1, - ) - position_ids[0, c_off : c_off + c_len] = torch.arange(c_len, device=device) - + id_parts += [item.prefix_ids.to(device), tokens] + mask_parts += [item.prefix_audio_mask.to(device), target_mask] + positions.append(torch.arange(c_len)) + doc_lens.append(c_len) item.flat_start = flat_cursor flat_cursor += t_len - cond_ranges.append( - torch.arange(c_off + c_len - t_len, c_off + c_len, device=device) - ) + cond_rows.append(torch.arange(cursor + c_len - t_len, cursor + c_len)) + cursor += c_len - if u_off < 0: + if not item.does_cfg: item.flat_u_start = -1 continue # Unconditional document: the target canvas alone. Dropping the # conditioning entirely *is* the null prompt here. - packed_ids[0, :, u_off : u_off + t_len] = tokens - audio_mask[0, u_off : u_off + t_len] = True - position_ids[0, u_off : u_off + t_len] = torch.arange(t_len, device=device) - + id_parts.append(tokens) + mask_parts.append(target_mask) + positions.append(torch.arange(t_len)) + doc_lens.append(t_len) item.flat_u_start = flat_u_cursor flat_u_cursor += t_len - uncond_ranges.append(torch.arange(u_off, u_off + t_len, device=device)) + uncond_rows.append(torch.arange(cursor, cursor + t_len)) + cursor += t_len + + def to_device(t: torch.Tensor) -> torch.Tensor: + if device.type == "cuda": + return t.pin_memory().to(device, non_blocking=True) + return t.to(device) return PackedCanvas( - packed_ids=packed_ids, - audio_mask=audio_mask, - position_ids=position_ids, + packed_ids=torch.cat(id_parts, dim=-1).unsqueeze(0), + audio_mask=torch.cat(mask_parts).unsqueeze(0), + position_ids=to_device(torch.cat(positions).unsqueeze(0)), doc_lens=doc_lens, - tgt_index=torch.cat(cond_ranges + uncond_ranges), + tgt_index=to_device(torch.cat(cond_rows + uncond_rows)), flat_target_total=flat_cursor, flat_uncond_total=flat_u_cursor, items=items, @@ -240,10 +252,10 @@ def build_packed_canvas( class OmniVoiceBackbone(nn.Module): """The patched reference model, driven one packed step at a time. - Holding the reference instance rather than borrowing its submodules is what - keeps the embedding merge, the audio head reshape and the fused attention - path byte-for-byte the upstream ones; every line of per-step math that could - drift is theirs. + Holding the reference instance keeps the embedding merge, the audio head + reshape and the fused weights the upstream ones. ``encode`` runs the layer + stack over them itself (see the module docstring); ``forward`` is the + reference's own packed forward, the path with backbone graphs off. """ def __init__(self, model: nn.Module, plan_dtype: torch.dtype): @@ -258,6 +270,68 @@ def __init__(self, model: nn.Module, plan_dtype: torch.dtype): # request attending with another's document boundaries. self._ctx_lock = threading.Lock() + def encode( + self, + input_ids: torch.Tensor, + audio_mask: torch.Tensor, + pos_ids: torch.Tensor, + attn, + ) -> torch.Tensor: + """Embeddings and the Qwen3 body over one packed row: ``[N, hidden]``. + + Token-major inputs (``input_ids`` ``[N, C]``, the rest ``[N]``, ``pos_ids`` + int32) so a captured replay can copy a shorter row into its bucket. + ``attn`` is the engine's ragged attention resource, planned for this + row's documents. + """ + import flashinfer + + hidden = self.model._prepare_embed_inputs( + input_ids.t().unsqueeze(0), audio_mask.unsqueeze(0) + )[0] + llm = self.model.llm + residual = None + for layer in llm.layers: + sa, mlp = layer.self_attn, layer.mlp + norm = layer.input_layernorm + if residual is None: + residual = hidden + x = flashinfer.norm.rmsnorm(hidden, norm.weight, eps=norm.variance_epsilon) + else: + # residual += hidden; hidden = norm(residual), in place. The kernel + # normalizes the sum before rounding it to half precision, where + # the reference rounds first: logits match it to that rounding, + # not bit for bit. + flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) + x = hidden + n, d = x.shape[0], sa.head_dim + heads_q, heads_k, _ = (width // d for width in sa._fi_qkv_split) + qkv = F.linear(x, sa._fi_w_qkv).view(n, -1, d) + # the q/k norms read their heads straight off the fused output + q = flashinfer.norm.rmsnorm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) + k = flashinfer.norm.rmsnorm( + qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon, + ) + flashinfer.rope.apply_rope_pos_ids_inplace( + q, k, pos_ids, rope_theta=sa._fi_rope_theta, interleave=False, + ) + hidden = sa.o_proj(attn.run(q, k, qkv[:, heads_q + heads_k:]).reshape(n, -1)) + norm = layer.post_attention_layernorm + flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) + hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(hidden, mlp._fi_w_gate_up))) + flashinfer.norm.fused_add_rmsnorm(hidden, residual, llm.norm.weight, eps=llm.norm.variance_epsilon) + return hidden + + def head(self, hidden: torch.Tensor, tgt_index: torch.Tensor) -> torch.Tensor: + """The audio head at the logits-consuming positions, laid out as + ``_forward_logits`` returns it: ``[1, C, len(tgt_index), V]``.""" + cfg = self.model.config + tgt_hidden = hidden[tgt_index] + logits = self.model.audio_heads(tgt_hidden) + return logits.view( + 1, tgt_hidden.shape[0], cfg.num_audio_codebook, cfg.audio_vocab_size + ).permute(0, 2, 1, 3) + @torch.inference_mode() def forward(self, canvas: PackedCanvas) -> torch.Tensor: """One packed forward; returns ``[1, C, 2 * flat_target_total, V]``. diff --git a/mstar/model/omnivoice/components/codec.py b/mstar/model/omnivoice/components/codec.py index 279bd3a83c..ee9731f1aa 100644 --- a/mstar/model/omnivoice/components/codec.py +++ b/mstar/model/omnivoice/components/codec.py @@ -88,11 +88,12 @@ def post_process( designed or automatic voice has no reference to match and is peak-normalised to 0.5 instead. - Imported from the ``omnivoice`` package rather than re-ported: the silence - trim is pydub-based and re-implementing it would drift from the reference - for no gain. + The fade and pad are the ``omnivoice`` package's own; the silence trim is + a sample-exact vectorized port of its pydub-based one (``silence.py``), + which cost ~14 ms of Python per second of audio. """ - from omnivoice.utils.audio import fade_and_pad_audio, remove_silence + from mstar.model.omnivoice.components.silence import remove_silence + from omnivoice.utils.audio import fade_and_pad_audio audio = waveform if audio.ndim == 1: diff --git a/mstar/model/omnivoice/components/silence.py b/mstar/model/omnivoice/components/silence.py new file mode 100644 index 0000000000..94d700b08e --- /dev/null +++ b/mstar/model/omnivoice/components/silence.py @@ -0,0 +1,146 @@ +"""The reference's silence trim (``omnivoice.utils.audio.remove_silence``), +vectorized and sample-exact. + +The reference runs pydub: an RMS window every 10 ms, each a fresh +``AudioSegment`` slice and an ``audioop.rms`` call, which costs about 14 ms of +Python per second of audio and ran inside the code2wav forward, stalling the +worker's loop for every finished request. This computes the same windows from +one prefix sum of squared int16 samples and replays pydub's millisecond +slicing (``int(ms * rate / 1000)`` samples, zero-filled past the end), its +range merging and its edge trim, so the output is the reference's to the +sample; ``test_omnivoice_silence`` checks it against pydub. +""" + +import itertools + +import numpy as np + +# pydub's dBFS: amplitude relative to the int16 full scale +_FULL_SCALE = 32768.0 + + +class _Segment: + """An int16 mono signal addressed the way ``pydub.AudioSegment`` slices it.""" + + def __init__(self, samples: np.ndarray, rate: int): + self.samples = samples + self.rate = rate + self.per_ms = rate / 1000.0 + squares = samples.astype(np.int64) ** 2 + self.prefix = np.concatenate([[0], np.cumsum(squares)]) + + def __len__(self) -> int: + # milliseconds, rounded as pydub rounds them + return round(1000 * (len(self.samples) / self.rate)) + + def bounds(self, start_ms, end_ms): + """Sample bounds and zero-filled length of ``segment[start_ms:end_ms]``.""" + n_ms = len(self) + start = (np.minimum(start_ms, n_ms) * self.per_ms).astype(np.int64) + end = (np.minimum(end_ms, n_ms) * self.per_ms).astype(np.int64) + return start, end + + def rms(self, start_ms, end_ms) -> np.ndarray: + """``audioop.rms`` of each ``[start_ms, end_ms)`` slice (an int, floored).""" + start, end = self.bounds(np.asarray(start_ms), np.asarray(end_ms)) + n = len(self.samples) + squares = (self.prefix[np.minimum(end, n)] - self.prefix[np.minimum(start, n)]).astype(np.float64) + width = (end - start).astype(np.float64) # past the end pydub pads zeros + with np.errstate(divide="ignore", invalid="ignore"): + rms = np.floor(np.sqrt(squares / width)) + return np.where(width > 0, rms, 0.0) + + def slice(self, start_ms: int, end_ms: int) -> np.ndarray: + start, end = (int(x) for x in self.bounds(np.asarray(start_ms), np.asarray(end_ms))) + data = self.samples[start:end] + if len(data) < end - start: + data = np.concatenate([data, np.zeros(end - start - len(data), dtype=data.dtype)]) + return data + + +def _silent_ranges(seg: _Segment, min_silence_len: int, threshold: float, seek_step: int) -> list[list[int]]: + """``pydub.silence.detect_silence``.""" + seg_len = len(seg) + if seg_len < min_silence_len: + return [] + last = seg_len - min_silence_len + starts = list(range(0, last + 1, seek_step)) + if last % seek_step: + starts.append(last) + starts = np.asarray(starts, dtype=np.int64) + silent = starts[seg.rms(starts, starts + min_silence_len) <= threshold].tolist() + if not silent: + return [] + ranges = [] + prev = range_start = silent[0] + for s in silent[1:]: + if s != prev + seek_step and s > prev + min_silence_len: + ranges.append([range_start, prev + min_silence_len]) + range_start = s + prev = s + ranges.append([range_start, prev + min_silence_len]) + return ranges + + +def _nonsilent_ranges(seg: _Segment, min_silence_len: int, threshold: float, seek_step: int) -> list[list[int]]: + """``pydub.silence.detect_nonsilent``.""" + silent = _silent_ranges(seg, min_silence_len, threshold, seek_step) + seg_len = len(seg) + if not silent: + return [[0, seg_len]] + if silent[0][0] == 0 and silent[0][1] == seg_len: + return [] + ranges, prev_end = [], 0 + for start, end in silent: + ranges.append([prev_end, start]) + prev_end = end + if end != seg_len: + ranges.append([prev_end, seg_len]) + if ranges[0] == [0, 0]: + ranges.pop(0) + return ranges + + +def _split_on_silence(seg: _Segment, min_silence_len: int, threshold: float, keep: int, seek_step: int) -> np.ndarray: + """``pydub.silence.split_on_silence``, its pieces joined back together.""" + ranges = [[s - keep, e + keep] for s, e in _nonsilent_ranges(seg, min_silence_len, threshold, seek_step)] + for a, b in itertools.pairwise(ranges): + if b[0] < a[1]: + a[1] = (a[1] + b[0]) // 2 + b[0] = a[1] + n_ms = len(seg) + pieces = [seg.slice(max(s, 0), min(e, n_ms)) for s, e in ranges] + return np.concatenate(pieces) if pieces else seg.samples[:0] + + +def _leading_silence(seg: _Segment, threshold_db: float, chunk: int = 10) -> int: + """``pydub.silence.detect_leading_silence``.""" + n_ms = len(seg) + starts = np.arange(0, n_ms, chunk, dtype=np.int64) + rms = seg.rms(starts, starts + chunk) + with np.errstate(divide="ignore"): + db = np.where(rms > 0, 20 * np.log10(rms / _FULL_SCALE), -np.inf) + loud = np.flatnonzero(db >= threshold_db) + trim = int(starts[loud[0]]) if len(loud) else len(starts) * chunk + return min(trim, n_ms) + + +def remove_silence( + audio: np.ndarray, sampling_rate: int, mid_sil: int = 300, lead_sil: int = 100, trail_sil: int = 300, +) -> np.ndarray: + """``omnivoice.utils.audio.remove_silence`` for mono ``(1, T)`` float audio.""" + if audio.shape[0] != 1: + raise ValueError(f"expected mono (1, T) audio, got {audio.shape}") + samples = (audio[0] * 32768.0).clip(-32768, 32767).astype(np.int16) + threshold_db = -50.0 + threshold = 10 ** (threshold_db / 20) * _FULL_SCALE + if mid_sil > 0: + samples = _split_on_silence(_Segment(samples, sampling_rate), mid_sil, threshold, mid_sil, 10) + + seg = _Segment(samples, sampling_rate) + start = max(0, _leading_silence(seg, threshold_db) - lead_sil) + samples = seg.slice(start, len(seg)) + seg = _Segment(samples[::-1], sampling_rate) + start = max(0, _leading_silence(seg, threshold_db) - trail_sil) + samples = seg.slice(start, len(seg))[::-1] + return (samples.astype(np.float32) / 32768.0)[np.newaxis, :] diff --git a/mstar/model/omnivoice/config.py b/mstar/model/omnivoice/config.py index 90f28ce9ac..0bef136070 100644 --- a/mstar/model/omnivoice/config.py +++ b/mstar/model/omnivoice/config.py @@ -117,6 +117,21 @@ class OmniVoiceConfig: # Set by _refresh_checkpoint_defaults; None until weights load. checkpoint_dtype: str | None = None + # Replay the backbone's layer stack from CUDA graphs, one per (requests, + # packed-token bucket), with attention through the engine's ragged + # resource. Eager, a step is launch-bound: ~400 kernels for ~2 ms of GPU + # work. ``model_kwargs: {backbone_graphs: false}`` runs the reference's + # own packed attention eagerly instead. + backbone_graphs: bool = True + backbone_capture_batch_sizes: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16]) + # Packed tokens per step (both CFG documents of every request); a step + # above the largest bucket runs eagerly. Past a few thousand tokens the + # step is GPU-bound, eager launches hide behind it, and a bucket's padding + # would only add work. + backbone_token_buckets: list[int] = field( + default_factory=lambda: [256, 512, 768, 1024, 1536, 2048, 2560, 3072, 3584, 4096] + ) + @property def audio_embedding_rows(self) -> int: """Rows in ``audio_embeddings`` — one codebook's vocab per row block.""" diff --git a/mstar/model/omnivoice/omnivoice_model.py b/mstar/model/omnivoice/omnivoice_model.py index 9c8dc9d477..c6d990c508 100644 --- a/mstar/model/omnivoice/omnivoice_model.py +++ b/mstar/model/omnivoice/omnivoice_model.py @@ -30,7 +30,7 @@ CurrentForwardConductorMetadata, StreamingConnectionState, ) -from mstar.engine.resources import NodeResourceSpec +from mstar.engine.resources import NodeResourceSpec, RaggedAttentionConfig, RaggedAttentionSpec from mstar.graph.base import ( GraphEdge, GraphNode, @@ -44,6 +44,7 @@ from mstar.model.omnivoice.components.backbone import ( OmniVoiceBackbone, assert_flashinfer_api, + assert_flashinfer_patched, ) from mstar.model.omnivoice.components.codec import load_codec from mstar.model.omnivoice.components.text import ( @@ -53,6 +54,7 @@ ) from mstar.model.omnivoice.config import OmniVoiceConfig from mstar.model.omnivoice.submodules import ( + BACKBONE_ATTN, UNMASK_LOOP_NAME, OmniVoiceBackboneSubmodule, OmniVoiceCode2WavSubmodule, @@ -88,6 +90,10 @@ def __init__( self.model_path_hf = os.environ.get("MSTAR_OMNIVOICE_MODEL_PATH") or model_path_hf self.cache_dir = cache_dir self.config = OmniVoiceConfig() + if "backbone_graphs" in kwargs: + self.config.backbone_graphs = bool(kwargs["backbone_graphs"]) + if "backbone_token_buckets" in kwargs: + self.config.backbone_token_buckets = sorted(int(n) for n in kwargs["backbone_token_buckets"]) self.skip_weight_loading = skip_weight_loading # Loaded lazily on the data worker so construction never hits the @@ -104,7 +110,8 @@ def __init__( # ------------------------------------------------------------------ def get_node_resources(self) -> list[NodeResourceSpec]: - """No engine-built resources: every node here is stateless. + """The backbone's ragged attention, when its layer stack is captured; + nothing else, since every node here is stateless. There is deliberately no KV cache. The canvas is rewritten each iteration and attention is bidirectional, so the prefix attends *into* @@ -113,7 +120,23 @@ def get_node_resources(self) -> list[NodeResourceSpec]: resource either: the reveal picks cells by confidence rank across the whole canvas rather than drawing one token per position. """ - return [] + if not self.config.backbone_graphs: + return [] + body = self.config.backbone + return [ + RaggedAttentionSpec( + resource_key=BACKBONE_ATTN, + nodes={"backbone"}, + config=RaggedAttentionConfig( + num_qo_heads=body.num_attention_heads, + num_kv_heads=body.num_key_value_heads, + head_dim=body.head_dim, + # the conditional and unconditional documents + max_segments_per_request=2, + dtype=getattr(torch, self.config.load_dtype), + ), + ) + ] def get_graph_walk_graphs(self) -> dict[str, GraphSection]: encode_reference = GraphNode( @@ -143,7 +166,7 @@ def _build_speech_walk(self, clone: bool) -> GraphSection: backbone_inputs = [ "prefix_ids", "prefix_audio_mask", - "target_len", + "initial_canvas", "audio_tokens", "step_index", ] @@ -384,7 +407,12 @@ def process_prompt( out: NameToTensorList = { "prefix_ids": [prefix_ids], "prefix_audio_mask": [prefix_audio_mask], - "target_len": [torch.tensor([target_len], dtype=torch.long)], + # the all-MASK canvas the unmask loop starts from: its shape carries + # the target length, so the backbone sizes the request without a + # host read of a device value + "initial_canvas": [ + torch.full((self.config.num_audio_codebook, target_len), self.config.audio_mask_id, dtype=torch.long) + ], } if ref_audio: out["ref_audio_inputs"] = [ref_waveform] @@ -495,8 +523,8 @@ def get_initial_forward_pass_args( "walk_schedule": schedule, "walk_step": 0, # target_len is NOT here: it is computed on the data worker (it - # needs the tokenizer and the duration estimator) and rides an edge - # to the backbone, rather than being re-derived at the conductor. + # needs the tokenizer and the duration estimator) and rides to the + # backbone as the initial canvas's shape. "num_step": num_step, "guidance_scale": float( model_kwargs.get("guidance_scale", defaults.guidance_scale) @@ -568,7 +596,7 @@ def _speech_gen_inputs( step counter at iteration 0. """ inputs: list[GraphEdge] = [] - for name in ("prefix_ids", "prefix_audio_mask", "target_len"): + for name in ("prefix_ids", "prefix_audio_mask", "initial_canvas"): edge = GraphEdge(next_node="backbone", name=name) edge.tensor_info = persist_signals.get(name, []) inputs.append(edge) @@ -731,6 +759,7 @@ def _create_submodule(self, node_name: str, device: str) -> NodeSubmodule | None # After .to(device): apply_flashinfer sizes its attention workspace # against model.device, and patching on CPU would allocate it there. fi.apply_flashinfer(reference, enable_cuda_graph=False) + assert_flashinfer_patched(reference) backbone = OmniVoiceBackbone(reference, plan_dtype=dtype).eval() return OmniVoiceBackboneSubmodule(backbone, self.config) diff --git a/mstar/model/omnivoice/submodules.py b/mstar/model/omnivoice/submodules.py index 15fec8cc88..a8955ac7f4 100644 --- a/mstar/model/omnivoice/submodules.py +++ b/mstar/model/omnivoice/submodules.py @@ -13,12 +13,22 @@ """ import logging +from collections.abc import Mapping +from functools import partial +from typing import Any import torch from torch import nn from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.accelerator_graph_config import ( + PiecewiseAcceleratorGraphConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, + PiecewisePackedConfig, +) +from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep from mstar.model.omnivoice.components.backbone import ( CanvasItem, OmniVoiceBackbone, @@ -41,6 +51,16 @@ logger = logging.getLogger(__name__) UNMASK_LOOP_NAME = "unmask_loop" +# The backbone's ragged attention resource, and its captured layer stack. +BACKBONE_ATTN = "backbone_attn" +BACKBONE_LAYERS = "backbone_layers" + + +def document_pairs(lengths: list[tuple[int, int]]) -> list[int]: + """Two documents per request, conditional then unconditional, in packing + order; a request with CFG off has a zero-length unconditional one, so every + request owns two segments whatever its guidance.""" + return [n for pair in lengths for n in pair] def _loop_step(fwd_info: CurrentForwardPassInfo, step_index: torch.Tensor) -> int: @@ -206,23 +226,17 @@ def prepare_inputs( ) if "audio_tokens" not in inputs or len(inputs["audio_tokens"]) == 0: - # Iteration 0: the loop-back edges arrive empty. Seed a fully - # masked canvas; there is no noise and no seed to honour, because - # the stochastic part of this model is the reveal order, not the - # starting state. target_len rides an edge rather than - # step_metadata: computing it needs the tokenizer and the duration - # estimator, both data-worker assets. Later iterations read it off - # the canvas shape instead, which costs no host read. - target_len = int(inputs["target_len"][0].reshape(-1)[0].item()) - tokens = torch.full( - (self.config.num_audio_codebook, target_len), - self.config.audio_mask_id, - dtype=torch.long, - device=device, - ) + # Iteration 0: the loop-back edges arrive empty, and the loop starts + # from the all-MASK canvas the data worker built (it knows the + # target length from the tokenizer and the duration estimator). + # There is no noise and no seed to honour: the stochastic part of + # this model is the reveal order, not the starting state. + tokens = inputs["initial_canvas"][0].to(device) + target_len = tokens.shape[-1] # 1-D, not 0-dim: the worker's output fanout reads dims[0] of # every routed edge. step_index = torch.zeros(1, dtype=torch.int64, device=device) + k = 0 else: step_index = inputs["step_index"][0] k = _loop_step(fwd_info, step_index) @@ -239,6 +253,8 @@ def prepare_inputs( tokens = inputs["audio_tokens"][0].to(device) target_len = tokens.shape[-1] + cond_len = prefix_ids.shape[-1] + target_len + uncond_len = target_len if float(meta["guidance_scale"]) != 0 else 0 return NodeInputs( tensor_inputs={ "prefix_ids": prefix_ids, @@ -246,12 +262,109 @@ def prepare_inputs( "audio_tokens": tokens, "step_index": step_index, }, - kwargs={"request_id": fwd_info.request_id}, - # Both CFG documents are packed, so the row's real contribution - # is the conditional length plus the target again. - input_seq_len=prefix_ids.shape[-1] + 2 * target_len, + # the iteration, for postprocess: at iteration 0 the loop counter + # is not there yet, and reading the step edge would sync + kwargs={"request_id": fwd_info.request_id, "doc_lens": (cond_len, uncond_len), "step": k}, + # The row's packed contribution: its conditional document plus, + # with CFG on, the target again as the unconditional one. + input_seq_len=cond_len + uncond_len, + ) + + def declare_step( + self, + graph_walk: str, + request_ids: list[str], + inputs: list[NodeInputs], + slot_lease: SlotLease | None = None, + piecewise_leases: Mapping[str, SlotLease] | None = None, + **kwargs, + ) -> SubmoduleStep | None: + """The eager path's attention layout; a leased replay plans its own.""" + if not self.config.backbone_graphs or (piecewise_leases or {}).get(BACKBONE_LAYERS): + return None + return self._attention_step( + request_ids, document_pairs([inp.kwargs["doc_lens"] for inp in inputs]) + ) + + @staticmethod + def _attention_step(request_ids: list[str], doc_lens: list[int]) -> SubmoduleStep: + """Each document attends within itself. ``doc_lens`` holds two entries + per real request; a padded replay appends one zero per padding row, and + a capture one length per row, so documents are spread over the ids in + order -- the ragged plan reads only the spans.""" + per_id = max(1, -(-len(doc_lens) // len(request_ids))) + return SubmoduleStep( + segments=[ + Segment(request_id=request_ids[i // per_id], label="main", span=n) + for i, n in enumerate(doc_lens) + ], + steps={BACKBONE_ATTN: AttentionStep(causal=False)}, ) + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, **kwargs: Any + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + """The layer stack (embeddings through the final norm), by request + count and packed-token bucket. The packed canvas is built eagerly before + it and the audio head runs after it, on the gathered target rows.""" + del autocast_dtype, tp_world_size, kwargs + if not self.config.backbone_graphs: + return {} + return { + BACKBONE_LAYERS: PiecewisePackedConfig( + capture_fn=self._capture_layers, + make_static_inputs=partial(self._layer_static_inputs, device=device), + declare_step=self._attention_step, + lease_before_step=True, + total_tokens=list(self.config.backbone_token_buckets), + capture_batch_sizes=list(self.config.backbone_capture_batch_sizes), + ) + } + + def _layer_static_inputs(self, shape: PiecewiseCaptureShape, device) -> dict[str, torch.Tensor]: + n = shape.total_tokens + return { + "input_ids": torch.zeros(n, self.config.num_audio_codebook, dtype=torch.long, device=device), + "audio_mask": torch.zeros(n, dtype=torch.bool, device=device), + "pos_ids": torch.zeros(n, dtype=torch.int32, device=device), + } + + def _capture_layers(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + s = inp.static_inputs + return { + "hidden": self.backbone.encode( + s["input_ids"], s["audio_mask"], s["pos_ids"], + self.node_resources[BACKBONE_ATTN], + ) + } + + @torch.inference_mode() + def _logits(self, engine_inputs: ModelInputsFromEngine, canvas: PackedCanvas) -> torch.Tensor: + """The step's target logits: from the captured layer stack when a + bucket fits, else eagerly through the same ragged attention.""" + if not self.config.backbone_graphs: + return self.backbone(canvas) + position_ids = canvas.position_ids[0] + static_inputs = { + "input_ids": canvas.packed_ids[0].t(), + "audio_mask": canvas.audio_mask[0], + "pos_ids": position_ids.to(torch.int32), + } + num_requests = len(canvas.items) + runner = engine_inputs.piecewise_runners.get(BACKBONE_LAYERS) + if runner is not None and runner.can_run(num_requests, int(position_ids.shape[0])): + hidden = runner.run( + static_inputs=static_inputs, + request_ids=list(engine_inputs.request_ids), + seq_lens=document_pairs([ + (item.cond_len, item.target_len if item.does_cfg else 0) for item in canvas.items + ]), + real_bs=num_requests, + ).get_view("hidden") + else: + hidden = self.backbone.encode(**static_inputs, attn=engine_inputs.resources[BACKBONE_ATTN]) + return self.backbone.head(hidden, canvas.tgt_index) + def preprocess( self, graph_walk: str, @@ -335,7 +448,7 @@ def forward_batched( assert items is not None and canvas is not None, ( "OmniVoice backbone requires preprocess output" ) - logits = self.backbone(canvas) + logits = self._logits(engine_inputs, canvas) outputs: dict[str, NameToTensorList] = {} for item in items: c_logits, u_logits = canvas.slice_logits(logits, item) @@ -363,12 +476,12 @@ def postprocess( Scoring, the reveal schedule and the in-place write are per-request and data-dependent, which is what keeps them out of ``forward_batched``. - The iteration index is the engine's loop counter (see ``_loop_step``); - everything else comes from step metadata. + The iteration index is the one ``prepare_inputs`` took from the loop + counter (see ``_loop_step``); everything else comes from step metadata. """ assert inputs is not None, "OmniVoice backbone postprocess needs its inputs" meta = request_info.step_metadata - k = _loop_step(request_info, inputs.tensor_inputs["step_index"]) + k = inputs.kwargs["step"] # Cloned because apply_reveal writes in place and this tensor is the # edge the engine routed in from the previous iteration; mutating it diff --git a/mstar/model/orpheus/submodules.py b/mstar/model/orpheus/submodules.py index cb4e4d24d9..3194142765 100644 --- a/mstar/model/orpheus/submodules.py +++ b/mstar/model/orpheus/submodules.py @@ -3,6 +3,7 @@ from typing import Any import torch +import torch.nn.functional as F from torch import nn from mstar.communication.tensors import NameToTensorList @@ -22,6 +23,7 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, @@ -51,6 +53,17 @@ def __init__( self.embed_tokens = language_model.model.embed_tokens self.lm_head = language_model.lm_head self.config = config + # Orpheus only ever emits ```` ids (SNAC codes and the + # speech/turn markers): the LM head computes the logits of rows + # ``[custom_token_base_id, vocab)`` alone, about 18% of the head, and + # pads the Llama text rows with -inf, so the sampler sees full-width + # logits and every token id stays in full-vocab space. Not under TP: + # the vocab-sharded head puts the custom rows on the last rank. + self._audio_vocab_start: int | None = None + if self.lm_head.tp_size == 1: + if config.stop_token_id < config.custom_token_base_id: + raise ValueError("Orpheus's stop token must be a custom token") + self._audio_vocab_start = config.custom_token_base_id PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] PREFILL_CAPTURE_BATCH_SIZES = [1, 4, 16, 64] @@ -147,7 +160,11 @@ def _forward( if graph_walk != "decode": hidden = attn.select_last_hidden(hidden) - logits = self.lm_head(hidden) + if self._audio_vocab_start is not None: + start = self._audio_vocab_start + logits = F.pad(F.linear(hidden, self.lm_head.weight[start:]), (start, 0), value=float("-inf")) + else: + logits = self.lm_head(hidden) return sampler.sample( engine_inputs.request_ids, logits=logits @@ -188,7 +205,10 @@ def forward_batched( text_inputs=text_inputs ) # row i is request i; a per-rid dict here would guard the compiled frame on the ids - return BatchedModelOutput(row_outputs={"new_token": new_tokens}) + return BatchedModelOutput( + row_outputs={"new_token": new_tokens}, + check_stop_buffers={"new_token": new_tokens}, + ) def postprocess( self, request_id: str, @@ -218,6 +238,33 @@ def check_stop( return set() + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for every row off one host copy of the step's tokens.""" + tokens = host_rows.buffers.get("new_token") + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(-1).tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + eos = self.config.stop_token_id + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + if (values[i] == eos and not info.resource_configs[SAMPLER].ignore_eos) or ( + info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens + ): + stops[rid] = {"decode_loop"} + return stops + + class SNACDecoderSubmodule(NodeSubmodule): """SNAC 24kHz streaming decoder submodule. diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index 0bc17fb1aa..f25913ac61 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -43,6 +43,7 @@ ProcessPromptOutput, TensorAndMetadata, ) +from mstar.model.media_limits import read_image from mstar.model.multimodal import ( TEXT, PromptPart, @@ -149,6 +150,15 @@ class PrefillStep: # one list per name: a `prefill_vision` step carries every span and image input_tensors: dict[str, list[TensorPointerInfo]] order: tuple[int, ...] = () + # each picked image's or video's (t, h, w) patch grid, as `image_grid_thw` + # holds it on the device; the nodes read these instead of syncing on it + image_grids: tuple[tuple[int, int, int], ...] = () + + +# model_kwargs keys: every image's / video's (t, h, w) patch grid, in attachment order +IMAGE_GRIDS_KEY, VIDEO_GRIDS_KEY = "image_grids", "video_grids" +# the step_metadata key the prefill_vision nodes read the picked grids from +STEP_GRIDS_KEY = "image_grids" # TODO: implement MoE variants @@ -159,6 +169,7 @@ class Qwen3_5DenseModel(Model): # 35-130 ms); at 4, preprocessing held at 7-10 ms and p95 at 49 ms. One # thread is steadier but slower (13 ms). PREPROCESS_TORCH_THREADS = 4 + INTERNAL_KWARGS = (IMAGE_GRIDS_KEY, VIDEO_GRIDS_KEY) def __init__( self, @@ -413,6 +424,36 @@ def video_processor(self): ) return self._video_processor + def pixel_normalizer(self) -> tuple[torch.Tensor, torch.Tensor]: + """Per-column mean and std that turn `process_prompt`'s byte patches + into the processors' normalized pixels: their fused rescale-and- + normalize, ``(x - mean / rescale) / (std / rescale)``, element for + element. One for images and videos, whose processors share it.""" + processor = self.image_processor + video = self.video_processor + for name in ("rescale_factor", "image_mean", "image_std", "temporal_patch_size", "patch_size"): + if getattr(video, name) != getattr(processor, name): + raise ValueError(f"image and video processors disagree on {name}") + scale = 1.0 / processor.rescale_factor + mean = torch.tensor(processor.image_mean) * scale + std = torch.tensor(processor.image_std) * scale + # a patch's columns are channel-major (C x T x P x P) + per_channel = processor.temporal_patch_size * processor.patch_size ** 2 + return mean.repeat_interleave(per_channel), std.repeat_interleave(per_channel) + + def load_image(self, filepath: str, device: str) -> TensorAndMetadata: + """The decoded image, left uint8 ``(C, H, W)``. + + The base class hands out floats in [0, 1], which `_as_hwc_uint8` only + turns back into bytes for the processor: two full-image passes that + cost more than the processor itself on a large image. + """ + try: + img = read_image(filepath, rgb=True) + except (RuntimeError, OSError) as exc: # torchvision's decoders raise RuntimeError, Pillow's OSError + raise ValueError("could not decode image: the file is corrupt or truncated") from exc + return TensorAndMetadata(img.to(device)) + def load_video(self, filepath: str, device: str) -> TensorAndMetadata: """Only the frames the video processor samples, on the host whatever ``device`` is (see `sample_video`).""" @@ -440,6 +481,9 @@ def _preprocess_video( videos=[frames.cpu()], video_metadata=[VideoMetadata(**metadata)], do_sample_frames=False, + # bytes, normalized on the device like images' (`pixel_normalizer`) + do_rescale=False, + do_normalize=False, return_tensors="pt", ) (grid,) = out["video_grid_thw"] @@ -531,8 +575,15 @@ def process_prompt( segments = split_around_spans(input_ids, spans) check_plan(prefill_plan(expand_video_parts(parts, groups)), spans, len(segments)) + # Resized and patchified here, still bytes: the vision encoder + # rescales and normalizes on the device (`pixel_normalizer`). That is + # a quarter of the bytes to ship, and the float work was most of this + # call (43 of 52 ms for a 2560 px image). images = [ - self.image_processor(images=[_as_hwc_uint8(image)], return_tensors="pt") + self.image_processor( + images=[_as_hwc_uint8(image)], return_tensors="pt", + do_rescale=False, do_normalize=False, + ) for image in raw_images ] videos = [ @@ -544,13 +595,24 @@ def process_prompt( if images: result["pixel_values"] = [out["pixel_values"] for out in images] result["image_grid_thw"] = [g for out in images for g in out["image_grid_thw"]] + # the raw images are not a graph input (the encoder reads + # `pixel_values`), so do not ship them to the workers + result["image_inputs"] = [] if videos: result["pixel_values_videos"] = [video.pixel_values for video in videos] result["video_grid_thw"] = [video.grid for video in videos] # the decoded frames are spent; don't hold them for the request's life result["video_inputs"] = [] - # always set, so a `video_groups` the client sent never reaches the schedule - return ProcessPromptOutput(result, {"video_groups": groups}) + def host(grids) -> list[tuple[int, ...]]: + return [tuple(int(v) for v in g.reshape(-1).tolist()) for g in grids] + + # always set, so a `video_groups` the client sent never reaches the + # schedule; the grids ride to the nodes as step metadata + return ProcessPromptOutput(result, { + "video_groups": groups, + IMAGE_GRIDS_KEY: host(result.get("image_grid_thw", [])), + VIDEO_GRIDS_KEY: host(result.get("video_grid_thw", [])), + }) def postprocess( @@ -584,6 +646,7 @@ def _prefill_schedule( input_modalities: list[str], signals: dict[str, list[TensorPointerInfo]], video_groups: list[int] | None = None, + host_grids: dict[str, list] | None = None, ) -> list[PrefillStep]: """The prefill walks a request runs. @@ -604,6 +667,12 @@ def _prefill_schedule( "image": signals.get("image_grid_thw", []), VIDEO: signals.get("video_grid_thw", []), } + host_grids = host_grids or {} + host_pools = { + "image": host_grids.get(IMAGE_GRIDS_KEY) or [], + VIDEO: host_grids.get(VIDEO_GRIDS_KEY) or [], + } + picked_grids: list[tuple[int, int, int]] = [] owners = group_owners(video_groups) if len(video_groups) != input_modalities.count(VIDEO): logger.warning( @@ -649,6 +718,8 @@ def _prefill_schedule( if first_group: picked["pixel_values"].append(pool[index]) picked["image_grid_thw"].append(grid_pools[step.modality][index]) + if index < len(host_pools[step.modality]): + picked_grids.append(tuple(host_pools[step.modality][index])) if not order: return [] @@ -661,7 +732,10 @@ def _prefill_schedule( # TODO: one walk holds every image and video, so a many-image prompt # makes one very wide step. Cap images per walk here and emit several # steps, each with its share of `order`. - return [PrefillStep("prefill_vision", picked, tuple(order))] + # all or nothing: the nodes fall back to the device grids without them + if len(picked_grids) != len(picked["image_grid_thw"]): + picked_grids = [] + return [PrefillStep("prefill_vision", picked, tuple(order), tuple(picked_grids))] def get_initial_forward_pass_args( self, @@ -671,9 +745,11 @@ def get_initial_forward_pass_args( input_signals: dict[str, list[TensorPointerInfo]], model_kwargs: dict | None = None, ) -> ForwardPassArgs: + model_kwargs = model_kwargs or {} schedule = self._prefill_schedule( input_modalities, input_signals, - (model_kwargs or {}).get("video_groups"), + model_kwargs.get("video_groups"), + {key: model_kwargs.get(key) for key in (IMAGE_GRIDS_KEY, VIDEO_GRIDS_KEY)}, ) if not schedule: # every step was skipped for want of a tensor @@ -694,15 +770,18 @@ def get_initial_forward_pass_args( metadata.kwargs["prefill_schedule"] = schedule[1:] inputs = self._walk_inputs(step) + step_metadata = { + "is_prefill": True, + "last_prefill": len(schedule) == 1, + "prefill_order": step.order, + } + if step.image_grids: + step_metadata[STEP_GRIDS_KEY] = step.image_grids return ForwardPassArgs( full_metadata=metadata, inputs=inputs, unpersist_tensors=sum((inp.tensor_info for inp in inputs), start=[]), - step_metadata={ - "is_prefill": True, - "last_prefill": len(schedule) == 1, - "prefill_order": step.order, - }, + step_metadata=step_metadata, ) @classmethod @@ -735,6 +814,7 @@ def get_partition_forward_pass_args( last_prefill = len(remaining) == 1 order: tuple[int, ...] = () + image_grids: tuple[tuple[int, int, int], ...] = () if metadata.is_prefill and remaining: # each prefill step carries its own tensors; only decode reads # back what the previous step emitted @@ -742,6 +822,7 @@ def get_partition_forward_pass_args( metadata.graph_walk = step.walk inputs = self._walk_inputs(step) order = step.order + image_grids = step.image_grids elif metadata.is_prefill: metadata.is_prefill = False metadata.graph_walk = "decode" @@ -759,15 +840,18 @@ def get_partition_forward_pass_args( request_done=True, ) + step_metadata = { + "is_prefill": metadata.is_prefill, + "last_prefill": last_prefill, + "prefill_order": order, + } + if image_grids: + step_metadata[STEP_GRIDS_KEY] = image_grids return ForwardPassArgs( full_metadata=metadata, inputs=inputs, unpersist_tensors=sum((inp.tensor_info for inp in inputs), start=[]), - step_metadata={ - "is_prefill": metadata.is_prefill, - "last_prefill": last_prefill, - "prefill_order": order, - }, + step_metadata=step_metadata, ) # ------------------------------------------------------------------ @@ -842,7 +926,7 @@ def _create_submodule( load_qwen3_5_vision_weights(tower, weights_dir, device=device) tower.requires_grad_(False).eval() logger.info("Loaded Qwen3.5 vision tower onto %s", device) - return VisionEncoderSubmodule(tower, self.vision_config) + return VisionEncoderSubmodule(tower, self.vision_config, self.pixel_normalizer()) from mstar.model.qwen3_5.weight_loader import load_qwen3_5_weights diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 44950b8e87..0792181286 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -64,7 +64,7 @@ Qwen3_5Config, Qwen3_5VisionConfig, ) -from mstar.model.qwen3_5.qwen3_5_model import TEXT_PART, VIDEO_PART +from mstar.model.qwen3_5.qwen3_5_model import STEP_GRIDS_KEY, TEXT_PART, VIDEO_PART from mstar.model.submodule_base import ( ARNodeInputs, ARNodeSubmodule, @@ -75,12 +75,36 @@ NodeSubmodule, device_loopback_enabled, ) -from mstar.utils.knobs import ingraph_decode_rope +from mstar.platforms.compile import with_inductor_overrides +from mstar.utils.device import DeviceCapability from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) +def _host_grids( + fwd_info: CurrentForwardPassInfo, device_grids: list[torch.Tensor], +) -> list[tuple[int, int, int]]: + """Each image's or video's (t, h, w), off the step metadata the conductor + filled from `process_prompt`; reading `image_grid_thw` back from the + device instead blocks the host on every kernel queued ahead of it.""" + host = fwd_info.step_metadata.get(STEP_GRIDS_KEY) + if host: + return [(int(t), int(h), int(w)) for t, h, w in host] + return [ + (int(t), int(h), int(w)) + for g in device_grids for t, h, w in g.reshape(-1, 3).tolist() + ] + + +def _to_device(t: torch.Tensor, device: torch.device) -> torch.Tensor: + """A host-built tensor onto the device without blocking the host: a copy + from pageable memory waits for the stream to drain.""" + if device.type == "cpu": + return t + return t.pin_memory().to(device, non_blocking=True) + + class LLMSubmodule(ARNodeSubmodule): # the decode walk reads its input ids off the sampler's slot masters reads_device_loopback = True @@ -244,7 +268,7 @@ def _vision_inputs( device = self.get_device() merge = self.vision_config.spatial_merge_size texts = inputs.get("text_inputs", []) - grids = inputs.get("image_grid_thw", []) + grids = _host_grids(fwd_info, inputs.get("image_grid_thw", [])) packed = inputs["vision_embeds"][0].to(device) start_embed, end_embed = self._sentinel_embeds() @@ -267,13 +291,11 @@ def _vision_inputs( text_i += 1 span = ids.shape[0] embeds.append(self.model.model.embed_tokens(ids)) - pos_ids.append(text_position_ids(span, pos, device)) + pos_ids.append(text_position_ids(span, pos)) tracked.append(ids) pos += span continue - grid = grids[grid_i] - grid = grid[0] if grid.dim() == 2 else grid - t, h, w = (int(v) for v in grid.tolist()) + t, h, w = grids[grid_i] if kind == VIDEO_PART: group += 1 if group == t: @@ -285,10 +307,11 @@ def _vision_inputs( span = t * (h // merge) * (w // merge) advance = vision_position_advance(grid, merge) embeds += [start_embed, packed[cursor:cursor + span], end_embed] + # built on the host, copied once below pos_ids += [ - text_position_ids(1, pos, device), - vision_position_ids(grid, merge, pos + 1, device), - text_position_ids(1, pos + 1 + advance, device), + text_position_ids(1, pos), + vision_position_ids(grid, merge, pos + 1), + text_position_ids(1, pos + 1 + advance), ] cursor += span # two sentinels either side, and the image between them @@ -298,7 +321,7 @@ def _vision_inputs( return ARNodeInputs( input_seq_len=input_embeds.shape[0], input_embeds=input_embeds, - custom_pos_ids=torch.cat(pos_ids, dim=1), + custom_pos_ids=_to_device(torch.cat(pos_ids, dim=1), device), tensor_inputs={ "mrope_advance": pos - start_pos, # Ids cannot ride `input_ids` — that is what `preprocess` keys @@ -437,14 +460,22 @@ def preprocess( for i, inp in enumerate(inputs) ]) elif routed: + # Through pinned memory, non-blocking: a pageable + # ``torch.tensor(..., device=cuda)`` is a synchronous copy + # that waits for the queued step, and every request's first + # decode step after its prefill takes this path. + index = torch.tensor( + routed, dtype=torch.long, + pin_memory=DeviceCapability.of(ids.device).pin_memory, + ).to(ids.device, non_blocking=True) ids.index_copy_( - 0, torch.tensor(routed, device=ids.device), + 0, index, torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), ) out["input_ids"] = ids if PHASE_PERIOD: phase_record("engine.preprocess.loopback", perf_counter() - t0) - if captured and ingraph_decode_rope(): + if captured: # the captured forward builds the rotary tables itself return out elif inputs[0].input_ids is not None: @@ -546,9 +577,10 @@ def _forward( captured = self._captured_decode(engine_inputs, graph_walk) if input_ids is None and input_embeds is None: input_ids = sampler.loopback_tokens(engine_inputs.request_ids) - if (captured and ingraph_decode_rope()) or cos_3d is None or sin_3d is None: - # MSTAR_INGRAPH_DECODE_ROPE: the tables come from the planned - # positions, inside the graph; the static cos/sin stay unread + if captured or cos_3d is None or sin_3d is None: + # A captured decode builds its tables from the planned positions, + # inside the graph, off the per-step critical path; the static + # cos/sin stay unread num_tokens = (input_ids if input_ids is not None else input_embeds).shape[0] cos_3d, sin_3d = self._decode_cos_sin(engine_inputs, num_tokens) @@ -726,30 +758,46 @@ def check_stop_batched( class VisionEncoderSubmodule(NodeSubmodule): """Qwen3.5's ViT, which turns pixel patches into LLM-width embeddings. - A prompt's images go through packed, in one call. Grid maths runs in + A step's images go through packed, in one call. Grid maths runs in `prepare_inputs` and attention boundaries are declared as step segments, so the compiled tower's only symbolic shape is its token count. - The block loop is captured one prompt per replay, which beats batching: - eager, the tower is launch-bound (~500 kernels, ~4 ms of GPU work in a - ~20 ms forward on 9B). So the encoder takes one request a step. + Eager, the tower is launch-bound (~500 kernels, ~4 ms of GPU work in a + ~20 ms forward on 9B), so the block loop is captured, and a step's + requests replay together as one packed row: every image and video frame + group of every request is its own attention segment. A replay's attention plan holds at most the vision attention resource's - ``max_segments_per_request`` segments (one per image and per video frame - group). A prompt with more takes no lease and runs the block loop - eagerly, however few its patches. + ``max_segments_per_request`` segments. A step with more takes no lease + and runs the block loop eagerly, however few its patches. """ - # Patch-count buckets, one prompt per replay. 16384 patches merge to the - # LLM's 4096 `prefill_vision` bucket; a larger prompt runs eagerly. + # Patch-count buckets of one replay. 16384 patches merge to the LLM's + # 4096 `prefill_vision` bucket; a larger step runs eagerly. BLOCK_LOOP_PATCH_BUCKETS = [256, 512, 1024, 2048, 4096, 8192, 16384] + # Requests per encoder step. They share one replay, so a burst of image + # requests pays the per-step host cost once per this many, not per image. + MAX_BATCH_SIZE = 8 + + # The engine would compile `forward_batched` whole, and its per-step + # bookkeeping (each request's slice of the output, the segment lengths) + # guards on the batch size: every new size recompiled (~0.6 s, on a live + # request). The tensor stages around the captured block loop are compiled + # here instead, dynamic in the patch count from the first call. A step + # past the top capture bucket runs the blocks eagerly: compiling that rare + # path stalled its first request for seconds. + disable_torch_compile = True def __init__( self, model: Qwen3_5VisionModel, config: Qwen3_5VisionConfig, + pixel_normalizer: tuple[torch.Tensor, torch.Tensor] | None = None, ): super().__init__() self.model = model self.config = config + self._compiled: tuple | None = None + # per-column (mean, std) for byte patches; see `Qwen3_5DenseModel.pixel_normalizer` + self._pixel_normalizer = pixel_normalizer def prepare_inputs( self, @@ -758,18 +806,15 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs: Any, ) -> NodeInputs: - # (num_images, 3): a bare [t, h, w] would read as three images device = self.get_device() merge = self.config.spatial_merge_size - grid = [ - (int(t), int(h), int(w)) - for g in inputs["image_grid_thw"] - for t, h, w in g.reshape(-1, 3).tolist() - ] + grid = _host_grids(fwd_info, inputs["image_grid_thw"]) # Built here, not in the compiled forward: a `.tolist()` there breaks - # the graph and makes h, w symints inductor cannot codegen from. + # the graph and makes h, w symints inductor cannot codegen from. And + # built on the host: on the device, the grid upload and the + # `repeat_interleave`s each block on the stream. indices, weights = vision_interpolation( - grid, self.config.num_grid_per_side, merge, device, + grid, self.config.num_grid_per_side, merge, torch.device("cpu"), ) num_patches = sum(t * h * w for t, h, w in grid) return NodeInputs( @@ -777,9 +822,11 @@ def prepare_inputs( input_seq_len=num_patches, tensor_inputs={ "pixel_values": torch.cat(inputs["pixel_values"], dim=0), - "indices": indices, - "weights": weights, - "position_ids": vision_grid_position_ids(grid, merge, device), + "indices": _to_device(indices, device), + "weights": _to_device(weights, device), + "position_ids": _to_device( + vision_grid_position_ids(grid, merge, torch.device("cpu")), device, + ), # not forward args: `declare_step` makes segments of the # lengths, `forward_batched` cuts the output at the count "seq_lengths": vision_seq_lengths(grid), @@ -814,8 +861,7 @@ def can_batch( return True def max_batch_size(self, graph_walk: str) -> int | None: - # one request a step, so it can take the captured block loop - return 1 + return self.MAX_BATCH_SIZE def get_piecewise_accelerator_graph_configs( self, device: torch.device, autocast_dtype: torch.dtype, @@ -850,6 +896,7 @@ def declare_step(request_ids: list[str], seq_lens: list[int]) -> SubmoduleStep: make_static_inputs=make_static_inputs, declare_step=declare_step, lease_before_step=True, + one_row_per_step=True, can_lease=self._fits_capture, total_tokens=self.BLOCK_LOOP_PATCH_BUCKETS, capture_batch_sizes=[1], @@ -869,14 +916,14 @@ def _fits_capture(self, inputs: list[NodeInputs]) -> bool: """ per_request = self.node_resources[VISION_ATTN].max_segments_per_request segments = sum(len(inp.tensor_inputs["seq_lengths"]) for inp in inputs) - return segments <= per_request * len(inputs) + # the whole step replays as one row (`one_row_per_step`) + return segments <= per_request def _capture_block_loop(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: """The captured region: reads the runner's buffers, never rebinds.""" s = inp.static_inputs return {"hidden": self.model.encode(s["hidden"], s["cos"], s["sin"])} - @torch.compiler.disable def _replay_block_loop( self, engine_inputs: ModelInputsFromEngine, @@ -893,8 +940,9 @@ def _replay_block_loop( forward driven outside the engine (no declaration either) eager. """ runner = engine_inputs.piecewise_runners.get(VISION_BLOCK_LOOP) + # the whole step is one row (`one_row_per_step`) if runner is None or engine_inputs.step is not None or not runner.can_run( - len(engine_inputs.request_ids), int(hidden.shape[0]), + 1, int(hidden.shape[0]), ): return None # one replay for every image, not one each: a second plan overwrites the schedule a queued replay still reads @@ -911,19 +959,41 @@ def preprocess( engine_inputs: ModelInputsFromEngine, inputs: list[NodeInputs], ) -> dict[str, Any]: - """The batch's one request (see `max_batch_size`), its images packed - in `declare_step`'s segment order.""" - (inp,) = inputs + """Every request's images and videos packed end to end, in + `declare_step`'s segment order.""" device = self.get_device() - tensors = inp.tensor_inputs + + def packed(name: str) -> torch.Tensor: + parts = [inp.tensor_inputs[name] for inp in inputs] + return (parts[0] if len(parts) == 1 else torch.cat(parts)).to(device) + return { - "pixel_values": tensors["pixel_values"].to(device), - "indices": tensors["indices"].to(device), - "weights": tensors["weights"].to(device), - "position_ids": tensors["position_ids"].to(device), - "seq_lengths": tuple(tensors["seq_lengths"]), + "pixel_values": self._normalize(packed("pixel_values")), + "indices": packed("indices"), + "weights": packed("weights"), + "position_ids": packed("position_ids"), + "seq_lengths": tuple( + n for inp in inputs for n in inp.tensor_inputs["seq_lengths"] + ), + # merged tokens per request, to cut the output back apart + "row_tokens": tuple( + inp.input_seq_len // self.config.spatial_merge_size ** 2 + for inp in inputs + ), } + def _normalize(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Byte patches to the processors' float pixels. Eager, not in the + compiled embed, so the arithmetic (a subtract, a true divide) is the + processors' exactly.""" + if pixel_values.dtype != torch.uint8: + return pixel_values + mean, std = self._pixel_normalizer + if mean.device != pixel_values.device: + mean, std = mean.to(pixel_values.device), std.to(pixel_values.device) + self._pixel_normalizer = (mean, std) + return (pixel_values.to(torch.float32) - mean) / std + def forward( self, graph_walk: str, @@ -942,13 +1012,34 @@ def forward_batched( weights: torch.Tensor, position_ids: torch.Tensor, seq_lengths: tuple[int, ...] = (), + row_tokens: tuple[int, ...] = (), **kwargs, ) -> dict[str, NameToTensorList]: # One call for every image; the ragged kernel's per-frame segments # keep them apart. - hidden, cos, sin = self.model.embed(pixel_values, indices, weights, position_ids) + if self._compiled is None: + # under the platform's inductor settings, as the engine's compiles are + self._compiled = ( + with_inductor_overrides(torch.compile(self.model.embed, dynamic=True)), + with_inductor_overrides(torch.compile(self.model.merger, dynamic=True)), + ) + embed, merger = self._compiled + hidden, cos, sin = embed(pixel_values, indices, weights, position_ids) encoded = self._replay_block_loop(engine_inputs, hidden, cos, sin, seq_lengths) if encoded is None: encoded = self.model.encode(hidden, cos, sin) - embeds = self.model.merger(encoded) - return {engine_inputs.request_ids[0]: {"vision_embeds": [embeds]}} + embeds = merger(encoded) + return self._split_rows(engine_inputs.request_ids, embeds, row_tokens) + + @staticmethod + def _split_rows( + request_ids: list[str], embeds: torch.Tensor, row_tokens: tuple[int, ...], + ) -> dict[str, NameToTensorList]: + """Each request's slice of the packed output.""" + if len(request_ids) == 1: + return {request_ids[0]: {"vision_embeds": [embeds]}} + out, start = {}, 0 + for rid, n in zip(request_ids, row_tokens, strict=True): + out[rid] = {"vision_embeds": [embeds[start:start + n]]} + start += n + return out diff --git a/mstar/model/qwen3_omni/components/cp_kernels.py b/mstar/model/qwen3_omni/components/cp_kernels.py new file mode 100644 index 0000000000..dbbf661578 --- /dev/null +++ b/mstar/model/qwen3_omni/components/cp_kernels.py @@ -0,0 +1,161 @@ +"""Code-predictor attention prologue in one launch. + +Per layer and depth step the code predictor normalized and roped q and k with one +kernel each (plus an int32 cast of the positions before each), then copied k and v +into its dense KV cache with two strided copies: six launches of a few microseconds +around a ~1 us job, 16 x 5 times per Talker step. This does all of it in one, with +the same per-element arithmetic as ``mstar.utils.attention._qk_norm_rope_kernel``. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _cp_qk_norm_rope_kv_kernel( + QKV, Wq, Wk, Pos, KV, + eps_q, eps_k, rope_theta, seq_len, cache_pos, + s_qkv, + s_kv_b, s_kv_c, s_kv_n, s_kv_h, + H_Q: tl.constexpr, H_KV: tl.constexpr, D: tl.constexpr, HALF_D: tl.constexpr, +): + m = tl.program_id(0) + h = tl.program_id(1) + offs = tl.arange(0, HALF_D) + row = QKV + m * s_qkv + # programs [0, H_Q) take the q heads, [H_Q, H_Q + H_KV) the k heads and their v + base = row + h * D + if h < H_Q: + w = Wq + eps = eps_q + else: + w = Wk + eps = eps_k + + x_first = tl.load(base + offs).to(tl.float32) + x_second = tl.load(base + offs + HALF_D).to(tl.float32) + sumsq = tl.sum(x_first * x_first, axis=0) + tl.sum(x_second * x_second, axis=0) + inv_rms = 1.0 / tl.sqrt(sumsq / D + eps) + x_first = x_first * inv_rms * tl.load(w + offs).to(tl.float32) + x_second = x_second * inv_rms * tl.load(w + offs + HALF_D).to(tl.float32) + + pos = tl.load(Pos + m).to(tl.float32) + inv_freq = tl.exp(-tl.log(rope_theta) * (offs.to(tl.float32) * 2.0 / D)) + angle = pos * inv_freq + cos = tl.cos(angle) + sin = tl.sin(angle) + out_first = x_first * cos - x_second * sin + out_second = x_second * cos + x_first * sin + + if h < H_Q: + tl.store(base + offs, out_first.to(QKV.dtype.element_ty)) + tl.store(base + offs + HALF_D, out_second.to(QKV.dtype.element_ty)) + else: + hk = h - H_Q + b = m // seq_len + n = cache_pos + m % seq_len + k_dst = KV + b * s_kv_b + n * s_kv_n + hk * s_kv_h + tl.store(k_dst + offs, out_first.to(KV.dtype.element_ty)) + tl.store(k_dst + offs + HALF_D, out_second.to(KV.dtype.element_ty)) + d = tl.arange(0, D) + v = tl.load(row + (H_Q + H_KV + hk) * D + d) + tl.store(k_dst + s_kv_c + d, v.to(KV.dtype.element_ty)) + + +def cp_qk_norm_rope_kv_write( + qkv: torch.Tensor, q_norm_w: torch.Tensor, k_norm_w: torch.Tensor, pos: torch.Tensor, + kv_layer: torch.Tensor, cache_pos: int, seq_len: int, eps_q: float, eps_k: float, + rope_theta: float, n_q_heads: int, n_kv_heads: int, head_dim: int, +) -> None: + """QK-norm + RoPE the q heads of ``qkv (tokens, (Hq + 2 Hkv) * D)`` in place, and + write the normed, roped k and the raw v of token ``b * seq_len + s`` into + ``kv_layer (bs, 2, S, Hkv, D)`` at sequence position ``cache_pos + s``. ``pos`` holds + one position per token.""" + tokens = qkv.shape[0] + assert qkv.stride(1) == 1 and kv_layer.stride(-1) == 1 + _cp_qk_norm_rope_kv_kernel[(tokens, n_q_heads + n_kv_heads)]( + qkv, q_norm_w, k_norm_w, pos, kv_layer, + eps_q, eps_k, rope_theta, seq_len, cache_pos, + qkv.stride(0), + kv_layer.stride(0), kv_layer.stride(1), kv_layer.stride(2), kv_layer.stride(3), + H_Q=n_q_heads, H_KV=n_kv_heads, D=head_dim, HALF_D=head_dim // 2, + num_warps=4, num_stages=1, + ) + + +@triton.jit +def _cp_embed_accumulate_kernel( + Tok, W, Emb, Sum, Codes, group, s_w, s_emb, s_sum, s_codes, + H: tl.constexpr, BLOCK: tl.constexpr, +): + b = tl.program_id(0) + tok = tl.load(Tok + b) + offs = tl.arange(0, BLOCK) + for h0 in tl.static_range(0, H, BLOCK): + e = tl.load(W + tok * s_w + h0 + offs) + tl.store(Emb + b * s_emb + h0 + offs, e) + acc = tl.load(Sum + b * s_sum + h0 + offs) + tl.store(Sum + b * s_sum + h0 + offs, acc + e.to(acc.dtype)) + tl.store(Codes + b * s_codes + group, tok.to(Codes.dtype.element_ty)) + + +def cp_embed_accumulate( + tokens: torch.Tensor, weight: torch.Tensor, emb_sum: torch.Tensor, + all_codes: torch.Tensor, group: int, +) -> torch.Tensor: + """For a depth step's sampled ``tokens (bs,)``: record them as ``all_codes[:, group]``, + look up their embeddings in ``weight (V, H)`` and add those to ``emb_sum (bs, H)`` in + its dtype. Returns the embeddings, the next depth step's input. One launch for what + was an index_select, a copy and a promoting add.""" + bs, hidden = tokens.shape[0], weight.shape[1] + assert weight.stride(1) == 1 and emb_sum.stride(1) == 1 and hidden % 256 == 0 + emb = torch.empty(bs, hidden, dtype=weight.dtype, device=weight.device) + _cp_embed_accumulate_kernel[(bs,)]( + tokens, weight, emb, emb_sum, all_codes, group, + weight.stride(0), emb.stride(0), emb_sum.stride(0), all_codes.stride(0), + H=hidden, BLOCK=256, num_warps=4, + ) + return emb + + + +@triton.jit +def _cp_attn_kernel(Q, K, V, O, cache_pos, sm_scale, + sq_b, sq_s, sq_h, sk_b, sk_n, sk_h, so_b, so_s, so_h, + GROUP: tl.constexpr, D: tl.constexpr, BLOCK_N: tl.constexpr): + b = tl.program_id(0) + h = tl.program_id(1) + s = tl.program_id(2) + hk = h // GROUP + d = tl.arange(0, D) + n = tl.arange(0, BLOCK_N) + q = tl.load(Q + b * sq_b + s * sq_s + h * sq_h + d).to(tl.float32) * sm_scale + valid = n <= cache_pos + s # causal within the depth sequence + kv = b * sk_b + n[:, None] * sk_n + hk * sk_h + d[None, :] + k = tl.load(K + kv, mask=valid[:, None], other=0.0).to(tl.float32) + qk = tl.where(valid, tl.sum(k * q[None, :], 1), float("-inf")) + p = tl.exp(qk - tl.max(qk, 0)) + v = tl.load(V + kv, mask=valid[:, None], other=0.0).to(tl.float32) + out = tl.sum(p[:, None] * v, 0) / tl.sum(p, 0) + tl.store(O + b * so_b + s * so_s + h * so_h + d, out.to(O.dtype.element_ty)) + + +def cp_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, cache_pos: int) -> torch.Tensor: + """Causal attention of the code predictor's ``q (bs, S, Hq, D)``, whose tokens sit + at depth positions ``cache_pos .. cache_pos + S - 1``, over its dense cache + ``(bs, max_depth, Hkv, D)``. The depth sequence is at most 16 long, so one + program holds a (batch row, head, token)'s whole key range.""" + bs, seq, hq, d = q.shape + hkv, max_depth = k_cache.shape[2], k_cache.shape[1] + assert k_cache.stride() == v_cache.stride() and q.stride(-1) == 1 and k_cache.stride(-1) == 1 + out = torch.empty(bs, seq, hq, d, device=q.device, dtype=q.dtype) + _cp_attn_kernel[(bs, hq, seq)]( + q, k_cache, v_cache, out, cache_pos, d ** -0.5, + q.stride(0), q.stride(1), q.stride(2), + k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), + out.stride(0), out.stride(1), out.stride(2), + GROUP=hq // hkv, D=d, BLOCK_N=triton.next_power_of_2(max_depth), num_warps=4, + ) + return out diff --git a/mstar/model/qwen3_omni/components/talker.py b/mstar/model/qwen3_omni/components/talker.py index 44217d5688..c4ed2afa5a 100644 --- a/mstar/model/qwen3_omni/components/talker.py +++ b/mstar/model/qwen3_omni/components/talker.py @@ -23,8 +23,8 @@ from mstar.model.components import ParallelSparseMoeBlockWithSharedExpert, RMSNorm from mstar.model.components.distributed import ParallelGatedMLP from mstar.model.qwen3_omni.components.attention import Qwen3OmniAttention +from mstar.model.qwen3_omni.components.cp_kernels import cp_attention, cp_qk_norm_rope_kv_write from mstar.model.qwen3_omni.config import TALKER_ATTN, TALKER_KV, TALKER_POS, Qwen3OmniModelConfig, TalkerTextConfig -from mstar.utils.attention import decode_attn_nhd, fused_qk_norm_rope # --------------------------------------------------------------------------- # Projection MLP (Thinker -> Talker) @@ -98,6 +98,7 @@ def __init__( comm_group=comm_group, ), comm_group=comm_group, + fused_softmax_topk=True, ) def forward( @@ -492,59 +493,40 @@ def forward_depth_unrolled( # the KVCacheEngine.has_autocast=False invariant). flat_pos = position_ids.reshape(-1) + # Each residual add rides in the next RMSNorm (FlashInfer's fused + # add + norm, in place on the stream's own tensors), so the input + # embeddings, which the caller owns, are added out of place first. + residual = None for layer_idx, layer in enumerate(self.model.layers): attn = layer.self_attn - residual = hidden_states + if residual is None: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + else: + hidden_states, residual = layer.input_layernorm.forward_residual(hidden_states, residual) - hidden_states = layer.input_layernorm(hidden_states) - - total_tokens = bs * seq_len qkv = F.linear(hidden_states, layer.self_attn.qkv_proj.weight) q_size = n_q_heads * head_dim - kv_size = n_kv_heads * head_dim - q, k, v = qkv.split((q_size, kv_size, kv_size), dim=-1) - q = q.view(total_tokens, n_q_heads, head_dim) - k = k.view(total_tokens, n_kv_heads, head_dim) - - rope_theta = float(attn.rope_theta) - q = fused_qk_norm_rope( - q, attn.q_norm.weight, - pos=flat_pos, - eps=attn.q_norm.variance_epsilon, - rope_theta=rope_theta - ) - k = fused_qk_norm_rope( - k, attn.k_norm.weight, - pos=flat_pos, - eps=attn.k_norm.variance_epsilon, - rope_theta=rope_theta - ) - - q = q.view(bs, seq_len, n_q_heads, head_dim) - k = k.view(bs, seq_len, n_kv_heads, head_dim) - v = v.view(bs, seq_len, n_kv_heads, head_dim) - - # Append K, V into dense cache at cache_pos .. cache_pos+seq_len. - kv_cache[layer_idx, :, 0, cache_pos:cache_pos + seq_len] = k - kv_cache[layer_idx, :, 1, cache_pos:cache_pos + seq_len] = v - - # Read K, V for all positions up to and including the new tokens. - valid_len = cache_pos + seq_len - k_full = kv_cache[layer_idx, :, 0] - v_full = kv_cache[layer_idx, :, 1] - - attn_out = decode_attn_nhd( - q, k_full, v_full, valid_len + # q normed and roped in place; k (normed, roped) and v into the cache + # at cache_pos .. cache_pos + seq_len. + cp_qk_norm_rope_kv_write( + qkv.view(bs * seq_len, -1), attn.q_norm.weight, attn.k_norm.weight, + flat_pos, kv_cache[layer_idx], cache_pos, seq_len, + attn.q_norm.variance_epsilon, attn.k_norm.variance_epsilon, + float(attn.rope_theta), n_q_heads, n_kv_heads, head_dim, ) - attn_out = attn_out.reshape(bs, seq_len, -1) - hidden_states = residual + attn.o_proj(attn_out) - - # Post-attention RMSNorm + dense MLP. - residual = hidden_states - hidden_states = layer.post_attention_layernorm(hidden_states) + q = qkv[..., :q_size].view(bs, seq_len, n_q_heads, head_dim) + + # Each new token attends causally over the depth positions so far. + attn_out = cp_attention(q, kv_cache[layer_idx, :, 0], kv_cache[layer_idx, :, 1], cache_pos) + attn_out = attn.o_proj(attn_out.reshape(bs, seq_len, -1)) + if layer_idx == 0: + residual = residual + attn_out + hidden_states = layer.post_attention_layernorm(residual) + else: + hidden_states, residual = layer.post_attention_layernorm.forward_residual(attn_out, residual) hidden_states = layer.mlp(hidden_states) - hidden_states = residual + hidden_states # Final norm. - hidden_states = self.model.norm(hidden_states) + hidden_states, _ = self.model.norm.forward_residual(hidden_states, residual) return hidden_states diff --git a/mstar/model/qwen3_omni/components/thinker.py b/mstar/model/qwen3_omni/components/thinker.py index 542fd991b3..982b1302ac 100644 --- a/mstar/model/qwen3_omni/components/thinker.py +++ b/mstar/model/qwen3_omni/components/thinker.py @@ -91,6 +91,7 @@ def __init__( num_experts_per_tok=tc.num_experts_per_tok, norm_topk_prob=tc.norm_topk_prob, comm_group=comm_group, + fused_softmax_topk=True, ) else: self.mlp = ParallelGatedMLP( diff --git a/mstar/model/qwen3_omni/config.py b/mstar/model/qwen3_omni/config.py index 3014cc9e7f..389cbb40f5 100644 --- a/mstar/model/qwen3_omni/config.py +++ b/mstar/model/qwen3_omni/config.py @@ -298,8 +298,8 @@ class Code2WavConfig: # Streaming chunk size (in codec frames) for the Talker→Code2Wav edge. # 25 frames keeps TTFT low (~1s). ``codec_left_context_frames`` is the # number of overlap frames prepended to non-first chunks by - # ``LeftContextChunkPolicy`` (warms up the causal ConvNet vocoder) and - # trimmed from the emitted waveform in ``Code2WavSubmodule``. + # ``ScheduledLeftContextChunkPolicy`` (warms up the vocoder) and trimmed + # from the emitted waveform in ``Code2WavSubmodule``. # # Named with a ``codec_`` prefix to avoid collision with HF's # ``Qwen3OmniMoeCode2WavConfig`` (which has no such field today but could @@ -307,6 +307,14 @@ class Code2WavConfig: # same bare name with different semantics). codec_chunk_frames: int = 25 codec_left_context_frames: int = 25 + # New frames in the first chunks, before ``codec_chunk_frames`` takes + # over: a small first chunk is the first audio after that many Talker + # steps instead of 25 (half the time to first audio). The vocoder's + # 72-frame attention window outreaches the 25-frame context, so the + # stream is not the full decode, and moving the window boundaries changes + # the waveform; boundaries stay as close to the full decode as with 25 + # (test_qwen3_omni_code2wav_real_weights.py). ``()`` restores 25. + codec_chunk_schedule: Tuple[int, ...] = (8,) attention_bias: bool = False attention_dropout: float = 0.0 codebook_dim: int = 512 diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index 15c338b864..509a106c21 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/mstar/model/qwen3_omni/qwen3_omni_model.py @@ -84,7 +84,7 @@ from mstar.model.qwen3_omni.video_reader import header_charged, install_budgeted_readers from mstar.model.submodule_base import NodeSubmodule from mstar.model.utils import Operation, WeightConverter, drop_stop_tokens -from mstar.streaming.chunk_policy import FixedChunkPolicy, LeftContextChunkPolicy +from mstar.streaming.chunk_policy import FixedChunkPolicy, ScheduledLeftContextChunkPolicy from mstar.streaming.topology import Connection, PartitionTopology, StreamingGraphEdge logger = logging.getLogger(__name__) @@ -572,7 +572,8 @@ def get_partition_topology(self) -> PartitionTopology: from_partition="Talker", to_partition="Code2Wav", edge_name="codec_tokens", - chunk_policy_factory=lambda: LeftContextChunkPolicy( + chunk_policy_factory=lambda: ScheduledLeftContextChunkPolicy( + schedule=self.config.code2wav.codec_chunk_schedule, chunk=self.config.code2wav.codec_chunk_frames, left_context=self.config.code2wav.codec_left_context_frames, ), @@ -1003,7 +1004,7 @@ def _get_code2wav_forward( metadata.graph_walk = "code2wav_chunk" step_metadata = {"consumed_tokens": chunk_size} - # Don't predict the last chunk from token counts: LeftContextChunkPolicy + # Don't predict the last chunk from token counts: the left-context policy # emits an extra flush pass for the retained overlap, so a count-based # guess completes the request before that final chunk is emitted. The # `available <= 0` check above fires once consumption actually catches up. diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index f1932a93a4..37469df36d 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -33,6 +33,7 @@ from mstar.engine.resources.attn.flashinfer import FlashInferManager from mstar.engine.resources.sampler.resource import SamplerResource from mstar.model.qwen3_omni.components.code2wav import Qwen3OmniMoeCode2Wav +from mstar.model.qwen3_omni.components.cp_kernels import cp_embed_accumulate from mstar.model.qwen3_omni.components.rope import ( compute_3d_cos_sin, compute_rope_freqs, @@ -55,8 +56,10 @@ Qwen3OmniModelConfig, ) from mstar.model.submodule_base import ( + EMPTY_INPUT_METADATA, ARNodeInputs, ARNodeSubmodule, + BatchedModelOutput, InputMetadata, ModelInputsFromEngine, NodeInputs, @@ -86,6 +89,21 @@ def _accepts_kwarg(fn: Any, name: str) -> bool: # =================================================================== + +class _CachedDeviceMixin: + """``get_device`` once: the base walks the module tree for a parameter on + every call, and prepare_inputs / preprocess call it per row per step + (~6% of the Talker GPU thread's held-GIL samples at c=32). These + submodules never move once loaded.""" + + _device: torch.device | None = None + + def get_device(self): + if self._device is None: + self._device = super().get_device() + return self._device + + class AudioEncoderSubmodule(NodeSubmodule): """Thin wrapper around the HF Whisper-style audio encoder. @@ -96,6 +114,11 @@ class AudioEncoderSubmodule(NodeSubmodule): Runs once per request (not batched across requests). """ + # The HF encoder splits its input by host-side chunk lengths (``.tolist()``), + # so compiled frames guard on the chunk count and recompile for nearly + # every new clip length, seconds each, on the request's critical path. + disable_torch_compile = True + def __init__(self, audio_encoder: nn.Module, config: Qwen3OmniModelConfig): super().__init__() self.audio_encoder = audio_encoder @@ -152,6 +175,46 @@ def forward( return {"audio_embeds": [audio_embeds]} + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: + return True + + def preprocess( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + inputs: list[NodeInputs], + ) -> dict[str, torch.Tensor | Any]: + # The encoder packs clips along time and windows its attention within + # each clip (``feature_lens``), so a batch is one call over the + # concatenated mel frames. + return { + "audio_features": torch.cat([inp.tensor_inputs["audio_features"] for inp in inputs], dim=-1), + "audio_seqlens": torch.cat([inp.tensor_inputs["audio_seqlens"].reshape(-1) for inp in inputs]), + } + + def forward_batched( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + audio_features: torch.Tensor, + audio_seqlens: torch.Tensor, + **kwargs, + ) -> dict[str, NameToTensorList]: + """Encode several requests' clips in one call (each ~10 ms on an H200, + mostly launches, whatever the clip length) and split the packed + embeddings back by each clip's output length.""" + from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + _get_feat_extract_output_lengths, + ) + + embeds = self.forward(graph_walk, engine_inputs, audio_features, audio_seqlens)["audio_embeds"][0] + # the encoder already reads its chunk lengths on the host + counts = _get_feat_extract_output_lengths(audio_seqlens).tolist() + return { + rid: {"audio_embeds": [part]} + for rid, part in zip(engine_inputs.request_ids, embeds.split(counts), strict=True) + } + # =================================================================== # 2. VisionEncoderSubmodule (enc_dec engine) @@ -257,7 +320,7 @@ def forward( # =================================================================== -class ThinkerSubmodule(ARNodeSubmodule): +class ThinkerSubmodule(_CachedDeviceMixin, ARNodeSubmodule): """Wraps the FlashInfer-based Thinker MoE transformer. Dispatches on graph_walk: @@ -292,6 +355,7 @@ def __init__( # request per step. Helps keep the captured graph's output-dict # contents self-evidently constant, too. self._decode_thinker_mask: torch.Tensor | None = None + self._decode_talker_mm_mask: torch.Tensor | None = None self._audio_bos_embed: torch.Tensor | None = None self._audio_eos_embed: torch.Tensor | None = None @@ -321,6 +385,13 @@ def _get_decode_thinker_mask(self, device: torch.device) -> torch.Tensor: ) return self._decode_thinker_mask + def _get_decode_talker_mm_mask(self, device: torch.device) -> torch.Tensor: + """The 1D multimodal mask ``postprocess`` emits for a decode token: + ``[False]``, lazily allocated per device. Never written after.""" + if self._decode_talker_mm_mask is None or self._decode_talker_mm_mask.device != device: + self._decode_talker_mm_mask = torch.zeros(1, dtype=torch.bool, device=device) + return self._decode_talker_mm_mask + def _get_talker_text_mask(self, input_ids: torch.Tensor) -> torch.Tensor: """ Cut system prompt and previous assistant parts out of the talker input @@ -981,17 +1052,13 @@ def forward_batched( [layer_0_embed, layer_0_embed], dim=-1, ) - outputs: dict[str, NameToTensorList] = {} - for i, rid in enumerate(request_ids): - req_out: NameToTensorList = { - "thinker_states": [thinker_states[i : i + 1]], - "new_token": [new_tokens[i:i+1]] - } - if masks_for_talker is not None and rid in masks_for_talker: - req_out["thinker_mask"] = [masks_for_talker[rid]] - outputs[rid] = req_out - - return outputs + # Row i is request i: one clone per output and one host copy of the + # tokens per step. A decode row's Talker mask is the constant one, + # which postprocess attaches. + return BatchedModelOutput( + row_outputs={"thinker_states": thinker_states, "new_token": new_tokens}, + check_stop_buffers={"new_token": new_tokens}, + ) def unpack_packed_outputs( self, @@ -1008,8 +1075,7 @@ def unpack_packed_outputs( reattach the per-request talker masks, which live on the original ARNodeInputs (the captured graph never saw them — masks vary in shape with text content). Drops per-rid emission for requests with - ``audio_output=False``, mirroring ``filter_batched_output``'s gating - for the decode path. + ``audio_output=False``, as postprocess does for the decode path. """ packed_states = static_output.get("__batched_thinker_states__") if packed_states is None: @@ -1059,7 +1125,15 @@ def postprocess( # shipping both hidden states for every token. ts_list = outputs.get("thinker_states") mask_list = outputs.get("thinker_mask") - if ts_list and mask_list: + if ts_list and request_info.graph_walk == "thinker_decode": + # A decode token is text the Talker keeps (``prepare_inputs`` + # gives it the constant ``[[0], [1]]`` mask), so it is layer 0 + # as is. The boolean indexing below would sync the host on + # every row of every decode step. + thinker_hidden = self.config.thinker_hidden_size + outputs["thinker_states"] = [ts_list[0][..., :thinker_hidden]] + outputs["thinker_mask"] = [self._get_decode_talker_mm_mask(ts_list[0].device)] + elif ts_list and mask_list: ts = ts_list[0] # (seq_len, 2*hidden) mask = mask_list[0] # (2, seq_len) thinker_hidden = self.config.thinker_hidden_size @@ -1095,34 +1169,12 @@ def check_stop( return {"thinker_decode_loop"} return set() - def filter_batched_output( - self, - request_info: CurrentForwardPassInfo, - rid_output: dict[str, list[torch.Tensor]], - ) -> dict[str, list[torch.Tensor]]: - """Drop ``thinker_states`` + ``thinker_mask`` for text-only requests. - - ``forward_batched`` always emits these keys so the captured CUDA - graph's output shape is static. Here, outside the captured - region, we gate them on the real request's ``audio_output`` flag - so the Talker edge stays unrouted for text-only requests (matches - the pre-capture eager-mode behaviour). - """ - if request_info is None: - return rid_output - if request_info.step_metadata.get("audio_output", True): - return rid_output - return { - k: v for k, v in rid_output.items() - if k not in ("thinker_states", "thinker_mask") - } - # =================================================================== # 4. TalkerSubmodule (ar engine) -- SECOND MOST COMPLEX # =================================================================== -class TalkerSubmodule(ARNodeSubmodule): +class TalkerSubmodule(_CachedDeviceMixin, ARNodeSubmodule): MAX_BATCH_SIZE = 32 def __init__( @@ -1143,6 +1195,11 @@ def __init__( self._tts_pad_embed_cached: torch.Tensor | None = None self._tts_bos_embed_cached: torch.Tensor | None = None self._tts_eos_embed_cached: torch.Tensor | None = None + # The same pad/eos before ``text_projection``: a decode step projects + # its rows' text inputs in the captured forward, so these stand in + # for the Thinker state once the Thinker has finished. + self._tts_pad_raw: torch.Tensor | None = None + self._tts_eos_raw: torch.Tensor | None = None # Lazy-built suppress mask for layer-0 logits. Shape (vocab_size,) # with True at positions to suppress. Cached on first forward. @@ -1155,6 +1212,8 @@ def __init__( # Lazy-built CodePredictor scratch KV cache, overwritten every step. self._cp_kv_cache: torch.Tensor | None = None + self._cp_positions: torch.Tensor | None = None + self._cp_pair_positions: torch.Tensor | None = None def _get_cp_kv_cache(self): # TODO: make this a resource @@ -1169,6 +1228,24 @@ def _get_cp_kv_cache(self): ) return self._cp_kv_cache + def _get_cp_positions(self) -> torch.Tensor: + """``[g, b, 0] = g``: the code predictor's position for depth step g, + the same for every row. Constant, so a step slices it instead of + counting up a buffer with a launch per depth step.""" + if self._cp_positions is None: + self._cp_positions = torch.arange( + self.num_codes, device=self.get_device(), + ).view(-1, 1, 1).expand(-1, self.MAX_BATCH_SIZE, 1).contiguous() + return self._cp_positions + + def _get_cp_pair_positions(self) -> torch.Tensor: + """``[b] = [0, 1]``: the positions of the depth sequence's opening pair.""" + if self._cp_pair_positions is None: + self._cp_pair_positions = torch.arange( + 2, device=self.get_device(), + ).expand(self.MAX_BATCH_SIZE, 2).contiguous() + return self._cp_pair_positions + def init_tts_embeds(self, thinker_embed_tokens: nn.Embedding) -> None: """Pre-compute TTS pad/bos/eos hidden states using the Thinker's embedding table + Talker text_projection @@ -1192,6 +1269,7 @@ def init_tts_embeds(self, thinker_embed_tokens: nn.Embedding) -> None: eos_raw = thinker_embed_tokens( torch.tensor([self.config.tts_eos_token_id], device=device) ).to(proj_dtype) + self._tts_pad_raw, self._tts_eos_raw = pad_raw, eos_raw self._tts_pad_embed_cached = self.model.text_projection(pad_raw).squeeze(0) self._tts_bos_embed_cached = self.model.text_projection(bos_raw).squeeze(0) self._tts_eos_embed_cached = self.model.text_projection(eos_raw).squeeze(0) @@ -1317,23 +1395,24 @@ def prepare_inputs( seq_len = input_embeds.shape[0] elif graph_walk == "talker_decode": - dtype = self.model.text_projection.linear_fc1.weight.dtype - input_embeds = inputs["talker_input_embeds"][0].to(dtype) - + # The previous step's codec embedding sum plus this step's text + # conditioning, which the forward projects for the whole batch: + # the Thinker's next token state (decode is always text), then + # tts_eos for one step, then tts_pad. thinker_states = inputs.get("thinker_states", []) rid = fwd_info.rid_handle if thinker_states: - # Decode is always text → text_projection. - text_hidden = self.model.text_projection( - thinker_states[0].to(dtype) - ) + text_raw = thinker_states[0] elif rid not in self._eos_embed_sent: - text_hidden = self._tts_eos_embed_cached + text_raw = self._tts_eos_raw self._eos_embed_sent.add(rid) else: - text_hidden = self._tts_pad_embed_cached - input_embeds += text_hidden - seq_len = 1 + text_raw = self._tts_pad_raw + return ARNodeInputs( + input_embeds=inputs["talker_input_embeds"][0], + input_seq_len=1, + tensor_inputs={"text_raw": text_raw}, + ) return ARNodeInputs( input_embeds=input_embeds, @@ -1390,20 +1469,12 @@ def preprocess( device = self.get_device() extra_args = {} - # codec scratch is decode-only; talker_prefill runs the LLM backbone - # alone and interning these would clash (codec_emb_sum's hidden dim - # aliases the token-bucket seq dim). - if graph_walk != "talker_prefill": - extra_args["suppress_mask"] = self._get_suppress_mask() - extra_args["all_codes"] = torch.zeros( - (len(inputs), self.num_codes), device=device, dtype=torch.long - ) - extra_args["codec_emb_sum"] = torch.zeros( - (len(inputs), self.cp_cfg.hidden_size), device=device - ) - extra_args["pos_buf"] = torch.zeros( - (len(inputs), 1), device=device, dtype=torch.long - ) + if graph_walk == "talker_decode": + dtype = self.model.text_projection.linear_fc1.weight.dtype + input_embeds = input_embeds.to(dtype) + extra_args["text_raw"] = torch.cat( + [inp.tensor_inputs["text_raw"] for inp in inputs], dim=0, + ).to(dtype) if graph_walk == "talker_last_prefill": extra_args["last_token_indices"] = ( torch.tensor(seq_lens, device=device, dtype=torch.long).cumsum(0) - 1 @@ -1413,6 +1484,15 @@ def preprocess( **extra_args } + def _add_text_conditioning( + self, input_embeds: torch.Tensor, text_raw: torch.Tensor | None, + ) -> torch.Tensor: + """A decode step's input: its codec embedding sum plus the projected + text conditioning, one ``text_projection`` for the batch.""" + if text_raw is None: + return input_embeds + return input_embeds + self.model.text_projection(text_raw) + def _forward_prefill( self, input_embeds: torch.Tensor, ): @@ -1425,10 +1505,6 @@ def _forward_decode_like( self, request_ids: list[str], talker_sampler: SamplerResource, code_pred_sampler: SamplerResource, - suppress_mask: torch.Tensor, - all_codes: torch.Tensor, - codec_emb_sum: torch.Tensor, - pos_buf: torch.Tensor, is_batched_decode: bool, input_embeds: torch.Tensor, last_token_indices: torch.Tensor | None = None, @@ -1456,7 +1532,7 @@ def _forward_decode_like( else: last_hidden = hidden logits = self.model.codec_head(last_hidden) - logits = logits.masked_fill(suppress_mask.unsqueeze(0), float("-inf")) + logits = logits.masked_fill(self._get_suppress_mask().unsqueeze(0), float("-inf")) # Apply the repetition penalty to the Talker's layer-0 codes (the code # predictor depth loop below samples through its own aux config and # stays penalty-free). The penalty is baked into the captured graph. @@ -1466,27 +1542,31 @@ def _forward_decode_like( # code predictor section embed = self.talker_code_emb(layer0_codes) # [bs, hidden] - codec_emb_sum.add_(embed) + bs = embed.shape[0] + # Scratch made here rather than passed in, so a captured step zeroes it + # in the graph instead of copying zeros into static inputs. + codec_emb_sum = embed.float() + all_codes = torch.zeros((bs, self.num_codes), device=embed.device, dtype=torch.long) all_codes[:, 0] = layer0_codes - bs = all_codes.shape[0] + positions = self._get_cp_positions() kv_cache = self._get_cp_kv_cache()[:, :bs] - # forward over [last_hidden] to update kv cache with the Talker's final hidden - # state as context for the code prediction. This returns nothing because the - # layer 0 code is already provided by the talker LLM + # The depth sequence opens with the Talker's hidden state and the layer-0 + # code's embedding (as in HF's code-predictor prefill): one two-token pass + # whose second token predicts layer 1, then one token per layer. cp = self.code_predictor codec_embedding = cp.model.codec_embedding - cp.forward_depth_unrolled( - last_hidden.unsqueeze(1), pos_buf, kv_cache, cache_pos=0, - ) - pos_buf += 1 - for group_idx in range(1, self.num_codes): - hidden = cp.forward_depth_unrolled( - embed.unsqueeze(1), pos_buf, kv_cache, cache_pos=group_idx, - ).squeeze(1) - pos_buf += 1 + if group_idx == 1: + hidden = cp.forward_depth_unrolled( + torch.stack([last_hidden, embed], dim=1), self._get_cp_pair_positions()[:bs], + kv_cache, cache_pos=0, + )[:, 1] + else: + hidden = cp.forward_depth_unrolled( + embed.unsqueeze(1), positions[group_idx, :bs], kv_cache, cache_pos=group_idx, + ).squeeze(1) logits = torch.matmul( hidden, cp.lm_head_weight[group_idx - 1].t() @@ -1497,9 +1577,9 @@ def _forward_decode_like( tokens = code_pred_sampler.sample( request_ids, logits, ) - all_codes[:, group_idx] = tokens - embed = codec_embedding[group_idx - 1](tokens) - codec_emb_sum.add_(embed) + embed = cp_embed_accumulate( + tokens, codec_embedding[group_idx - 1].weight, codec_emb_sum, all_codes, group_idx, + ) return { "talker_input_embeds": [codec_emb_sum], @@ -1518,6 +1598,7 @@ def forward( return self._forward_prefill( input_embeds=input_embeds ) + input_embeds = self._add_text_conditioning(input_embeds, kwargs.pop("text_raw", None)) return self._forward_decode_like( request_ids=engine_inputs.request_ids, talker_sampler=engine_inputs.resources[TALKER_SAMPLER], @@ -1562,6 +1643,7 @@ def forward_batched( } last_token_indices = kwargs.pop("last_token_indices", None) + input_embeds = self._add_text_conditioning(input_embeds, kwargs.pop("text_raw", None)) if graph_walk == "talker_last_prefill": assert last_token_indices is not None, ( "talker_last_prefill forward_batched requires " @@ -1578,14 +1660,18 @@ def forward_batched( **kwargs ) - outputs = { - rid: { - "talker_input_embeds": [fwd_out["talker_input_embeds"][0][i:i+1]], - "codec_tokens": [fwd_out["codec_tokens"][0][i]], - "new_token": [fwd_out["new_token"][0][i]], - } for i, rid in enumerate(engine_inputs.request_ids) - } - return outputs + # Row i is request i: the engine clones each output once for the batch + # and the stop check reads one host copy, instead of three clones and + # a device-to-host copy per request. + layer0_codes = fwd_out["new_token"][0] + return BatchedModelOutput( + row_outputs={ + "talker_input_embeds": fwd_out["talker_input_embeds"][0], + "codec_tokens": fwd_out["codec_tokens"][0], + "layer0_codes": layer0_codes, + }, + check_stop_buffers={"layer0_codes": layer0_codes}, + ) def postprocess( self, request_id: str, @@ -1649,6 +1735,10 @@ def get_accelerator_graph_configs( device=device, dtype=torch.bfloat16 ), input_seq_len=1, + tensor_inputs={"text_raw": torch.zeros( + (1, self.config.thinker_hidden_size), + device=device, dtype=torch.bfloat16, + )}, ), capture_batch_sizes=[1, 2, 4, 8, 16, 32], compile=True @@ -1708,13 +1798,6 @@ def __init__(self, code2wav_model: Qwen3OmniMoeCode2Wav, config: Qwen3OmniModelC super().__init__() self.code2wav = code2wav_model self.config = config - # Per-request set of request_ids that have already emitted their first - # audio chunk. The first chunk has no prior audio to overlap with, so - # its output must NOT be trimmed — the left-context trim only applies - # to subsequent chunks. Matches HF chunked_decode's ``context_size = - # left_context_size if start_index - left_context_size > 0 else start_index`` - # logic, where the first iteration has context_size=0. - self._first_chunk_emitted: set[str] = set() # Pre-compute the total upsample factor. HF defines this as # ``np.prod(upsample_rates + upsampling_ratios)`` — both tuples @@ -1730,10 +1813,6 @@ def __init__(self, code2wav_model: Qwen3OmniMoeCode2Wav, config: Qwen3OmniModelC self.full_seqlen = self.config.code2wav.codec_left_context_frames + \ self.config.code2wav.codec_chunk_frames - def cleanup_request(self, request_id): - super().cleanup_request(request_id) - self._first_chunk_emitted.discard(request_id) - @staticmethod def _chunk(rid, per_request_input_metadata: Mapping[str, InputMetadata] | None): meta = (per_request_input_metadata or {}).get(rid) @@ -1794,6 +1873,9 @@ def prepare_inputs( num_quantizers = self.config.code2wav.num_quantizers # 16 codec_eos = self.config.talker.codec_eos_token_id codec_tokens = inputs["codec_tokens"][0] + if codec_tokens.dim() == 3: + # stacked (1, groups) frames, the Talker's batched rows + codec_tokens = codec_tokens.flatten(0, 1) # Reshape to (num_frames, num_code_groups) if flat if codec_tokens.dim() == 1: @@ -1803,8 +1885,15 @@ def prepare_inputs( else: codec_tokens = codec_tokens.unsqueeze(0) - # Filter out codec_eos frames - if codec_tokens.dim() == 2 and codec_tokens.shape[0] > 0: + chunk = kwargs.get("input_metadata", EMPTY_INPUT_METADATA).stream_chunks.get("codec_tokens") + # The Talker's codec_eos frame ends its stream. A stream's last chunk + # drops it here, which reads the frames on the host. Any other window + # can hold it only as its last new frame (it arrived before the + # producer's done signal); reading every window for that would sync + # the GPU thread on each chunk, so ``forward_batched`` decodes the + # frame as padding and silences its samples instead: such a reply ends + # with one frame (80 ms) of zeros, the price of not syncing per chunk. + if (chunk is None or chunk.is_final) and codec_tokens.dim() == 2 and codec_tokens.shape[0] > 0: eos_mask = codec_tokens[:, 0] == codec_eos if eos_mask.any(): codec_tokens = codec_tokens[~eos_mask] @@ -1825,15 +1914,19 @@ def prepare_inputs( # Transpose to (Q, T) codec_tokens = codec_tokens.T # (Q, T) - # The frame count rides with this pass rather than in request state, - # so postprocess reads this window's count even if the request's next - # window is prepared first. + # The leading frames already decoded by an earlier window: none for + # the first, then up to ``codec_left_context_frames`` (fewer while the + # stream is shorter than that, as in HF's ``chunked_decode``). + context = chunk.context_items if chunk is not None else 0 + # The frame and context counts ride with this pass rather than in + # request state, so postprocess reads this window's counts even if the + # request's next window is prepared first. return NodeInputs( tensor_inputs={ "codec_tokens": codec_tokens, "position_ids": torch.arange(codec_tokens.shape[1], device=codec_tokens.device) }, - kwargs={"frames": orig_seq_len}, + kwargs={"frames": orig_seq_len, "context": context}, ) def preprocess( @@ -1873,12 +1966,11 @@ def forward_batched( ) -> dict[str, NameToTensorList]: """Decode a batch of padded windows to untrimmed int16 PCM. - The Talker→Code2Wav StreamBuffer uses ``LeftContextChunkPolicy``: - the first popped chunk for a request contains ``codec_chunk_frames`` - fresh frames with no overlap; every subsequent chunk contains - ``codec_chunk_frames + codec_left_context_frames`` frames where the - leading ``codec_left_context_frames`` are overlap from the previous - chunk's tail. The overlap lets the vocoder warm up at chunk + The Talker→Code2Wav StreamBuffer uses ``ScheduledLeftContextChunkPolicy``: + the first popped chunk for a request contains ``codec_chunk_schedule[0]`` + fresh frames with no overlap; every later chunk carries its new frames + behind up to ``codec_left_context_frames`` frames of overlap from the + frames already delivered. The overlap lets the vocoder warm up at chunk boundaries; ``postprocess`` cuts each request's new audio out of the window. @@ -1894,9 +1986,19 @@ def forward_batched( for rid in request_ids } + # A codec_eos frame left in a live window (see prepare_inputs) decodes + # as the zero padding it would have been, and the samples from its + # frame on, which a dropped frame would not have produced, are silenced. + eos = codec_tokens[:, :1, :] == self.config.talker.codec_eos_token_id + codec_tokens = codec_tokens.masked_fill(eos, 0) wavs = self.code2wav( codec_tokens, position_ids ) + frames = codec_tokens.shape[-1] + tail = frames * self.total_upsample - wavs.shape[-1] + first_eos = torch.where(eos.any(-1), eos.int().argmax(-1), frames).unsqueeze(-1) # (bs, 1, 1) + sample = torch.arange(wavs.shape[-1], device=wavs.device) + wavs = torch.where(sample < first_eos * self.total_upsample - tail, wavs, 0.0) results: dict[str, NameToTensorList] = {} for i, rid in enumerate(request_ids): @@ -1951,9 +2053,7 @@ def postprocess( raise ValueError("Code2Wav postprocess needs the pass's inputs for its frame count") orig_seq_len = int(inputs.kwargs["frames"]) - cfg_ctx = self.config.code2wav.codec_left_context_frames - left_context_size = 0 if request_id not in self._first_chunk_emitted else cfg_ctx - self._first_chunk_emitted.add(request_id) + left_context_size = int(inputs.kwargs["context"]) audio = outputs["audio_chunk"][0].reshape(-1) if audio.numel() == 0: outputs["audio_chunk"][0] = audio diff --git a/mstar/model/qwen3_tts/components/talker.py b/mstar/model/qwen3_tts/components/talker.py index cc2dc4f116..8eadd74300 100644 --- a/mstar/model/qwen3_tts/components/talker.py +++ b/mstar/model/qwen3_tts/components/talker.py @@ -132,13 +132,25 @@ def forward( # make inductor specialize on the int. The resources own the pages, so # module code stays clear of the allocator. self.layers[0].self_attn.attend.bind_step(label) + # Each residual add is deferred into the next norm (one FlashInfer + # fused add+RMSNorm instead of an add and a norm). The residual + # stream starts as, and is updated in place over, ``input_embeds``. + residual = None for layer_idx, layer in enumerate(self.layers): layer.self_attn.attend.set_layer_idx(layer_idx) - hidden_states = layer(hidden_states) + if residual is None: + residual = hidden_states + h = layer.input_layernorm(hidden_states) + else: + h, residual = layer.input_layernorm.forward_residual(hidden_states, residual) + h = layer.self_attn(h) + h, residual = layer.post_attention_layernorm.forward_residual(h, residual) + hidden_states = layer.mlp(h) # Sequence lengths still advance once per forward, after every layer # wrote K/V for the same packed range — the runner does it now, on the # step this forward was declared from. - return self.norm(hidden_states) + normed, _ = self.norm.forward_residual(hidden_states, residual) + return normed class Qwen3TTSTalkerModel(nn.Module): @@ -294,9 +306,16 @@ def forward_depth_unrolled( end = cache_pos + seq_len flat_position_ids = position_ids.reshape(-1).to(torch.int32) + # Residual adds are deferred into the next norm, as in the Talker; + # the residual stream is updated in place over this position's input, + # which no caller reads afterwards. + residual = None for layer_idx, layer in enumerate(self.model.layers): - residual = hidden_states - hidden_states = layer.input_layernorm(hidden_states) + if residual is None: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + else: + hidden_states, residual = layer.input_layernorm.forward_residual(hidden_states, residual) attn = layer.self_attn qkv = attn.qkv_proj(hidden_states) q_size = attn.num_heads * attn.head_dim @@ -333,10 +352,10 @@ def forward_depth_unrolled( end, ) attn_output = attn_output.reshape(batch_size, seq_len, -1) - hidden_states = residual + attn.o_proj(attn_output) - - residual = hidden_states - hidden_states = layer.post_attention_layernorm(hidden_states) - hidden_states = residual + layer.mlp(hidden_states) + hidden_states, residual = layer.post_attention_layernorm.forward_residual( + attn.o_proj(attn_output), residual, + ) + hidden_states = layer.mlp(hidden_states) - return self.model.norm(hidden_states) + normed, _ = self.model.norm.forward_residual(hidden_states, residual) + return normed diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 55ef9622b9..6f21a6010f 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -85,6 +85,7 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, @@ -118,8 +119,11 @@ class TalkerSubmodule(ARNodeSubmodule): # for this path. The whole decode forward (CodePredictor loop included) is # still CUDA-graph captured. disable_torch_compile = True - MAX_BATCH_SIZE = 32 - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # The largest decode bucket is also the cap on rows per Talker step: a + # step costs nearly the same at 64 rows as at 32 (it is launch-bound), so + # with 32 as the top, 64 streams ran as two alternating batches of 32. + DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 24, 32, 48, 64] + MAX_BATCH_SIZE = DECODE_CAPTURE_BATCH_SIZES[-1] # Prefill prompts are short (a sentence plus the ChatML frame, an optional # instruction and speaker slot): a few token buckets cover them, and the # eager alternative costs ~20 ms of launch overhead per request. @@ -445,11 +449,12 @@ def prepare_inputs( ) # ``talker_input_embeds`` is the recurrent graph edge emitted by # the previous frame, not a token ID that needs another lookup. + # The text condition is added for the whole batch in + # ``preprocess`` (one kernel instead of one per row). input_embeds = inputs["talker_input_embeds"][0].to( - device=self.get_device(), + device=text_condition.device, dtype=text_condition.dtype, ).reshape(1, -1) - input_embeds = input_embeds + text_condition state.add("generation_step", step + 1) else: raise ValueError(f"Unknown Qwen3-TTS Talker walk: {graph_walk!r}") @@ -458,10 +463,13 @@ def prepare_inputs( int(state.get("generated_frames", 0)) < self.config.generation.min_new_tokens ) + kwargs = {"suppress_eos": suppress_eos} + if graph_walk == "talker_decode": + kwargs["text_condition"] = text_condition return ARNodeInputs( input_embeds=input_embeds, input_seq_len=input_embeds.shape[0], - kwargs={"suppress_eos": suppress_eos}, + kwargs=kwargs, ) def preprocess( @@ -486,13 +494,12 @@ def preprocess( if graph_walk == "talker_decode": if any(seq_len != 1 for seq_len in seq_lens): raise ValueError("Qwen3-TTS decode expects one token per request") - # one token per request, so the indices are a fixed arange and a - # single-request batch is already packed - packed_embeds = ( - input_embeds[0] - if len(input_embeds) == 1 - else torch.cat(input_embeds, dim=0) - ) + # one token per request, so the indices are a fixed arange + conditions = [item.kwargs["text_condition"] for item in inputs] + if len(input_embeds) == 1: + packed_embeds = input_embeds[0] + conditions[0] + else: + packed_embeds = torch.cat(input_embeds, dim=0) + torch.stack(conditions) last_token_indices = self._get_decode_last_token_indices(len(inputs)) else: packed_embeds = torch.cat(input_embeds, dim=0) @@ -685,14 +692,16 @@ def _codec_stream_items(self, graph_walk: str, request_id: str, frame: torch.Ten ``left_context_frames`` matter, since that is all the context a window ever carries, and every reference frame the stream carries costs a codec pass whose audio ``CodecSubmodule`` then trims. Decode - steps (the captured path) always push exactly one frame. + steps (the captured path) always push exactly one frame. Every item is + ``[1, groups]``, the shape of a row of the batched outputs, so one + stream's items always stack. """ if graph_walk != "talker_prefill_clone": return [frame] reference = self.request_state(request_id).get("reference_frames") if reference is None: return [frame] - return [*reference[-self.config.codec.left_context_frames:].unbind(0), frame] + return [*reference[-self.config.codec.left_context_frames:].split(1), frame] def forward( self, @@ -707,10 +716,10 @@ def forward( output = self._run_frame( engine_inputs, input_embeds, last_token_indices, suppress_eos ) - request_id = engine_inputs.request_ids[0] + del graph_walk return { "talker_input_embeds": [output["talker_input_embeds"]], - "codec_tokens": self._codec_stream_items(graph_walk, request_id, output["codec_tokens"][0]), + "codec_tokens": [output["codec_tokens"][:1]], "new_token": [output["new_token"]], } @@ -722,19 +731,19 @@ def forward_batched( last_token_indices: torch.Tensor, suppress_eos: torch.Tensor, **kwargs: Any, - ) -> dict[str, NameToTensorList]: + ) -> dict[str, NameToTensorList] | BatchedModelOutput: del kwargs output = self._run_frame( engine_inputs, input_embeds, last_token_indices, suppress_eos ) - return { - request_id: { - "talker_input_embeds": [output["talker_input_embeds"][i:i + 1]], - "codec_tokens": self._codec_stream_items(graph_walk, request_id, output["codec_tokens"][i]), - "new_token": [output["new_token"][i]], - } - for i, request_id in enumerate(engine_inputs.request_ids) - } + del graph_walk + # One frame per row: the engine clones each batch tensor once and the + # stop check copies one host buffer, rather than three of each per + # request. A clone prefill's reference frames join in ``postprocess``. + return BatchedModelOutput( + row_outputs=output, + check_stop_buffers={"layer0_codes": output["new_token"]}, + ) def postprocess( self, @@ -750,7 +759,11 @@ def postprocess( as a fallback. This keeps EOS detection correct even if an execution path filters the sampler-only output before slow-path ``check_stop``. """ - del request_info, kwargs + del kwargs + if request_info.graph_walk == "talker_prefill_clone" and "codec_tokens" in outputs: + outputs["codec_tokens"] = self._codec_stream_items( + request_info.graph_walk, request_id, outputs["codec_tokens"][-1], + ) if "new_token" in outputs: outputs["layer0_codes"] = outputs.pop("new_token") elif "layer0_codes" not in outputs and "codec_tokens" in outputs: @@ -792,6 +805,33 @@ def check_stop( return {"talker_decode_loop"} return set() + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for every row off one host copy of the group-0 codes.""" + codes = host_rows.buffers.get("layer0_codes") + if not torch.is_tensor(codes) or codes.dim() == 0: + return None + values = codes.reshape(-1).tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + eos = self.talker_config.codec_eos_token_id + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + generated = int(self.request_state(rid).get("generated_frames", 0)) + max_tokens = info.step_metadata.get("talker_max_tokens", info.max_tokens) + ignore_eos = bool(getattr(info.resource_configs.get(TALKER_SAMPLER), "ignore_eos", False)) + if (values[i] == eos and not ignore_eos) or generated >= max_tokens: + stops[rid] = {"talker_decode_loop"} + return stops + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Admit compatible prefill/decode requests to continuous batching. @@ -817,9 +857,8 @@ def get_accelerator_graph_configs( Dynamic EOS suppression is carried by ``ARNodeInputs`` and packed into a graph input, so replay never consults capture-slot dummy request state. - Prefill replays a packed capture of the smallest token bucket that - holds the batch; the clone prefill stays eager because it also pushes - the reference clip's frames into the codec stream. + Prefill (the clone prefill too) replays a packed capture of the + smallest token bucket that holds the batch. """ del tp_world_size dtype = self.model.model.codec_embedding.weight.dtype @@ -834,6 +873,9 @@ def prefill_input(num_tokens: int) -> ARNodeInputs: prefill = PackedAcceleratorGraphConfig( capture_graph_walk="talker_prefill", + # The clone prefill is the same forward over a longer prompt; its + # reference frames join the codec stream in ``postprocess``. + replay_graph_walks=["talker_prefill", "talker_prefill_clone"], capture_token_lengths=self.PREFILL_TOKEN_BUCKETS, make_node_input=prefill_input, capture_batch_sizes=self.PREFILL_CAPTURE_BATCH_SIZES, @@ -850,8 +892,13 @@ def prefill_input(num_tokens: int) -> ARNodeInputs: ), input_seq_len=1, # Padding slots clone this template, so every replay input must - # declare the same dynamic key as a real prepared request. - kwargs={"suppress_eos": True}, + # declare the same dynamic keys as a real prepared request. + kwargs={ + "suppress_eos": True, + "text_condition": torch.zeros( + self.talker_config.hidden_size, dtype=dtype, device=device, + ), + }, ), capture_batch_sizes=self.DECODE_CAPTURE_BATCH_SIZES, compile=True, @@ -920,7 +967,7 @@ def can_use_accelerator_graphs( ) -> bool: """Replay the whole decode graph or a packed prefill; sampling params are read from buffers, so no request's settings can disqualify it.""" - if batch.graph_walk not in ("talker_decode", "talker_prefill") or not self.can_batch( + if batch.graph_walk not in ("talker_decode", "talker_prefill", "talker_prefill_clone") or not self.can_batch( batch, model_inputs ): return False @@ -961,8 +1008,14 @@ class CodecSubmodule(ARNodeSubmodule): # 12 Hz decoder, 2.8 s of audio) whatever the window, so its activations stay small # enough to capture batches of 32 next to the Talker; ``can_batch`` keeps # the ceiling. - MAX_BATCH_SIZE = 32 - CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # The codec's cost is linear in rows (about 2.3 ms a full window at any + # batch size on an H200), so every padding row of a bucket costs a real + # row's decode: the buckets are dense. + CAPTURE_BATCH_SIZES = [1, 2, 3, 4, 5, 6, 8, 10, 12, 16, 20, 24, 28, 32] + MAX_BATCH_SIZE = CAPTURE_BATCH_SIZES[-1] + # The opt-in watermark's captures are per chunk length and cost far more + # memory each (see docs/models.rst), so they keep power-of-two buckets. + WATERMARK_BATCH_SIZES = [1, 2, 4, 8, 16, 32] def __init__( self, @@ -1052,12 +1105,13 @@ def prepare_inputs( codes = inputs["codec_tokens"][0].to( device=self.get_device(), dtype=torch.long ) - if codes.ndim == 1: - codes = codes.view(-1, self.config.num_code_groups) - if codes.ndim != 2: + # Items are one frame each, ``[groups]`` or ``[1, groups]`` (a row + # view of the Talker's batch output), so the window flattens to frames. + if codes.shape[-1] != self.config.num_code_groups: raise ValueError( f"Expected codec tokens with shape (frames, groups), got {codes.shape}" ) + codes = codes.reshape(-1, self.config.num_code_groups) num_items = codes.shape[0] # EOS belongs to Talker loop control and is not a valid codec codebook # index for waveform reconstruction. It is only ever the last frame of @@ -1252,7 +1306,7 @@ def get_piecewise_accelerator_graph_configs( ] if not lengths: return {} - return self.watermarker.piecewise_configs(lengths, self.CAPTURE_BATCH_SIZES) + return self.watermarker.piecewise_configs(lengths, self.WATERMARK_BATCH_SIZES) def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Any mix of windows batches (``preprocess`` pads to the widest).""" @@ -1263,6 +1317,17 @@ def max_batch_size(self, graph_walk: str) -> int: del graph_walk return self.MAX_BATCH_SIZE + def split_batches_by_capture_key(self, graph_walk: str) -> bool: + """Batch each window bucket apart. + + A mixed batch pads every row to its widest window, and the codec's + cost is per padded frame: a ramp window (1 to 28 frames) batched with + a full one decodes 35 frames instead of a few. Kept apart, the small + ramp windows, a stream's first audio among them, also do not wait + behind a batch of full windows. + """ + return graph_walk == "codec_chunk" + def cg_key_info( self, graph_walk: str, diff --git a/mstar/model/vjepa2/components/ac_predictor.py b/mstar/model/vjepa2/components/ac_predictor.py index c2ca8baaa4..8c81f46c06 100644 --- a/mstar/model/vjepa2/components/ac_predictor.py +++ b/mstar/model/vjepa2/components/ac_predictor.py @@ -141,6 +141,7 @@ def forward( h_pos: torch.Tensor | None = None, w_pos: torch.Tensor | None = None, time_pos: torch.Tensor | None = None, + rope_tables: dict | None = None, ) -> torch.Tensor: if label is not None and self.attn is not None: if d_pos is None: @@ -150,7 +151,7 @@ def forward( return self.forward_cached( x=x, d_pos=d_pos, h_pos=h_pos, w_pos=w_pos, time_pos=time_pos, - action_tokens=action_tokens, + action_tokens=action_tokens, rope_tables=rope_tables, ) b, n, c = x.size() @@ -231,6 +232,7 @@ def forward_cached( w_pos: torch.Tensor, # [H*W] — width positions time_pos: torch.Tensor | None, # [action_tokens] or None action_tokens: int, + rope_tables: dict | None = None, # see rotate_queries_or_keys_BNHD ) -> torch.Tensor: """Single-frame cached attention with pre-computed position tensors. @@ -261,14 +263,14 @@ def forward_cached( # 3-axis RoPE for spatial tokens s = 0 - qd = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.d_dim], d_pos) - kd = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.d_dim], d_pos) + qd = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.d_dim], d_pos, rope_tables) + kd = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.d_dim], d_pos, rope_tables) s += self.d_dim - qh = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.h_dim], h_pos) - kh = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.h_dim], h_pos) + qh = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.h_dim], h_pos, rope_tables) + kh = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.h_dim], h_pos, rope_tables) s += self.h_dim - qw = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.w_dim], w_pos) - kw = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.w_dim], w_pos) + qw = rotate_queries_or_keys_BNHD(q_spatial[..., s:s + self.w_dim], w_pos, rope_tables) + kw = rotate_queries_or_keys_BNHD(k_spatial[..., s:s + self.w_dim], w_pos, rope_tables) s += self.w_dim if s < hd: q_spatial = torch.cat([qd, qh, qw, q_spatial[..., s:]], dim=-1) @@ -279,8 +281,8 @@ def forward_cached( # Temporal RoPE for action tokens if action_tokens > 0 and time_pos is not None: - qd = rotate_queries_or_keys_BNHD(q_action[..., :self.d_dim], time_pos) - kd = rotate_queries_or_keys_BNHD(k_action[..., :self.d_dim], time_pos) + qd = rotate_queries_or_keys_BNHD(q_action[..., :self.d_dim], time_pos, rope_tables) + kd = rotate_queries_or_keys_BNHD(k_action[..., :self.d_dim], time_pos, rope_tables) q_action = torch.cat([qd, q_action[..., self.d_dim:]], dim=-1) k_action = torch.cat([kd, k_action[..., self.d_dim:]], dim=-1) q = torch.cat([q_action, q_spatial], dim=1) @@ -340,11 +342,13 @@ def forward( h_pos: torch.Tensor | None = None, w_pos: torch.Tensor | None = None, time_pos: torch.Tensor | None = None, + rope_tables: dict | None = None, ) -> torch.Tensor: x = x + self.attn( self.norm1(x), attn_mask=attn_mask, t=t, h=h, w=w, action_tokens=action_tokens, t_0=t_0, label=label, d_pos=d_pos, h_pos=h_pos, w_pos=w_pos, time_pos=time_pos, + rope_tables=rope_tables, ) x = x + self.mlp(self.norm2(x)) return x @@ -514,6 +518,8 @@ def fn(x: torch.Tensor) -> torch.Tensor: if label is not None: # cursors on the shared resources; see AttentionCallable blocks[0].attn.attend.bind_step(label) + # every block rotates by this step's positions: one sin/cos table each + rope_tables: dict = {} for blk_num, blk in enumerate(blocks): if label is not None: blk.attn.attend.set_layer_idx(blk_num) @@ -525,6 +531,7 @@ def fn(x: torch.Tensor) -> torch.Tensor: action_tokens=cond_tokens, label=label, d_pos=d_pos, h_pos=h_pos, w_pos=w_pos, time_pos=time_pos, + rope_tables=rope_tables, ) return x @@ -568,6 +575,7 @@ def forward( if label is not None: # cursors on the shared resources; see AttentionCallable self.predictor_blocks[0].attn.attend.bind_step(label) + rope_tables: dict | None = {} if label is not None else None for blk_num, blk in enumerate(self.predictor_blocks): if label is not None: blk.attn.attend.set_layer_idx(blk_num) @@ -581,6 +589,7 @@ def forward( t_0=t_0, label=label, d_pos=d_pos, h_pos=h_pos, w_pos=w_pos, time_pos=time_pos, + rope_tables=rope_tables, ) # the advance is the runner's now, off the step declaration diff --git a/mstar/model/vjepa2/components/rope_utils.py b/mstar/model/vjepa2/components/rope_utils.py index 9544c40886..496a1fbbfb 100644 --- a/mstar/model/vjepa2/components/rope_utils.py +++ b/mstar/model/vjepa2/components/rope_utils.py @@ -42,10 +42,14 @@ def rotate_queries_or_keys(x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: return (x * emb_cos) + (y * emb_sin) -def rotate_queries_or_keys_BNHD(x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: - _, _, _, D = x.size() - omega = torch.arange(D // 2, dtype=x.dtype, device=x.device) - omega /= D / 2.0 +def rope_tables_BNHD( + pos: torch.Tensor, dim: int, dtype: torch.dtype, device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """``(sin, cos)`` for ``rotate_queries_or_keys_BNHD`` over the last ``dim`` dims, + computed exactly as it computes them, so callers can build them once per step + and reuse them across heads, q/k and layers.""" + omega = torch.arange(dim // 2, dtype=dtype, device=device) + omega /= dim / 2.0 omega = 1.0 / 10000**omega if pos.dim() == 1: @@ -55,8 +59,26 @@ def rotate_queries_or_keys_BNHD(x: torch.Tensor, pos: torch.Tensor) -> torch.Ten emb_sin = freq.sin().repeat(1, 1, 1, 2) # [B, N, num_heads, D] emb_cos = freq.cos().repeat(1, 1, 1, 2) + return emb_sin, emb_cos + +def apply_rope_BNHD(x: torch.Tensor, emb_sin: torch.Tensor, emb_cos: torch.Tensor) -> torch.Tensor: y = x.unflatten(-1, (-1, 2)) y1, y2 = y.unbind(dim=-1) y = torch.stack((-y2, y1), dim=-1).flatten(-2) return (x * emb_cos) + (y * emb_sin) + + +def rotate_queries_or_keys_BNHD( + x: torch.Tensor, pos: torch.Tensor, tables: dict | None = None, +) -> torch.Tensor: + """``tables`` memoizes the sin/cos per (position tensor, dim, dtype) for the + caller's step: every layer of a rollout step rotates by the same positions.""" + _, _, _, D = x.size() + if tables is None: + return apply_rope_BNHD(x, *rope_tables_BNHD(pos, D, x.dtype, x.device)) + key = (id(pos), D, x.dtype) + cached = tables.get(key) + if cached is None: + cached = tables[key] = rope_tables_BNHD(pos, D, x.dtype, x.device) + return apply_rope_BNHD(x, *cached) diff --git a/mstar/model/vjepa2/config.py b/mstar/model/vjepa2/config.py index 84b77daef0..52d875eb38 100644 --- a/mstar/model/vjepa2/config.py +++ b/mstar/model/vjepa2/config.py @@ -14,6 +14,10 @@ KV_CACHE = "kv" ATTN = "attn" +# Frames of context an AC rollout request's video is trimmed to, matching +# upstream's DROID training config (configs/train/vitg16/droid-256px-8f.yaml). +AC_ROLLOUT_NUM_FRAMES = 8 + @dataclass class VJepa2ACPredictorConfig: diff --git a/mstar/model/vjepa2/submodules.py b/mstar/model/vjepa2/submodules.py index c5e79a9d23..cadc111b73 100644 --- a/mstar/model/vjepa2/submodules.py +++ b/mstar/model/vjepa2/submodules.py @@ -30,6 +30,8 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.accelerator_graph_config import ( + AcceleratorGraphConfig, + BatchedAcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig, PiecewiseBatchedConfig, PiecewiseCallInputs, @@ -42,7 +44,7 @@ from mstar.model.vjepa2.components.ac_predictor import VisionTransformerPredictorAC from mstar.model.vjepa2.components.predictor import VJEPA2Predictor from mstar.model.vjepa2.components.vit_encoder import VJEPA2Encoder -from mstar.model.vjepa2.config import ATTN, KV_CACHE, VJepa2Config +from mstar.model.vjepa2.config import AC_ROLLOUT_NUM_FRAMES, ATTN, KV_CACHE, VJepa2Config # the AC predictor's transformer block loop, as a piecewise capture region BLOCK_LOOP_REGION = "block_loop" @@ -143,11 +145,53 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs ) -> NodeInputs: - return NodeInputs( - tensor_inputs={ - "video_frames": _normalize_video_frames(inputs["video_frames"][0]) - } - ) + frames = _normalize_video_frames(inputs["video_frames"][0]) + return NodeInputs(tensor_inputs={"video_frames": frames}, input_seq_len=frames.size(1)) + + # AC rollout clips: 8 frames, each encoded as a 2-frame self-tubelet of 256 + # tokens, which leaves the ViT-g forward launch-bound (about the same time + # at 1 clip as at 16), so it replays a captured graph. + _AC_ROLLOUT_WALKS = ("prefill_video_rollout", "prefill_video_rollout_streaming") + + def _ac_rollout_key(self) -> tuple[int, int]: + return (AC_ROLLOUT_NUM_FRAMES, self.config.crop_size) + + def get_accelerator_graph_configs( + self, device: torch.device, tp_world_size: int = 1 + ) -> list[AcceleratorGraphConfig]: + del tp_world_size + if self.config.predictor_kind != "ac": + return [] + frames, size = self._ac_rollout_key() + clip = torch.zeros((1, frames, self.config.in_chans, size, size), dtype=torch.float32, device=device) + return [BatchedAcceleratorGraphConfig( + capture_graph_walk=self._AC_ROLLOUT_WALKS[0], + replay_graph_walks=list(self._AC_ROLLOUT_WALKS), + single_request_inputs=NodeInputs(tensor_inputs={"video_frames": clip}, input_seq_len=frames), + capture_batch_sizes=[1, 2, 4], + capture_forward_method="forward_batched", + additional_key_info=self._ac_rollout_key(), + input_seq_dims={"pixel_values_videos": 0}, + # The engine's own compile of the eager forward; max-autotune's + # kernels move this encoder's output by ~4% (relative L2). + compile_mode="default", + )] + + def cg_key_info( + self, graph_walk: str, per_request_info: dict[str, CurrentForwardPassInfo], + per_request_input_metadata=None, **kwargs, + ) -> tuple[int, int] | None: + """The AC rollout capture for a batch of full 8-frame clips at the crop + size, else None (eager): a shorter clip encodes a different shape.""" + del per_request_input_metadata, kwargs + if self.config.predictor_kind != "ac" or graph_walk not in self._AC_ROLLOUT_WALKS: + return None + frames, size = self._ac_rollout_key() + want = (frames, self.config.in_chans, size, size) + for info in per_request_info.values(): + if tuple(info.step_metadata.get("video_shape") or ()) != want: + return None + return self._ac_rollout_key() def preprocess( self, diff --git a/mstar/model/vjepa2/vjepa2_model.py b/mstar/model/vjepa2/vjepa2_model.py index 033346ca88..ef462ab7f7 100644 --- a/mstar/model/vjepa2/vjepa2_model.py +++ b/mstar/model/vjepa2/vjepa2_model.py @@ -54,12 +54,18 @@ ) from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata, video_metadata_dict -from mstar.model.media_limits import VideoBudget +from mstar.model.media_limits import VIDEO_DECODE_THREADS, VideoBudget from mstar.model.submodule_base import NodeSubmodule from mstar.model.vjepa2.components.ac_predictor import VisionTransformerPredictorAC from mstar.model.vjepa2.components.predictor import VJEPA2Predictor from mstar.model.vjepa2.components.vit_encoder import VJEPA2Encoder -from mstar.model.vjepa2.config import ATTN, KV_CACHE, VJepa2ACPredictorConfig, VJepa2Config +from mstar.model.vjepa2.config import ( + AC_ROLLOUT_NUM_FRAMES, + ATTN, + KV_CACHE, + VJepa2ACPredictorConfig, + VJepa2Config, +) from mstar.model.vjepa2.submodules import ( VJepa2ACPredictorSubmodule, VJepa2ACRolloutPredictorSubmodule, @@ -579,7 +585,7 @@ def load_video(self, filepath: str, device: str) -> TensorAndMetadata: target_frames = self.config.frames_per_clip logger.info("load_video: opening %s", filepath) - decoder = VideoDecoder(filepath) + decoder = VideoDecoder(filepath, num_ffmpeg_threads=VIDEO_DECODE_THREADS) logger.info("load_video: decoder opened; reading metadata") metadata_obj = getattr(decoder, "metadata", None) @@ -633,7 +639,7 @@ def load_video(self, filepath: str, device: str) -> TensorAndMetadata: frames = torch.stack([decoder[i] for i in indices]) logger.info("load_video: decoded frames shape=%s", tuple(frames.shape)) - video = frames.float() / 255.0 + video = frames.float().div_(255.0) try: metadata = video_metadata_dict(metadata_obj) if metadata_obj is not None else {} @@ -730,7 +736,6 @@ def process_prompt( # (per notebook Cell 5: z_hat = z[:, :tokens_per_frame]); the # other 7 frames' tokens are computed and discarded — matched # FLOPs to upstream's reference inference path. - AC_ROLLOUT_NUM_FRAMES = 8 if "video_frames" in out: out["video_frames"] = [frames[:AC_ROLLOUT_NUM_FRAMES] for frames in out["video_frames"]] @@ -895,6 +900,10 @@ def get_initial_forward_pass_args( inputs.append(edge) step_metadata: dict = {"is_prefill": True} + if input_signals.get("video_frames"): + # The encoder's capture bucket is chosen before its inputs exist + # (VJepa2EncoderSubmodule.cg_key_info), so the clip shape rides here. + step_metadata["video_shape"] = tuple(int(d) for d in input_signals["video_frames"][0].dims) if walk in (self.PREFILL_VIDEO_ROLLOUT, self.PREFILL_VIDEO_ROLLOUT_STREAMING): # Per-request horizon enforced by ``VJepa2RolloutPredictorSubmodule`` # via ``check_stop`` once iter_idx + 1 reaches it. The diff --git a/mstar/model/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index 8a8fb97515..61b317a4a6 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -2,8 +2,9 @@ Exact port of diffusers 0.39.0 ``WanTransformer3DModel``, restricted to the TI2V-5B checkpoint's configuration: no image-conditioning branch, and the -per-token timestep path (the dit submodule always feeds a ``[batch, seq]`` -timestep grid). +per-token timestep path (``expand_timesteps``; the dit submodule feeds one +timestep per latent frame, which is the same thing since it is constant within +a frame). Numerics contract (bf16 weights with fp32 islands): @@ -29,11 +30,13 @@ """ import math +from typing import NamedTuple import torch import torch.nn.functional as F from torch import nn +from mstar.model.components.diffusion.attention import RaggedAttentionFn, joint_attention from mstar.model.components.mlp import MLP from mstar.model.wan22.config import Wan22Config @@ -148,6 +151,14 @@ def _apply_rotary_emb( return out.type_as(hidden_states) +class Wan22AttentionFns(NamedTuple): + """The engine's ragged attention callables for one step, per attention; a + None (or no ``Wan22AttentionFns`` at all) runs SDPA, the reference kernel.""" + + self_attn: RaggedAttentionFn | None + cross_attn: RaggedAttentionFn | None + + class Wan22DiTAttention(nn.Module): """Multi-head attention with across-heads RMSNorm on q/k (the checkpoint's ``rms_norm_across_heads``: one 3072-wide norm before the @@ -170,6 +181,7 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attn_fn: RaggedAttentionFn | None = None, ) -> torch.Tensor: kv_source = hidden_states if encoder_hidden_states is None else encoder_hidden_states query = self.q_norm(self.q_proj(hidden_states)) @@ -184,13 +196,9 @@ def forward( query = _apply_rotary_emb(query, *rotary_emb) key = _apply_rotary_emb(key, *rotary_emb) - # [B, S, H, D] -> [B, H, S, D] SDPA -> back (the reference's - # dispatch_attention_fn native path; bidirectional, no mask). - out = F.scaled_dot_product_attention( - query.permute(0, 2, 1, 3), - key.permute(0, 2, 1, 3), - value.permute(0, 2, 1, 3), - ).permute(0, 2, 1, 3) + # Bidirectional, no mask: SDPA (the reference's dispatch_attention_fn + # native path) without a ragged callable, one segment per row with one. + out = joint_attention(query, key, value, ragged=attn_fn) return self.o_proj(out.flatten(2, 3)) @@ -225,6 +233,26 @@ def forward( return temb, timestep_proj, self.text_embedder(encoder_hidden_states) +def _groups(tokens: torch.Tensor, groups: int) -> torch.Tensor: + """``[B, seq, dim]`` viewed as ``[B, groups, seq / groups, dim]``: the token + groups that share one timestep, hence one adaLN modulation row. Per-token + timesteps are ``groups == seq``; the serving path passes one per latent frame + (``expand_timesteps`` is constant within a frame), which embeds and projects + the timestep once per frame instead of once per token, with the same + elementwise math on every token.""" + return tokens.unflatten(1, (groups, -1)) + + +def _ffn_gelu_epilogue(mlp: MLP, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)`` with the tanh-GELU in the up-projection's cuBLASLt epilogue: one + kernel instead of a GEMM plus a pass over its [tokens, 14336] output. The GELU + reads the fp32 accumulator, so this rounds once fewer than the reference.""" + hidden = torch._addmm_activation( + mlp.linear_in.bias, x.flatten(0, -2), mlp.linear_in.weight.t(), use_gelu=True, + ) + return mlp.linear_out(hidden.unflatten(0, x.shape[:-1])) + + class Wan22DiTBlock(nn.Module): """One transformer block: adaLN-modulated self-attention, cross-attention to the text stream, adaLN-modulated FFN (reference ``WanTransformerBlock``, @@ -244,6 +272,9 @@ def __init__(self, dim: int, ffn_dim: int, num_heads: int, eps: float): self.ffn = MLP(dim, ffn_dim, dim, activation="gelu_tanh", bias=True) self.norm3 = WanFP32LayerNorm(dim, eps, elementwise_affine=False) self.scale_shift_table = nn.Parameter(torch.empty(1, 6, dim)) + # Serving sets this (Wan22DiT.fuse_ffn_gelu); the reference path keeps + # the separate GELU. + self.fuse_ffn_gelu = False def forward( self, @@ -251,25 +282,39 @@ def forward( encoder_hidden_states: torch.Tensor, timestep_proj: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], + attn_fns: Wan22AttentionFns | None = None, ) -> torch.Tensor: - # [1, 1, 6, dim] + [B, seq, 6, dim] -> six [B, seq, dim] fp32 tensors. + # [1, 1, 6, dim] + [B, G, 6, dim] -> six [B, G, 1, dim] fp32 tensors, one + # row per group of seq / G consecutive tokens (see _groups). + groups = timestep_proj.shape[1] shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( (self.scale_shift_table.unsqueeze(0) + timestep_proj.float()).chunk(6, dim=2) ) - shift_msa, scale_msa, gate_msa = shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2) - c_shift_msa, c_scale_msa, c_gate_msa = c_shift_msa.squeeze(2), c_scale_msa.squeeze(2), c_gate_msa.squeeze(2) - norm_hidden = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) - attn_output = self.self_attn(norm_hidden, rotary_emb=rotary_emb) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + grouped = _groups(hidden_states, groups) + norm_hidden = (self.norm1(grouped.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.self_attn( + norm_hidden.flatten(1, 2), rotary_emb=rotary_emb, attn_fn=attn_fns.self_attn if attn_fns else None, + ) + hidden_states = ( + (grouped.float() + _groups(attn_output, groups) * gate_msa).type_as(hidden_states).flatten(1, 2) + ) norm_hidden = self.norm2(hidden_states.float()).type_as(hidden_states) - attn_output = self.cross_attn(norm_hidden, encoder_hidden_states=encoder_hidden_states) + attn_output = self.cross_attn( + norm_hidden, encoder_hidden_states=encoder_hidden_states, + attn_fn=attn_fns.cross_attn if attn_fns else None, + ) hidden_states = hidden_states + attn_output - norm_hidden = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states) - ff_output = self.ffn(norm_hidden) - return (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + grouped = _groups(hidden_states, groups) + norm_hidden = ( + (self.norm3(grouped.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states).flatten(1, 2) + ) + ff_output = _ffn_gelu_epilogue(self.ffn, norm_hidden) if self.fuse_ffn_gelu else self.ffn(norm_hidden) + return ( + (grouped.float() + _groups(ff_output, groups).float() * c_gate_msa).type_as(hidden_states).flatten(1, 2) + ) class Wan22DiT(nn.Module): @@ -304,6 +349,13 @@ def __init__(self, config: Wan22Config): self.proj_out = nn.Linear(dim, config.out_channels * math.prod(config.patch_size)) self.scale_shift_table = nn.Parameter(torch.empty(1, 2, dim)) + def set_fuse_ffn_gelu(self, enabled: bool) -> None: + """Run every block's FFN GELU in the GEMM epilogue (see + ``_ffn_gelu_epilogue``); off by default, so the module alone stays the + reference port.""" + for block in self.blocks: + block.fuse_ffn_gelu = enabled + @property def dtype(self) -> torch.dtype: """Bulk compute dtype (the non-island weights); callers cast inputs @@ -328,17 +380,21 @@ def forward( hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, + attn_fns: Wan22AttentionFns | None = None, ) -> torch.Tensor: """``hidden_states`` ``[B, C, F, H, W]`` (bf16 latents), ``timestep`` ``[B, post-patch seq]`` per-token grid (``expand_timesteps`` — the only timestep form mstar's dit submodule produces), text embeds - ``[B, 512, text_dim]`` bf16. Returns ``[B, C_out, F, H, W]`` bf16.""" + ``[B, 512, text_dim]`` bf16. Returns ``[B, C_out, F, H, W]`` bf16. + ``attn_fns`` routes attention through the engine's ragged resources.""" batch_size, _, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size grid = (num_frames // p_t, height // p_h, width // p_w) - if timestep.ndim != 2: + seq_len = math.prod(grid) + if timestep.ndim != 2 or seq_len % timestep.shape[1]: raise ValueError( - f"Wan22DiT requires a per-token [batch, seq] timestep grid; got ndim={timestep.ndim}." + "Wan22DiT requires a [batch, groups] timestep, one per group of consecutive " + f"tokens (per token or per latent frame); got {tuple(timestep.shape)} for {seq_len} tokens." ) rotary_emb = self.rope(grid, hidden_states.device) @@ -350,14 +406,14 @@ def forward( ) for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, attn_fns) # Output head: 2-chunk adaLN in fp32 ([1, 1, 2, dim] island table + # bf16 temb promotes), then project and unpatchify. shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) - shift, scale = shift.squeeze(2), scale.squeeze(2) - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) - hidden_states = self.proj_out(hidden_states) + grouped = _groups(hidden_states, temb.shape[1]) + hidden_states = (self.norm_out(grouped.float()) * (1 + scale) + shift).type_as(hidden_states) + hidden_states = self.proj_out(hidden_states.flatten(1, 2)) hidden_states = hidden_states.reshape(batch_size, *grid, p_t, p_h, p_w, -1) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) diff --git a/mstar/model/wan22/config.py b/mstar/model/wan22/config.py index 86b639f0f6..8e059ec2ae 100644 --- a/mstar/model/wan22/config.py +++ b/mstar/model/wan22/config.py @@ -74,6 +74,11 @@ class Wan22Config: # build with this False so the eager path stays the bit-exact reference. compile_dit: bool = True + # The DiT's attention kernel: "flashinfer" runs both attentions through the + # engine's ragged resources (FlashInfer's ragged prefill, FA3 on Hopper); + # "sdpa" runs torch SDPA, the reference pipeline's kernel. + attention_backend: str = "flashinfer" + # VAE-decode tiling policy: "auto" decodes untiled when free VRAM comfortably # exceeds the estimated untiled peak for the output size and tiles otherwise; # "tiled"/"untiled" force a path. WAN22_VAE_DECODE_TILING overrides at runtime. @@ -94,6 +99,10 @@ def __post_init__(self): # rather than read as on. if not isinstance(self.vae_decode_fast, bool): raise ValueError(f"Wan2.2 vae_decode_fast must be true or false, got {self.vae_decode_fast!r}.") + if self.attention_backend not in ("flashinfer", "sdpa"): + raise ValueError( + f"Wan2.2 attention_backend must be 'flashinfer' or 'sdpa', got {self.attention_backend!r}." + ) self.fast_decode_sizes() def fast_decode_sizes(self) -> frozenset[tuple[int, int]]: diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 977d1791e0..505215edbf 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -1,4 +1,4 @@ -"""Wan2.2-TI2V-5B node submodules (all four declare no engine resources). +"""Wan2.2-TI2V-5B node submodules (nothing is cached; the dit declares a ragged attention resource). text_encoder -> Wan22TextEncoderSubmodule (UMT5-XXL, thin HF wrapper) vae_encoder -> Wan22VaeEncoderSubmodule (Wan2.2-VAE encode, I2V only) @@ -20,13 +20,17 @@ import logging import os from contextlib import contextmanager +from typing import NamedTuple import torch from torch import nn from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep +from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule +from mstar.model.wan22.components.dit import Wan22AttentionFns from mstar.model.wan22.components.unipc import ( SOLVER_ORDER, UniPCState, @@ -324,12 +328,21 @@ def forward( # dit # --------------------------------------------------------------------------- +class Wan22AttnLayout(NamedTuple): + """One step's self-attention spans for a request: each of the ``branches`` + guidance branches this rank computes attends within its video tokens.""" + + video_tokens: int + branches: int + + class Wan22DitSubmodule(_SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule): """Dense 5B video DiT with the UniPC step run inline. One loop iteration is one batch-2 transformer forward (the positive and - negative prompt embeddings stacked on the batch axis) plus one UniPC - predictor/corrector update from ``components/unipc.py``. + negative prompt embeddings stacked on the batch axis; on a two-rank sp + instance, one branch per rank) plus one UniPC predictor/corrector update + from ``components/unipc.py``. Loop-carried edges, all float32 over the latent grid: @@ -355,6 +368,10 @@ class Wan22DitSubmodule(_SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule): # circular import. _I2V_WALK = "video_gen_i2v" + # The ragged self-attention resource (Wan22Model.get_node_resources) and its span. + SELF_ATTN = "dit_self_attn" + VIDEO_SPAN = "video" + # The WHOLE forward cannot be one compiled region: Inductor would fuse the UniPC # step's `sample - sigma * output` and pass the 0-dim CPU-resident sigma as a # device pointer, which fails on the first call. Moving sigma to CUDA would fix @@ -364,10 +381,22 @@ class Wan22DitSubmodule(_SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule): # leaving the solver, CFG combine, and everything touching sigma eager. disable_torch_compile = True - def __init__(self, transformer: nn.Module, config: Wan22Config): + def __init__(self, transformer: nn.Module, config: Wan22Config, sp_group=None): super().__init__() self.transformer = transformer self.config = config + # Two ranks split the CFG batch: each runs one branch at batch 1 and the + # predictions are all-gathered, so the step costs one branch's forward plus a + # latent-sized gather. Every rank then runs the same combine and UniPC step on + # replicated latents. Any other shape (guidance off, more ranks) computes the + # whole batch on every rank, which is correct but no faster. + self.sp_group = sp_group if sp_group is not None and sp_group.world_size > 1 else None + if self.sp_group is not None and self.sp_group.world_size != 2: + logger.warning( + "Wan2.2 dit: sp_size=%d; only 2 ranks split the CFG batch, so every rank " + "computes the whole step.", self.sp_group.world_size, + ) + self._attn_fns: Wan22AttentionFns | None = None self._record_fp32_islands() # Compile the inner DiT region only (tensors in, tensor out; nothing # CPU-resident crosses it), IN PLACE so ``self.transformer`` keeps its @@ -378,6 +407,9 @@ def __init__(self, transformer: nn.Module, config: Wan22Config): self._compile_dit = bool(config.compile_dit) self._compiled_shapes: set[tuple[int, ...]] = set() if self._compile_dit and transformer is not None: + # Serving numerics, like the compile itself: the reference gates build + # with compile_dit off and keep the exact GELU path. + transformer.set_fuse_ffn_gelu(True) transformer.forward = torch.compile( transformer.forward, fullgraph=False, dynamic=False, ) @@ -435,7 +467,58 @@ def prepare_inputs( tensor_inputs["time_index"] = time_index tensor_inputs["unipc_model_outputs"] = inputs["unipc_model_outputs"][0] tensor_inputs["unipc_last_sample"] = inputs["unipc_last_sample"][0] - return NodeInputs(tensor_inputs=tensor_inputs) + _, _, t_lat, h_lat, w_lat = tensor_inputs["latents"].shape + p_t, p_h, p_w = self.config.patch_size + guidance_scale = float(fwd_info.step_metadata.get("guidance_scale", self.config.guidance_scale)) + layout = Wan22AttnLayout( + video_tokens=(t_lat // p_t) * (h_lat // p_h) * (w_lat // p_w), + branches=self._branches(guidance_scale), + ) + return NodeInputs( + tensor_inputs=tensor_inputs, resource_step_info=layout, input_seq_len=layout.video_tokens, + ) + + def _cfg_split(self, guidance_scale: float) -> bool: + """Whether this step's two guidance branches run one per rank. + + ``guidance_scale`` is request metadata, identical on every rank of the + instance, so all ranks agree on whether the forward's gather runs.""" + return guidance_scale > 1.0 and self.sp_group is not None and self.sp_group.world_size == 2 + + def _branches(self, guidance_scale: float) -> int: + """Guidance branches this rank's forward computes.""" + return 2 if guidance_scale > 1.0 and not self._cfg_split(guidance_scale) else 1 + + # attention ------------------------------------------------------------ + def bind_node_resources(self, resources) -> None: + super().bind_node_resources(resources) + self._attn_fns = None # the memoized callables hold the old resources + + def _attention_fns(self) -> Wan22AttentionFns | None: + """The ragged callables, built once per resource binding (the compiled + transformer guards on their identity), or None to run SDPA.""" + if self._attn_fns is None: + self_attn = self.node_resources.get(self.SELF_ATTN) + if self_attn is None: + return None + self._attn_fns = Wan22AttentionFns( + self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN, opaque=True), cross_attn=None, + ) + logger.info("Wan2.2 DiT: self-attention through the engine's ragged attention resource") + return self._attn_fns + + def declare_step( + self, graph_walk: str, request_ids: list[str], inputs: list[NodeInputs], + slot_lease: SlotLease | None = None, piecewise_leases=None, **kwargs, + ) -> SubmoduleStep | None: + if self.config.attention_backend != "flashinfer": + return None + video = [ + Segment(request_id=rid, label=self.VIDEO_SPAN, span=inp.resource_step_info.video_tokens) + for rid, inp in zip(request_ids, inputs, strict=True) + for _ in range(inp.resource_step_info.branches) + ] + return SubmoduleStep(steps={self.SELF_ATTN: AttentionStep(segments=tuple(video), causal=False)}) def forward( self, @@ -469,9 +552,10 @@ def forward( model_input = latents model_input = model_input.to(self.transformer.dtype) - # expand_timesteps: per-token timestep over the post-patch grid - # ([:, ::2, ::2] == patch_size (1, 2, 2)); I2V frame-0 tokens get 0. - temp_ts = (mask[0, 0][:, ::2, ::2] * t).flatten() + # expand_timesteps: the per-token timestep over the post-patch grid is + # constant within a latent frame (I2V frame-0 tokens get 0), so the + # DiT takes it once per frame and broadcasts the modulation. + temp_ts = mask[0, 0, :, 0, 0] * t noise_pred = self._noise_prediction( model_input, temp_ts, text_embeds_pos, text_embeds_neg, guidance_scale ) @@ -521,7 +605,8 @@ def _noise_prediction( reference's two sequential forwards. The combine runs in the transformer's output dtype (bf16), like the reference.""" do_cfg = guidance_scale > 1.0 - batch = 2 if do_cfg else 1 + cfg_split = self._cfg_split(guidance_scale) + batch = self._branches(guidance_scale) hidden_states = model_input.repeat(batch, 1, 1, 1, 1) timestep = temp_ts.unsqueeze(0).expand(batch, -1) # dynamic=False keys the compiled trace on the full input shape (batch @@ -537,16 +622,23 @@ def _noise_prediction( "this resolution now (one-time pause); later requests reuse it.", shape_key, ) - encoder_hidden_states = ( - torch.cat([text_embeds_pos, text_embeds_neg], dim=0) if do_cfg else text_embeds_pos - ).to(self.transformer.dtype) + if cfg_split: + # rank 0 the conditional branch, rank 1 the unconditional one + encoder_hidden_states = text_embeds_pos if self.sp_group.rank == 0 else text_embeds_neg + elif do_cfg: + encoder_hidden_states = torch.cat([text_embeds_pos, text_embeds_neg], dim=0) + else: + encoder_hidden_states = text_embeds_pos out = self.transformer( hidden_states=hidden_states, timestep=timestep, - encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states=encoder_hidden_states.to(self.transformer.dtype), + attn_fns=self._attention_fns(), ) if not do_cfg: return out + if cfg_split: + out = self.sp_group.all_gather(out.contiguous(), dim=0) noise_cond, noise_uncond = out[0:1], out[1:2] return noise_uncond + guidance_scale * (noise_cond - noise_uncond) diff --git a/mstar/model/wan22/wan22_model.py b/mstar/model/wan22/wan22_model.py index 38084ca16f..58616b981c 100644 --- a/mstar/model/wan22/wan22_model.py +++ b/mstar/model/wan22/wan22_model.py @@ -1,6 +1,6 @@ """Wan22Model: Wan2.2-TI2V-5B text/image-to-video generation. -Architecture (4 nodes, no engine resources — nothing here is cached): +Architecture (4 nodes; nothing here is cached, the dit's self-attention is a ragged resource): text_encoder - UMT5-XXL prompt encoder (thin HF wrapper) vae_encoder - Wan2.2-VAE encode of the conditioning frame (I2V) dit - dense 5B video DiT; CFG as one batch-2 forward, @@ -39,7 +39,9 @@ CurrentForwardConductorMetadata, StreamingConnectionState, ) +from mstar.distributed.base import ShardingConfig from mstar.engine.resources import NodeResourceSpec +from mstar.engine.resources.attn.ragged.config import RaggedAttentionConfig, RaggedAttentionSpec from mstar.graph.base import ( GraphEdge, GraphNode, @@ -84,6 +86,7 @@ def __init__( skip_weight_loading: bool = False, vae_decode_fast: bool = True, vae_decode_fast_sizes: tuple[str, ...] | None = None, + attention_backend: str = "flashinfer", **kwargs, ): if variant != WAN22_VARIANT_TI2V_5B: @@ -95,6 +98,7 @@ def __init__( self.cache_dir = cache_dir self.config = Wan22Config( variant=variant, vae_decode_fast=vae_decode_fast, vae_decode_fast_sizes=vae_decode_fast_sizes, + attention_backend=attention_backend, ) # Dummy mode: get_submodule returns None for every node, so engines run # without weights, GPU or network. @@ -115,7 +119,8 @@ def __init__( # ------------------------------------------------------------------ def get_node_resources(self) -> list[NodeResourceSpec]: - """No engine-built resources: every node here is stateless. + """The DiT's self-attention as a cacheless ragged resource (with + ``attention_backend: flashinfer``); nothing here keeps state. Nothing in this model reads K/V back across steps. The text encoder and both VAE passes are one-shot forwards, and the DiT recomputes @@ -123,7 +128,36 @@ def get_node_resources(self) -> list[NodeResourceSpec]: the only state that crosses a step is the UniPC solver's, which rides the loop-back edges as ordinary tensors rather than a cache. """ - return [] + if self.config.attention_backend != "flashinfer": + return [] + cfg = self.config + # Self-attention only: the cross-attention over the 512 text tokens runs + # SDPA (cuDNN), which measured faster than FlashInfer's ragged kernel at + # that shape (0.46 vs 0.59 ms per layer at 1280x704x81 on an H200), while + # FA3 beats it on the self-attention (12.4 vs 14.8 ms). + return [ + RaggedAttentionSpec( + resource_key=Wan22DitSubmodule.SELF_ATTN, nodes={"dit"}, + config=RaggedAttentionConfig( + num_qo_heads=cfg.num_attention_heads, + num_kv_heads=cfg.num_attention_heads, + head_dim=cfg.attention_head_dim, + # one segment per guidance branch + max_segments_per_request=2, + # the dit's sp ranks split guidance branches, not heads + shard_heads=False, + dtype=torch.bfloat16, + ), + ), + ] + + def get_default_sharding_config(self) -> ShardingConfig: + # The dit accepts an sp group: with guidance on, each of two ranks runs one + # CFG branch and they all-gather the two noise predictions (see + # Wan22DitSubmodule). Signals stay replicated, so no shard_dim entries. + return ShardingConfig( + groups=[], tp_enabled_nodes=set(), shard_dim={}, sp_enabled_nodes={"dit"}, + ) def get_graph_walk_graphs(self) -> dict[str, GraphSection]: # -- encode_text: UMT5 forward; both CFG embeddings persist at the @@ -432,7 +466,10 @@ def postprocess(self, output: torch.Tensor, modality: str, request_kwargs: dict if fps <= 0: raise ValueError(f"Wan2.2 fps must be positive; got {raw_fps!r}") - frames = output[0].permute(1, 2, 3, 0).cpu().numpy() + # One multithreaded torch copy into frame-major layout: handed the + # permuted view, PyAV copies each frame out of it single-threaded + # (~290 ms of a ~600 ms mux for 1280x704x81). + frames = output[0].permute(1, 2, 3, 0).contiguous().cpu().numpy() buffer = io.BytesIO() container = av.open(buffer, mode="w", format="mp4") # PyAV rejects float rates, so the rate is handed over as an exact @@ -648,18 +685,17 @@ def get_submodule( self, node_name: str, device: str = "cpu", tp_group=None, autocast_dtype: torch.dtype | None = None, sp_group=None, ) -> torch.nn.Module | None: - # ``autocast_dtype`` and ``sp_group`` exist for interface parity only. - # Weights load in the checkpoint's own dtypes, and wan22 declares no - # sp-enabled nodes. + # ``autocast_dtype`` exists for interface parity only: weights load in the + # checkpoint's own dtypes. ``sp_group`` reaches the dit only. if node_name in self._submodule_cache: return self._submodule_cache[node_name] - submodule = self._create_submodule(node_name, device) + submodule = self._create_submodule(node_name, device, sp_group=sp_group) self._submodule_cache[node_name] = submodule if submodule is not None: logger.info("Loaded Wan2.2 submodule for node %s", node_name) return submodule - def _create_submodule(self, node_name: str, device: str = "cpu") -> NodeSubmodule | None: + def _create_submodule(self, node_name: str, device: str = "cpu", sp_group=None) -> NodeSubmodule | None: """Construct one node's submodule from the checkpoint. The wrapped components load on CPU and the EngineManager moves them; the @@ -696,7 +732,7 @@ def _create_submodule(self, node_name: str, device: str = "cpu") -> NodeSubmodul transformer = build_wan22_dit( self.config, self.model_path_hf, device=device, cache_dir=self.cache_dir, ) - return Wan22DitSubmodule(transformer, self.config) + return Wan22DitSubmodule(transformer, self.config, sp_group=sp_group) logger.warning("Wan2.2 has no submodule for node %r; running it dummy.", node_name) return None diff --git a/mstar/model/z_image/submodules.py b/mstar/model/z_image/submodules.py index 2b6b2b1bf3..e4226e9131 100644 --- a/mstar/model/z_image/submodules.py +++ b/mstar/model/z_image/submodules.py @@ -47,6 +47,7 @@ logger = logging.getLogger(__name__) TEXT_INPUTS, TEXT_EMBEDS, IMAGE_OUTPUT = "text_inputs", "text_embeds", "image_output" +TEXT_MASK = "text_mask" CAP_PAD_MASK = "cap_pad_mask" @@ -92,19 +93,49 @@ class ZImageTextEncoderSubmodule(NodeSubmodule): encoder's ``max_sequence_length`` (512), as the reference pipeline pads them: the real tokens' values do not depend on the padding mathematically, but the attention kernel's accumulation order does, so the padded length is part of bit-exactness. Requests of any - length batch together.""" + length batch together. + + Rows are padded on the device in ``prepare_inputs``, and the batched forward emits the + whole ``[B, 512, 2560]`` batch, cut per request in ``unpack_packed_outputs``. The forward + then has fixed shapes and no host syncs, so it is captured per batch size + (``capture_batch_sizes``). A batch with a prompt that fills all 512 tokens runs the + encoder's unmasked path eagerly.""" disable_torch_compile = True + # the batched forward's whole output, cut per request after the forward + PACKED_EMBEDS = "__text_embeds__" - def __init__(self, encoder: nn.Module, config: ZImageConfig, max_batch_size: int = 16): + def __init__( + self, encoder: nn.Module, config: ZImageConfig, max_batch_size: int = 16, + capture_walk: str | None = None, capture_batch_sizes: Sequence[int] = (), + ): super().__init__() self.encoder = encoder self.config = config self._max_batch_size = max_batch_size + self._capture_walk = capture_walk + self._capture_batch_sizes = [b for b in capture_batch_sizes if b <= max_batch_size] + + def _pad_row(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``ids [L]`` -> right-padded ``(ids, mask)``, each ``[max_sequence_length]`` int64, + built where ``ids`` lives (device ops only, no host round trip).""" + width = self.config.text_encoder.max_sequence_length + n = int(ids.shape[0]) + padded = torch.full((width,), self.config.text_encoder.pad_token_id, dtype=torch.long, device=ids.device) + padded[:n].copy_(ids) + mask = torch.zeros(width, dtype=torch.long, device=ids.device) + mask[:n] = 1 + return padded, mask def prepare_inputs(self, graph_walk, fwd_info, inputs: NameToTensorList, **kwargs) -> NodeInputs: ids = inputs[TEXT_INPUTS][0] - return NodeInputs(tensor_inputs={TEXT_INPUTS: ids}, input_seq_len=int(ids.shape[0])) + padded, mask = self._pad_row(ids) + return NodeInputs( + tensor_inputs={TEXT_INPUTS: padded, TEXT_MASK: mask}, + input_seq_len=self.config.text_encoder.max_sequence_length, + # the real token count, from the shape: where this row's output is cut + resource_step_info=int(ids.shape[0]), + ) def can_batch(self, batch, model_inputs) -> bool: return len(model_inputs) > 1 @@ -112,35 +143,67 @@ def can_batch(self, batch, model_inputs) -> bool: def max_batch_size(self, graph_walk: str): return self._max_batch_size - def preprocess(self, graph_walk, engine_inputs, inputs: list[NodeInputs]) -> dict: - lengths = [int(inp.tensor_inputs[TEXT_INPUTS].shape[0]) for inp in inputs] + def cg_key_info(self, graph_walk: str, per_request_info, per_request_input_metadata=None, **kwargs): + """True when every real row is a padded prompt (``text_len`` in the step metadata): the + masked path, the only captured one. Padding rows of a captured bucket carry no metadata.""" + del graph_walk, per_request_input_metadata, kwargs width = self.config.text_encoder.max_sequence_length - pad_id = self.config.text_encoder.pad_token_id - ids = torch.full((len(inputs), width), pad_id, dtype=torch.long) - mask = torch.zeros((len(inputs), width), dtype=torch.long) - for row, inp in enumerate(inputs): - n = lengths[row] - ids[row, :n] = inp.tensor_inputs[TEXT_INPUTS].to(torch.long) - mask[row, :n] = 1 - return {TEXT_INPUTS: ids, "text_mask": mask, "lengths": lengths} - - def _encode(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + lengths = [info.step_metadata.get("text_len") for info in per_request_info.values() if info.step_metadata] + return bool(lengths) and all(n is not None and n < width for n in lengths) + + def preprocess(self, graph_walk, engine_inputs, inputs: list[NodeInputs]) -> dict: + lengths = [inp.resource_step_info for inp in inputs] + return { + TEXT_INPUTS: torch.stack([inp.tensor_inputs[TEXT_INPUTS] for inp in inputs]), + TEXT_MASK: torch.stack([inp.tensor_inputs[TEXT_MASK] for inp in inputs]), + "lengths": lengths, + # any padded row puts the batch on the masked path, as the reference batches it + "padded": any(n < self.config.text_encoder.max_sequence_length for n in lengths), + } + + def _encode(self, ids: torch.Tensor, mask: torch.Tensor, padded: bool | None = None) -> torch.Tensor: device = self.get_device() - hidden = self.encoder(ids.to(device), mask.to(device)).to(self.encoder.dtype) + mask = mask.to(device) + hidden = self.encoder(ids.to(device), mask, padded=padded).to(self.encoder.dtype) # zero the pad rows: the DiT replaces them, and the parity tests compare real rows only - return hidden * mask.to(device=device, dtype=hidden.dtype)[..., None] + return hidden * mask.to(dtype=hidden.dtype)[..., None] - def forward(self, graph_walk, engine_inputs, text_inputs, text_mask, lengths, **kwargs): - hidden = self._encode(text_inputs, text_mask) + def forward(self, graph_walk, engine_inputs, text_inputs, text_mask, lengths, padded=None, **kwargs): + hidden = self._encode(text_inputs, text_mask, padded) return {TEXT_EMBEDS: [hidden[:, : padded_length(lengths[0])]]} - def forward_batched(self, graph_walk, engine_inputs, text_inputs, text_mask, lengths, **kwargs): - hidden = self._encode(text_inputs, text_mask) + def forward_batched(self, graph_walk, engine_inputs, text_inputs, text_mask, padded=None, **kwargs): + return {self.PACKED_EMBEDS: self._encode(text_inputs, text_mask, padded)} + + def unpack_packed_outputs(self, static_output, request_ids, real_seq_lens, inputs, per_request_info): + hidden = static_output.get(self.PACKED_EMBEDS) + if hidden is None: + return {} + # cloned: a captured forward's output buffer is overwritten by the next replay return { - rid: {TEXT_EMBEDS: [hidden[i : i + 1, : padded_length(lengths[i])]]} - for i, rid in enumerate(engine_inputs.request_ids) + rid: {TEXT_EMBEDS: [hidden[i : i + 1, : padded_length(inp.resource_step_info)].clone()]} + for i, (rid, inp) in enumerate(zip(request_ids, inputs, strict=True)) } + def get_accelerator_graph_configs(self, device: torch.device, tp_world_size: int = 1): + if self._capture_walk is None or not self._capture_batch_sizes: + return [] + from mstar.engine.accelerator_graph_config import BatchedAcceleratorGraphConfig + + # a padding row: a one-token prompt, so every query row attends to something + ids, mask = self._pad_row(torch.zeros(1, dtype=torch.long, device=device)) + single = NodeInputs(tensor_inputs={TEXT_INPUTS: ids, TEXT_MASK: mask}, + input_seq_len=int(ids.shape[0]), resource_step_info=1) + return [BatchedAcceleratorGraphConfig( + capture_graph_walk=self._capture_walk, + single_request_inputs=single, + additional_key_info=True, + compile=False, + capture_batch_sizes=list(self._capture_batch_sizes), + caps_eager_batch_size=False, + input_seq_dims={TEXT_INPUTS: 0, TEXT_MASK: 0}, + )] + # --------------------------------------------------------------------------- # dit diff --git a/mstar/model/z_image/z_image_model.py b/mstar/model/z_image/z_image_model.py index 47e2d736d1..5d6c44212b 100644 --- a/mstar/model/z_image/z_image_model.py +++ b/mstar/model/z_image/z_image_model.py @@ -299,8 +299,12 @@ def _create_submodule(self, node_name: str, device) -> NodeSubmodule | None: from mstar.model.z_image.weight_loader import build_text_encoder, build_transformer, build_vae if node_name == "text_encoder": - return ZImageTextEncoderSubmodule(build_text_encoder(self.config, self.snapshot, device), self.config, - max_batch_size=self.max_batch_size) + return ZImageTextEncoderSubmodule( + build_text_encoder(self.config, self.snapshot, device), self.config, + max_batch_size=self.max_batch_size, + capture_walk=ENCODE_TEXT_WALK if self.accelerator_graph else None, + capture_batch_sizes=self.capture_batch_sizes, + ) if node_name == "vae_decoder": return ZImageVaeDecoderSubmodule( build_vae(self.config, self.snapshot, device), self.config, max_batch_size=self.max_batch_size, diff --git a/mstar/utils/fused_moe/kernels.py b/mstar/utils/fused_moe/kernels.py index b1d67fa566..50b8f7d0a5 100644 --- a/mstar/utils/fused_moe/kernels.py +++ b/mstar/utils/fused_moe/kernels.py @@ -25,6 +25,7 @@ from __future__ import annotations +import functools from typing import Any, Dict import torch @@ -619,6 +620,52 @@ def moe_sum_reduce_triton( # --------------------------------------------------------------------------- +def _tile(bn: int, bk: int, warps: int, stages: int) -> Dict[str, int]: + return {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": 1, + "num_warps": warps, "num_stages": stages} + + +# (E, N of the gate/up GEMM, K = hidden) -> rows of (max tokens, gate/up tile, down tile), +# swept on an H200 with uniform random routing for Qwen3-Omni's Thinker +# (128 experts, 2048 hidden, 768 inter, top-8) and Talker (1024 / 384, top-6). +# Above the last row the default applies; at those batch sizes the GEMMs read most +# experts and the default is within a few percent. +_TUNED: Dict[tuple, tuple] = { + (128, 1536, 2048): ( + (1, _tile(32, 128, 4, 4), _tile(64, 64, 4, 3)), + (2, _tile(64, 64, 4, 4), _tile(64, 64, 4, 4)), + (4, _tile(64, 64, 4, 3), _tile(64, 64, 4, 3)), + (8, _tile(64, 128, 4, 4), _tile(64, 128, 4, 3)), + ), + (128, 768, 1024): ( + (1, _tile(64, 64, 4, 5), _tile(64, 64, 8, 3)), + (2, _tile(64, 64, 4, 3), _tile(32, 64, 4, 3)), + (4, _tile(64, 64, 4, 4), _tile(64, 64, 4, 3)), + ), +} + + +@functools.cache +def _tuned_tiles_apply() -> bool: + """Only on the NVIDIA Hopper parts the table was swept on; every other GPU, + ROCm included, keeps the generic heuristics.""" + from mstar.platforms import current_platform + + return current_platform.is_cuda() and current_platform.cuda_capability() == (9, 0) + + +def get_tuned_configs(M: int, E: int, N: int, K: int) -> tuple[Dict[str, int], Dict[str, int]] | None: + """Swept ``(gate/up, down)`` tiles for a known shape and batch size, else None. + Both share ``BLOCK_SIZE_M``: the token alignment is computed once for the two.""" + rows = _TUNED.get((E, N, K)) + if not rows or not _tuned_tiles_apply(): + return None + for max_m, gate_up, down in rows: + if M <= max_m: + return dict(gate_up), dict(down) + return None + + def get_default_config(M: int, E: int, N: int, K: int, top_k: int) -> Dict[str, int]: """Pick Triton tile sizes based on problem shape. diff --git a/mstar/utils/fused_moe/router.py b/mstar/utils/fused_moe/router.py new file mode 100644 index 0000000000..b866614800 --- /dev/null +++ b/mstar/utils/fused_moe/router.py @@ -0,0 +1,64 @@ +"""Softmax top-k routing in one kernel. + +``softmax -> topk -> renormalize`` as separate ops is five launches per MoE layer +(``torch.topk`` alone is two, a radix select and a sort), which at decode batch sizes +costs more than the router GEMM. One program per token does it in registers. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _softmax_topk_kernel(logits, w_out, id_out, stride_t, + E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, + NORM: tl.constexpr): + """fp32 softmax over the E logits, the K largest probabilities in descending order, + optionally divided by their sum. Equal probabilities (common: the logits are bf16) go + to the lower expert first; ``torch.topk`` orders such ties arbitrarily, so the two can + pick different experts at a tie, both equally valid.""" + t = tl.program_id(0) + e = tl.arange(0, EP) + valid = e < E + x = tl.load(logits + t * stride_t + e, mask=valid, other=float("-inf")).to(tl.float32) + x = x - tl.max(x, 0) + num = tl.where(valid, tl.exp(x), 0.0) + prob = num / tl.sum(num, 0) + cand = tl.where(valid, prob, -1.0) + kk = tl.arange(0, KP) + ids = tl.zeros([KP], tl.int32) + ws = tl.zeros([KP], tl.float32) + for j in tl.static_range(K): + idx = tl.argmax(cand, 0) + ws = tl.where(kk == j, tl.max(cand, 0), ws) + ids = tl.where(kk == j, idx, ids) + cand = tl.where(e == idx, -1.0, cand) + if NORM: + ws = ws / tl.sum(tl.where(kk < K, ws, 0.0), 0) + tl.store(w_out + t * K + kk, ws, mask=kk < K) + tl.store(id_out + t * K + kk, ids.to(tl.int64), mask=kk < K) + + +@torch.library.custom_op("mstar::moe_softmax_topk", mutates_args=()) +def softmax_topk(logits: torch.Tensor, top_k: int, normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: + """Top-k of ``softmax(logits (T, E))`` as ``(fp32 weights, int64 expert ids)``, the + weights renormalized to sum to one when ``normalize``. Opaque to dynamo, so a compiled + router keeps it whole.""" + T, E = logits.shape + assert logits.stride(1) == 1 + topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) + topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) + if T: + _softmax_topk_kernel[(T,)]( + logits, topk_w, topk_ids, logits.stride(0), E=E, EP=triton.next_power_of_2(E), + K=top_k, KP=triton.next_power_of_2(top_k), NORM=normalize, num_warps=1) + return topk_w, topk_ids + + +@softmax_topk.register_fake +def _(logits, top_k, normalize): + T = logits.shape[0] + return (logits.new_empty(T, top_k, dtype=torch.float32), + logits.new_empty(T, top_k, dtype=torch.int64)) diff --git a/mstar/utils/fused_moe/runner.py b/mstar/utils/fused_moe/runner.py index c223cbf461..30ad7c1a7a 100644 --- a/mstar/utils/fused_moe/runner.py +++ b/mstar/utils/fused_moe/runner.py @@ -16,6 +16,7 @@ from mstar.utils.fused_moe.kernels import ( act_and_mul_triton, get_default_config, + get_tuned_configs, invoke_fused_moe_kernel, invoke_fused_moe_kernel_fp8_w8a8, moe_sum_reduce_triton, @@ -99,6 +100,10 @@ def fused_experts( topk_weights = topk_weights.contiguous() config = get_default_config(M=num_tokens, E=E, N=two_inter, K=hidden, top_k=top_k) + config_down = config + tuned = get_tuned_configs(num_tokens, E, two_inter, hidden) + if tuned is not None: + config, config_down = tuned compute_type = _tl_compute_type(hidden_states.dtype) # 1. Token permute + per-expert block alignment. @@ -156,7 +161,7 @@ def fused_experts( num_tokens_post_padded=num_tokens_post_padded, mul_routed_weight=True, top_k=1, - config=config, + config=config_down, compute_type=compute_type, ) diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 705375f309..bc3c2f2daf 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -28,14 +28,6 @@ def sampler_slots() -> int: return int(os.environ.get("MSTAR_SAMPLER_SLOTS", "4096")) - -def ingraph_decode_rope() -> bool: - """``MSTAR_INGRAPH_DECODE_ROPE`` (default 0): a captured decode step - builds its rotary tables inside the graph from the position resource's - planned positions instead of staging them before the replay.""" - return os.environ.get("MSTAR_INGRAPH_DECODE_ROPE", "0") == "1" - - def kv_chain_lazy_steps() -> int: """``MSTAR_KV_CHAIN_LAZY_STEPS`` (default 16): how many decode steps of sampled tokens the KV manager batches before extending the requests' diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 70b0c96e27..93bc416c6f 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -3284,6 +3284,12 @@ def _pp_stage(name: str) -> None: range_pop(synchronize=False) range_push("worker.postprocess.register_outputs", synchronize=False) self._register_outputs(route_output) + # Registration, just above, is the last reader of the step's host + # copies (the sends below run after it); an output held for a later + # send (a loop's accumulated outputs) must not read them, since the + # pinned buffers behind them are this step's and the next one + # rewrites them. + self.tensor_manager.release_host_copies(stored.flat_uuids) _pp_stage("register_outputs") # send outputs diff --git a/test/chatterbox/test_t3_components.py b/test/chatterbox/test_t3_components.py index 0bfae4022f..5aef26e75f 100644 --- a/test/chatterbox/test_t3_components.py +++ b/test/chatterbox/test_t3_components.py @@ -224,11 +224,10 @@ def test_backbone_prefill_and_decode_match_reference(pair: Pair): # three greedy decode steps, batched on our side, cached on both sides next_tokens = [logit.argmax() for logit in ref_logits] for step in range(3): - speech_pos = torch.full((1,), step + 1, dtype=torch.long) step_rows, ref_step_logits = [], [] for i in range(2): tok = next_tokens[i].view(1) - emb = pair.model.embed_speech(tok, speech_pos) + emb = pair.model.embed_speech(tok, step + 1) step_rows.append(emb) if t3.learned_pos_emb: ref_emb = pair.ref.speech_emb(tok[None]) + pair.ref.speech_pos_emb.get_fixed_embedding(step + 1) diff --git a/test/integration/test_qwen3_omni_code2wav_real_weights.py b/test/integration/test_qwen3_omni_code2wav_real_weights.py index 4f753d472e..ee03e25556 100644 --- a/test/integration/test_qwen3_omni_code2wav_real_weights.py +++ b/test/integration/test_qwen3_omni_code2wav_real_weights.py @@ -2,14 +2,16 @@ Streams real Talker codes (six TTS replies served by M*, in ``data/qwen3_omni_talker_codes.safetensors``) through the edge's -``LeftContextChunkPolicy`` and ``Code2WavSubmodule``, and compares the -concatenated audio with one decode of each whole utterance. +``ScheduledLeftContextChunkPolicy`` and ``Code2WavSubmodule``, and compares +the concatenated audio with one decode of each whole utterance. Streaming can't match the full decode exactly: 25 frames of left context cut into the vocoder's 72-frame attention window (24-28 dB SNR over an utterance). What it must do is produce the full decode's length, chunk by chunk, and stay about as close to it at chunk boundaries as mid-chunk. A boundary that drops -or repeats samples misaligns the rest of the stream. +or repeats samples misaligns the rest of the stream. The served schedule's +small first chunk moves every later boundary, so its boundaries are also +checked against the plain 25-frame schedule's. The module never downloads weights: it skips unless CUDA is available and the Qwen3-Omni checkpoint is in a local Hugging Face cache. Only the Code2Wav @@ -78,7 +80,7 @@ def _find_cached_snapshot() -> Path | None: from mstar.model.qwen3_omni.submodules import Code2WavSubmodule # noqa: E402 from mstar.model.submodule_base import InputMetadata, ModelInputsFromEngine # noqa: E402 from mstar.model.utils import ModuleAndPrefix, load_weights_from_hf_shards # noqa: E402 -from mstar.streaming.chunk_policy import LeftContextChunkPolicy # noqa: E402 +from mstar.streaming.chunk_policy import ScheduledLeftContextChunkPolicy # noqa: E402 from mstar.streaming.stream_buffer import StreamBuffer # noqa: E402 @@ -102,13 +104,13 @@ def code2wav(): torch.backends.cudnn.benchmark = benchmark -def stream(submodule: Code2WavSubmodule, frames: torch.Tensor) -> list[torch.Tensor]: +def stream(submodule: Code2WavSubmodule, frames: torch.Tensor, schedule: tuple[int, ...]) -> list[torch.Tensor]: """Per-step audio for ``frames`` ``[T, 16]`` streamed through the edge's - policy and the submodule.""" + policy, with ``schedule`` ahead of the steady chunks, and the submodule.""" cfg = submodule.config.code2wav buffer = StreamBuffer( request_id="r", edge_name="codec_tokens", from_partition="Talker", - policy=LeftContextChunkPolicy(chunk=cfg.codec_chunk_frames, left_context=cfg.codec_left_context_frames), + policy=ScheduledLeftContextChunkPolicy(schedule, cfg.codec_chunk_frames, cfg.codec_left_context_frames), ) for i, frame in enumerate(frames): buffer.pre_read_register(f"t{i}") @@ -147,22 +149,36 @@ def window_snr_db(audio: torch.Tensor, reference: torch.Tensor, center: int) -> return 10 * math.log10(signal / max(error, 1e-9)) -def boundary_and_mid_snr(audio: torch.Tensor, reference: torch.Tensor, cfg: Code2WavConfig, upsample: int, tail: int): +def boundary_and_mid_snr(audio: torch.Tensor, reference: torch.Tensor, chunk_samples: list[int]): """SNR around every chunk boundary of the stream, and halfway into every later chunk long enough to have a middle.""" boundaries, middles = [], [] - chunk = cfg.codec_chunk_frames * upsample end = min(audio.numel(), reference.numel()) - HALF_WINDOW - boundary = chunk - tail - while boundary < end: + boundary = chunk_samples[0] + for samples in chunk_samples[1:]: + if boundary >= end: + break boundaries.append(window_snr_db(audio, reference, boundary)) - if boundary + chunk // 2 < end: - middles.append(window_snr_db(audio, reference, boundary + chunk // 2)) - boundary += chunk + if boundary + samples // 2 < end: + middles.append(window_snr_db(audio, reference, boundary + samples // 2)) + boundary += samples return boundaries, middles -def test_real_codes_stream_to_the_full_decode_length_without_boundary_dips(code2wav): +def expected_chunk_samples(num_frames: int, schedule: tuple[int, ...], chunk: int, upsample: int, tail: int): + """Samples per non-empty chunk: the first ends ``tail`` short, every later + one carries its new frames' audio (the tail of the one before included).""" + sizes, delivered = [], 0 + while delivered < num_frames: + new = min(schedule[len(sizes)] if len(sizes) < len(schedule) else chunk, num_frames - delivered) + sizes.append(new * upsample - (tail if not sizes else 0)) + delivered += new + return sizes + + +def stream_all(code2wav, schedule: tuple[int, ...]) -> tuple[list[float], list[float]]: + """Every utterance streamed with ``schedule``: checks the chunk lengths and + returns the boundary and mid-chunk SNRs against the full decode.""" model, eager_forward, config = code2wav submodule = Code2WavSubmodule(model, Qwen3OmniModelConfig(code2wav=config)) upsample = submodule.total_upsample @@ -175,25 +191,25 @@ def test_real_codes_stream_to_the_full_decode_length_without_boundary_dips(code2 reference = (full.clamp(-1, 1) * 32767).to(torch.int16).cpu() tail = num_frames * upsample - reference.numel() - chunks = stream(submodule, frames) + chunks = stream(submodule, frames, schedule) audio = torch.cat(chunks) - boundaries, middles = boundary_and_mid_snr(audio, reference, config, upsample, tail) + sizes = [c.numel() for c in chunks if c.numel()] + boundaries, middles = boundary_and_mid_snr(audio, reference, sizes) print( - f"{name}: samples per chunk {[c.numel() for c in chunks]}, total {audio.numel()} " + f"{name} schedule {schedule}: samples per chunk {sizes}, total {audio.numel()} " f"(full decode {reference.numel()}), boundary SNR {[round(s, 1) for s in boundaries]} dB, " f"mid-chunk SNR {[round(s, 1) for s in middles]} dB" ) assert tail == 555 assert audio.numel() == reference.numel() == num_frames * upsample - tail - # the first chunk ends ``tail`` short; every later one carries 25 frames - chunk = config.codec_chunk_frames * upsample - expected = [chunk - tail] + [chunk] * max((num_frames - 1) // config.codec_chunk_frames - 1, 0) - if audio.numel() > sum(expected): - expected.append(audio.numel() - sum(expected)) - assert [c.numel() for c in chunks if c.numel()] == expected + assert sizes == expected_chunk_samples(num_frames, schedule, config.codec_chunk_frames, upsample, tail) all_boundaries += boundaries all_middles += middles + return all_boundaries, all_middles + +def test_real_codes_stream_to_the_full_decode_length_without_boundary_dips(code2wav): + all_boundaries, all_middles = stream_all(code2wav, ()) boundary_median = statistics.median(all_boundaries) middle_median = statistics.median(all_middles) print( @@ -207,3 +223,21 @@ def test_real_codes_stream_to_the_full_decode_length_without_boundary_dips(code2 # does everything after it. assert boundary_median >= middle_median - 12 assert min(all_boundaries) >= 10 + + +def test_the_served_schedule_keeps_boundaries_as_close_as_steady_chunks(code2wav): + """The served edge's small first chunk shifts every later window boundary, + so the waveform differs from the 25-frame schedule's; every window's new + frames still have up to 25 frames of context, so its boundaries should be + about as close to the full decode.""" + schedule = code2wav[2].codec_chunk_schedule + assert schedule and schedule[0] < code2wav[2].codec_chunk_frames + steady, _ = stream_all(code2wav, ()) + served, _ = stream_all(code2wav, schedule) + print( + f"boundary SNR vs full decode: schedule {schedule} median {statistics.median(served):.1f} dB " + f"(min {min(served):.1f}), 25-frame median {statistics.median(steady):.1f} dB (min {min(steady):.1f})" + ) + assert len(served) >= 10 + assert statistics.median(served) >= statistics.median(steady) - 3 + assert min(served) >= 10 diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py new file mode 100644 index 0000000000..421e3207dc --- /dev/null +++ b/test/modular/test_bagel_cfg_prune.py @@ -0,0 +1,315 @@ +"""BAGEL image_gen runs only the CFG branches a step needs. + +Guidance on, the old step always ran main, cfg_text and cfg_img as one +batched forward, though the combine never reads ``v_cfg_img`` when the image +scale is <= 1 (the text-to-image default), and outside the CFG interval it +reduces to ``v_main``. ``cfg_branches`` decides 1/2/3 on the host from the +loop counter; ``declare_step`` plans that many labels and the forward runs +that many. These pin the decision table, its agreement with the old +device-side interval test, the combine reducing to what the old one computed, +and the declaration matching the branch count. CPU only. +""" + +from __future__ import annotations + +import dataclasses +import types + +import pytest +import torch + +from mstar.engine.resources import KVStep, PositionStep +from mstar.model.bagel.config import BagelModelConfig +from mstar.model.bagel.submodules import ( + CFG_BRANCHES, + IMAGE_GEN_LOOP, + LLMSubmodule, + active_labels, + cfg_branches, + combine_cfg_velocity, + image_gen_timestep, +) +from mstar.model.submodule_base import ARNodeInputs + +_DEFAULTS = {f.name: f.default for f in dataclasses.fields(BagelModelConfig)} +N = _DEFAULTS["num_timesteps"] +SHIFT = _DEFAULTS["timestep_shift"] +T2I = dict( # text-to-image request defaults + cfg_interval=tuple(_DEFAULTS["cfg_interval"]), + text_scale=_DEFAULTS["cfg_text_scale"], + img_scale=_DEFAULTS["cfg_img_scale"], +) +EDIT = dict(cfg_interval=(0.0, 1.0), text_scale=4.0, img_scale=2.0) # bagel_model edit overrides +STEPS = range(N - 1) # time_index 0 .. N-2 + + +def _device_timestep(k: int) -> torch.Tensor: + """What `_forward_image_gen` computes from its time_index input.""" + time_index = torch.full((8,), float(k), dtype=torch.float32) + t_uniform = 1.0 - time_index / (N - 1) + return LLMSubmodule._apply_timestep_shift(t=t_uniform, shift=SHIFT) + + +def _old_in_interval(k: int, interval) -> bool: + t_val = _device_timestep(k)[0] + lo, hi = interval + return bool((t_val > lo) & (t_val <= hi)) + + +def _branches(k: int, cfg_interval, text_scale, img_scale) -> int: + return cfg_branches(image_gen_timestep(k, N, SHIFT), cfg_interval, text_scale, img_scale) + + +# --- the decision --------------------------------------------------------- + + +def test_host_timestep_matches_the_device_formula(): + for k in STEPS: + assert image_gen_timestep(k, N, SHIFT) == _device_timestep(k)[0].item() + + +def test_text_to_image_defaults_run_two_branches_in_interval_and_one_outside(): + counts = [_branches(k, **T2I) for k in STEPS] + assert set(counts) == {1, 2}, "t2i (img scale 1.0) must never run cfg_img" + # the schedule is decreasing in t, so the interval is a prefix of the steps + n_in = counts.count(2) + assert counts == [2] * n_in + [1] * (len(counts) - n_in) + assert 0 < n_in < len(counts) + + +def test_image_edit_runs_all_three_branches_every_step(): + assert [_branches(k, **EDIT) for k in STEPS] == [3] * len(STEPS) + + +def test_guidance_off_and_text_off_cases(): + assert _branches(0, cfg_interval=(0.0, 1.0), text_scale=1.0, img_scale=1.0) == 1 + # image guidance alone still needs cfg_text (labels are a prefix) + assert _branches(0, cfg_interval=(0.0, 1.0), text_scale=1.0, img_scale=1.5) == 3 + assert _branches(0, cfg_interval=(0.0, 1.0), text_scale=4.0, img_scale=1.0) == 2 + + +@pytest.mark.parametrize("interval", [ + (0.4, 1.0), (0.0, 1.0), (0.1, 0.7), (0.3333333333, 0.6666666666), + # bounds landing exactly on schedule values: (lo, hi] is open below + "on_grid", +]) +def test_interval_agrees_with_the_old_device_test(interval): + if interval == "on_grid": + interval = (image_gen_timestep(30, N, SHIFT), image_gen_timestep(5, N, SHIFT)) + for k in STEPS: + in_old = _old_in_interval(k, interval) + n = _branches(k, interval, 4.0, 2.0) + assert (n > 1) == in_old, (k, interval) + + +# --- the combine ---------------------------------------------------------- + + +def _old_combine(k, v_main, v_text, v_img, interval, text_scale, img_scale, renorm, rmin): + """The always-three-branch step: scales gated by the interval on device.""" + t_val = _device_timestep(k)[0] + lo, hi = interval + in_cfg = ((t_val > lo) & (t_val <= hi)).float() + eff_text = text_scale * in_cfg + 1.0 * (1 - in_cfg) + eff_img = img_scale * in_cfg + 1.0 * (1 - in_cfg) + return combine_cfg_velocity(v_main, v_text, v_img, eff_text, eff_img, renorm, rmin) + + +def _pruned_combine(k, v_main, v_text, v_img, interval, text_scale, img_scale, renorm, rmin): + n = _branches(k, interval, text_scale, img_scale) + if n == 1: + return v_main + return combine_cfg_velocity( + v_main, v_text, v_img if n == 3 else None, text_scale, img_scale, renorm, rmin, + ) + + +@pytest.mark.parametrize("renorm", ["global", "channel", "text_channel"]) +@pytest.mark.parametrize("rmin", [0.0, 0.5]) +@pytest.mark.parametrize("request_kind", ["t2i", "edit", "img_only_in_window"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_pruned_combine_matches_the_old_combine(renorm, rmin, request_kind, dtype): + params = { + "t2i": T2I, + "edit": EDIT, + "img_only_in_window": dict(cfg_interval=(0.4, 1.0), text_scale=4.0, img_scale=1.5), + }[request_kind] + gen = torch.Generator().manual_seed(0) + v_main, v_text, v_img = ( + torch.randn(64, 16, generator=gen).to(dtype) for _ in range(3) + ) + args = (params["cfg_interval"], params["text_scale"], params["img_scale"], renorm, rmin) + for k in STEPS: + old = _old_combine(k, v_main, v_text, v_img, *args) + new = _pruned_combine(k, v_main, v_text, v_img, *args) + assert new.dtype == old.dtype + if _branches(k, *args[:3]) > 1: + # in the interval: the same ops on the same branches + assert torch.equal(new, old), (k, renorm) + else: + # outside: the old combine was v_main up to rounding + tol = 1e-6 if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(new, old, rtol=tol, atol=tol) + + +# --- the declaration and the forward's inputs ----------------------------- + + +def _stub(**overrides): + return types.SimpleNamespace( + node_name="LLM", + CFG_BATCHED_LABEL=LLMSubmodule.CFG_BATCHED_LABEL, + _get_active_labels=lambda walk, cfg: active_labels(walk, cfg, "LLM"), + **overrides, + ) + + +def _image_gen_row(n_branches: int | None, span: int = 6) -> ARNodeInputs: + row = ARNodeInputs(input_seq_len=span) + row.resource_step_info = True # requires_cfg + row.custom_pos_ids = { + label: torch.full((span,), 10 + i, dtype=torch.int32) + for i, label in enumerate(("main", "cfg_text", "cfg_img")) + } + if n_branches is not None: + row.kwargs[CFG_BRANCHES] = n_branches + return row + + +@pytest.mark.parametrize("n_branches,labels", [ + (None, ("main", "cfg_text", "cfg_img")), # pruning off / unstamped + (3, ("main", "cfg_text", "cfg_img")), + (2, ("main", "cfg_text")), + (1, ("main",)), +]) +def test_declared_step_plans_exactly_the_branches_run(n_branches, labels): + span = 6 + step = LLMSubmodule.declare_step( + _stub(), "image_gen", ["r0"], [_image_gen_row(n_branches, span)], + ) + kv: KVStep = step.steps["kv"] + rope: PositionStep = step.steps["rope"] + assert tuple(seg.label for seg in kv.segments) == labels + assert kv.commit is False + assert rope.advance == (0,) * len(labels) + if len(labels) > 1: + assert kv.combined_labels == {labels: LLMSubmodule.CFG_BATCHED_LABEL} + pos = rope.pos_ids[LLMSubmodule.CFG_BATCHED_LABEL] + assert pos.tolist() == [10 + i for i in range(len(labels)) for _ in range(span)] + else: + assert not kv.combined_labels + assert rope.pos_ids["main"].tolist() == [10] * span + assert step.steps["attn"].causal is False + + +@pytest.mark.parametrize("n_branches", [None, 1, 2, 3]) +def test_preprocess_hands_the_forward_the_declared_count(n_branches): + row = _image_gen_row(n_branches) + out = LLMSubmodule.preprocess( + _stub(), "image_gen", + engine_inputs=types.SimpleNamespace( + per_request_info={"r0": types.SimpleNamespace(step_metadata={})}, + ), + inputs=[row], + ) + assert "kwargs" not in out + assert out.get(CFG_BRANCHES) == n_branches + assert out["requires_cfg"] is True + + +def _request(k, first_iter, prune=True, **meta): + cfg = types.SimpleNamespace( + num_timesteps=N, timestep_shift=SHIFT, + cfg_interval=_DEFAULTS["cfg_interval"], + cfg_text_scale=_DEFAULTS["cfg_text_scale"], + cfg_img_scale=_DEFAULTS["cfg_img_scale"], + ) + stub = _stub(config=cfg, _cfg_prune=prune, _warned_loop_counter=False) + info = types.SimpleNamespace( + step_metadata={"requires_cfg": True, **meta}, + dynamic_loop_iter_counts={} if k is None else {IMAGE_GEN_LOOP: k}, + ) + return LLMSubmodule._image_gen_cfg_branches(stub, info, first_iter) + + +def test_branch_count_follows_the_loop_counter_and_request_kwargs(): + last = N - 2 + assert _request(0, True) == 2 + assert _request(last, False) == 1 + # per-request kwargs (as step metadata) override the config defaults + assert _request(last, False, cfg_interval=[0.0, 1.0], cfg_img_scale=2.0) == 3 + assert _request(0, True, cfg_img_scale=2.0, cfg_interval=[0.0, 1.0]) == 3 + + +def test_falls_back_to_all_branches_when_off_or_counter_is_not_tracking(): + assert _request(5, False, prune=False) is None + assert _request(None, False) is None # loop counter missing on a later step + assert _request(3, True) is None # seeded latents but counter says step 3 + assert _request(0, True, requires_cfg=False) is None + + +def test_image_gen_walk_names_the_loop_the_counter_is_read_from(): + """An unnamed Loop gets a random uuid, and `dynamic_loop_iter_counts` + would never carry IMAGE_GEN_LOOP: every step would fall back.""" + from mstar.model.bagel.bagel_model import BagelModel + + model = BagelModel.__new__(BagelModel) + model.config = types.SimpleNamespace( + num_timesteps=N, temperature=1.0, top_k=0, top_p=1.0, + repetition_penalty=1.0, ignore_eos=False, vocab_size=32, + ) + model._has_cfg_parallel = False + model._has_llm_disaggregation = False + model._image_gen_remote_handoff = False + loops = model.get_graph_walk_graphs()["image_gen"].get_loops() + assert IMAGE_GEN_LOOP in loops + assert loops[IMAGE_GEN_LOOP].max_iters == N - 1 + + +# --- parallel CFG (image_gen_cfg) --------------------------------------------- + + +def test_parallel_branch_nodes_sit_out_the_branches_a_step_does_not_read(): + from mstar.model.bagel.submodules import cfg_branch_skipped + + # main always runs; cfg_text runs from 2 branches; cfg_img only at 3 + assert [cfg_branch_skipped("main", n) for n in (1, 2, 3)] == [False, False, False] + assert [cfg_branch_skipped("cfg_text", n) for n in (1, 2, 3)] == [True, False, False] + assert [cfg_branch_skipped("cfg_img", n) for n in (1, 2, 3)] == [True, True, False] + # no count stamped (pruning off): every branch runs + assert not any(cfg_branch_skipped(label, None) for label in ("main", "cfg_text", "cfg_img")) + + +def test_parallel_combine_reduces_to_the_old_combine(): + """combine_cfg with a host branch count gives what the device-side interval test gave.""" + from mstar.model.bagel.submodules import CombineCFGSubmodule + + cfg = types.SimpleNamespace(**{k: _DEFAULTS[k] for k in ( + "num_timesteps", "timestep_shift", "cfg_text_scale", "cfg_img_scale", + "cfg_interval", "cfg_renorm_type", "cfg_renorm_min", + )}) + node = CombineCFGSubmodule(cfg) + g = torch.Generator().manual_seed(0) + v = [torch.randn(16, 8, generator=g) for _ in range(3)] + latents = torch.randn(16, 8, generator=g) + for k in (0, N - 2): + time_index = torch.full((16,), float(k)) + n = _branches(k, **T2I) + old = node.forward("image_gen_cfg", None, *v, latents, time_index)["latents"][0] + new = node.forward("image_gen_cfg", None, *v, latents, time_index, cfg_branches=n)["latents"][0] + torch.testing.assert_close(new, old, rtol=0, atol=1e-6) + + +def test_image_gen_cfg_walk_names_the_loop_too(): + from mstar.model.bagel.bagel_model import BagelModel + + model = BagelModel.__new__(BagelModel) + model.config = types.SimpleNamespace( + num_timesteps=N, temperature=1.0, top_k=0, top_p=1.0, + repetition_penalty=1.0, ignore_eos=False, vocab_size=32, + ) + model._has_cfg_parallel = True + model._has_llm_disaggregation = False + model._image_gen_remote_handoff = False + loops = model.get_graph_walk_graphs()["image_gen_cfg"].get_loops() + assert IMAGE_GEN_LOOP in loops diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py new file mode 100644 index 0000000000..9bbea4ff4f --- /dev/null +++ b/test/modular/test_bagel_postprocess.py @@ -0,0 +1,21 @@ +"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path sent.""" +import io + +import numpy as np +import pytest +import torch + +from mstar.model.bagel.bagel_model import BagelModel + +PIL = pytest.importorskip("PIL.Image") + + +def test_image_png_decodes_to_the_truncated_pixels(): + g = torch.Generator().manual_seed(0) + output = torch.rand(1, 3, 48, 64, generator=g) + model = BagelModel.__new__(BagelModel) + png = model.postprocess(output, "image", {}) + decoded = np.asarray(PIL.open(io.BytesIO(png)).convert("RGB")) + # the old path: permute to HWC, scale, truncate to uint8 + expected = (output[0].permute(1, 2, 0) * 255).to(torch.uint8).numpy() + assert np.array_equal(decoded, expected) diff --git a/test/modular/test_chatterbox_model.py b/test/modular/test_chatterbox_model.py index 710575d3f6..6e57c1f5ae 100644 --- a/test/modular/test_chatterbox_model.py +++ b/test/modular/test_chatterbox_model.py @@ -711,8 +711,8 @@ def build_prefill_embeds(self, cond_emb, text_ids, *, uncond=False): parts = [cond_emb, text, bos] + ([bos] if self.config.duplicate_bos_in_prefill else []) return torch.cat(parts) - def embed_speech(self, ids, positions): - return self.speech_emb(ids) + self.speech_pos_emb(positions) + def embed_speech(self, ids, position): + return self.speech_emb(ids) + self.speech_pos_emb.weight[position] def _t3_submodule(variant="chatterbox"): diff --git a/test/modular/test_conductor_replica_assignment.py b/test/modular/test_conductor_replica_assignment.py new file mode 100644 index 0000000000..e2effee144 --- /dev/null +++ b/test/modular/test_conductor_replica_assignment.py @@ -0,0 +1,52 @@ +"""Data-parallel replica choice: a new request goes to the replica serving the +fewest live requests, so concurrent requests spread over the replicas instead of +landing on one at random.""" +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.conductor.conductor import Conductor # noqa: E402 + + +def _wg(group_id, ranks, instance_ranks=()): + return SimpleNamespace(_group_id=group_id, ranks=list(ranks), _instance_ranks=[list(r) for r in instance_ranks]) + + +def _conductor(worker_graphs): + conductor = object.__new__(Conductor) + conductor.worker_graphs = worker_graphs + conductor.requests = {} + return conductor + + +def _admit(conductor, rid): + assignment = conductor._assign_worker_graphs_to_workers() + conductor.requests[rid] = SimpleNamespace(worker_graph_to_workers=assignment) + return assignment + + +def test_requests_alternate_over_dp_replicas(): + # two node groups replicated on ranks 0 and 1, one wg each + conductor = _conductor({0: _wg(0, [0, 1]), 1: _wg(0, [0, 1])}) + first = _admit(conductor, "a") + second = _admit(conductor, "b") + # a group's wgs stay on one replica, and the second request takes the idle one + assert first[0] == first[1] and second[0] == second[1] + assert first[0] != second[0] + del conductor.requests["a"] + assert _admit(conductor, "c")[0] == first[0] + + +def test_lockstep_instances_are_balanced_whole(): + # one sp2 node replicated as two instances: ranks (0, 1) and (2, 3) + conductor = _conductor({0: _wg(0, [0, 1, 2, 3], instance_ranks=[(0, 1), (2, 3)])}) + picks = {tuple(_admit(conductor, rid)[0]) for rid in ("a", "b")} + assert picks == {("worker_0", "worker_1"), ("worker_2", "worker_3")} + + +def test_single_replica_needs_no_load_scan(): + conductor = _conductor({0: _wg(0, [3])}) + conductor.requests = None # never read when there is nothing to choose + assert conductor._assign_worker_graphs_to_workers() == {0: ["worker_3"]} diff --git a/test/modular/test_fish_s2_pro_model.py b/test/modular/test_fish_s2_pro_model.py index fc2e75cd61..c261949411 100644 --- a/test/modular/test_fish_s2_pro_model.py +++ b/test/modular/test_fish_s2_pro_model.py @@ -550,7 +550,13 @@ def old_flags(flags): assert_same_tensor(packed["last_token_indices"], torch.tensor(seq_lens).cumsum(0) - 1) assert_same_tensor(packed["suppress_eos"], old_flags([True, True, True])) - (graph_config,) = ar.get_accelerator_graph_configs(torch.device("cpu")) + prefill_config, graph_config = ar.get_accelerator_graph_configs(torch.device("cpu")) + # Plain prefills replay packed captures; their padding rows suppress EOS. + assert prefill_config.capture_graph_walk == "ar_prefill" + # a voice clone's prefill is the same forward: it replays these captures too + assert prefill_config.replay_graph_walks == ["ar_prefill", "ar_prefill_clone"] + assert prefill_config.make_node_input(5).kwargs == {"suppress_eos": True} + assert graph_config.capture_graph_walk == "ar_decode" template = graph_config.single_request_inputs assert template.kwargs == {"suppress_eos": False} and not template.tensor_inputs padded = [decode_row(True), decode_row(False), *graph_config.get_node_inputs(2, 0)] diff --git a/test/modular/test_flux2_klein_model.py b/test/modular/test_flux2_klein_model.py index 04234906ef..d65093819c 100644 --- a/test/modular/test_flux2_klein_model.py +++ b/test/modular/test_flux2_klein_model.py @@ -133,11 +133,14 @@ def test_dummy_mode_returns_no_submodules(): def test_process_prompt_tokenizes_with_chat_template_to_fixed_length(): model = _make_model() out = model.process_prompt("a cat", ["text"], ["image"]) - ids, mask = out["text_inputs"][0], out["text_mask"][0] + ids, mask = out.new_input_tensors["text_inputs"][0], out.new_input_tensors["text_mask"][0] assert ids.shape == (512,) and mask.shape == (512,) and ids.dtype == torch.long n = int(mask.sum()) assert bytes(ids[:n].tolist()) == b"a cat" assert int(mask[n:].sum()) == 0 + # the text encoder's attention path is settled here, on the host + assert out.metadata == {"text_padded": True} + assert "text_padded" in model.INTERNAL_KWARGS @pytest.mark.parametrize("kwargs,message", [ diff --git a/test/modular/test_higgs_audio_decode_hooks.py b/test/modular/test_higgs_audio_decode_hooks.py new file mode 100644 index 0000000000..cb527a01c6 --- /dev/null +++ b/test/modular/test_higgs_audio_decode_hooks.py @@ -0,0 +1,78 @@ +"""Higgs LLM decode loop: one token signal kept on the device and sent to the +client inline, and a batched stop check that agrees with the per-request one.""" + +from types import SimpleNamespace + +import torch + +from mstar.graph.special_destinations import EMIT_TO_CLIENT +from mstar.model.higgs_audio.config import SAMPLER, HiggsAudioModelConfig +from mstar.model.higgs_audio.higgs_audio_model import HiggsAudioModel +from mstar.model.higgs_audio.submodules import HiggsAudioLLMSubmodule +from mstar.model.submodule_base import HostRows + + +def _submodule(slot_masters=True): + sub = HiggsAudioLLMSubmodule(torch.nn.Linear(1, 1), HiggsAudioModelConfig()) + sub.node_resources = {SAMPLER: SimpleNamespace(has_slot_masters=slot_masters)} + return sub + + +def _info(walk="decode", decoded=0, max_tokens=8, ignore_eos=False): + return SimpleNamespace( + graph_walk=walk, + dynamic_loop_iter_counts={"decode_loop": decoded}, + max_tokens=max_tokens, + resource_configs={SAMPLER: SimpleNamespace(ignore_eos=ignore_eos)}, + step_metadata={}, + ) + + +def test_decode_loop_routes_one_signal_to_itself_and_the_client(): + model = object.__new__(HiggsAudioModel) + model.config = HiggsAudioModelConfig() + outputs = model.get_graph_walk_graphs()["decode"].section.outputs + assert {(edge.next_node, edge.name) for edge in outputs} == { + ("LLM", "text_inputs"), (EMIT_TO_CLIENT, "text_inputs"), + } + + +def test_decode_keeps_its_token_on_the_device_only_with_slot_masters(): + sub = _submodule() + assert sub.reads_device_loopback + assert sub.device_loopback_signals("decode") == {"text_inputs"} + assert sub.device_loopback_signals("prefill_audio") == frozenset() + assert sub.inline_client_signals("decode") == {"text_inputs": "new_token"} + assert sub.inline_client_signals("prefill_audio") == {} + row = sub.prepare_inputs("decode", _info(), {}, has_loopback=True) + assert row.input_ids is None and row.input_seq_len == 1 + assert sub.uniform_row_inputs("decode", has_loopback=True).input_seq_len == 1 + assert sub.uniform_row_inputs("decode", has_loopback=False) is None + + eager = _submodule(slot_masters=False) + assert eager.device_loopback_signals("decode") == frozenset() + assert eager.uniform_row_inputs("decode", has_loopback=True) is None + + +def test_batched_stop_check_agrees_with_the_per_request_one(): + sub = _submodule() + eos = next(iter(sub.config.stop_token_ids)) + infos = { + "a": _info(), + "b": _info(), # stop token + "c": _info(decoded=7), # out of budget + "d": _info(ignore_eos=True), # stop token ignored + } + tokens = torch.tensor([[11], [eos], [11], [eos], [11]]) # a padded replay's extra row last + want = { + rid: stops for i, rid in enumerate(infos) + if (stops := sub.check_stop(rid, infos[rid], {"new_token": [tokens[i]]})) + } + assert want == {"b": {"decode_loop"}, "c": {"decode_loop"}} + rows = HostRows(("a", "b", "c", "d"), {"new_token": tokens}) + assert sub.check_stop_batched(["d", "c", "b", "a"], infos, rows, step_walk="decode") == want + assert sub.check_stop_batched(list(infos), infos, rows, step_walk=None) == want + # only decode rows end the loop + assert sub.check_stop_batched(list(infos), infos, rows, step_walk="prefill_audio") == {} + prefill = {**infos, "b": _info(walk="prefill_audio")} + assert sub.check_stop_batched(list(prefill), prefill, rows, step_walk=None) == {"c": {"decode_loop"}} diff --git a/test/modular/test_higgs_audio_tower.py b/test/modular/test_higgs_audio_tower.py new file mode 100644 index 0000000000..a254edd22c --- /dev/null +++ b/test/modular/test_higgs_audio_tower.py @@ -0,0 +1,30 @@ +"""Higgs audio tower: a short clip padded to the chunk width and masked +(``valid_frames``) encodes as the unpadded clip does.""" + +import pytest +import torch + +from mstar.model.higgs_audio.components.audio_tower import HiggsAudioFeatureProjector, HiggsAudioTower +from mstar.model.higgs_audio.config import HiggsAudioModelConfig + + +@pytest.mark.parametrize("width", [3, 4, 7, 37, 64, 99, 100]) +def test_masked_padded_tower_matches_unpadded(width): + torch.manual_seed(0) + config = HiggsAudioModelConfig( + hidden_size=48, audio_d_model=32, audio_encoder_layers=2, audio_encoder_attention_heads=4, + audio_encoder_ffn_dim=64, audio_num_mel_bins=16, audio_max_source_positions=64, + ) + tower, projector = HiggsAudioTower(config).eval(), HiggsAudioFeatureProjector(config).eval() + frames = 100 + clip = torch.randn(1, config.audio_num_mel_bins, width) + other = torch.randn(1, config.audio_num_mel_bins, frames) + padded = torch.cat([torch.nn.functional.pad(clip, (0, frames - width)), other]) + with torch.no_grad(): + want = projector(tower(clip))[0] + got = projector(tower(padded, torch.tensor([width, frames]))) + full = projector(tower(other))[0] + n = config.encoder_output_length(width) + assert n == want.shape[0] + assert torch.allclose(got[0, :n], want[:n], atol=1e-5, rtol=1e-4) + assert torch.allclose(got[1], full, atol=1e-5, rtol=1e-4) diff --git a/test/modular/test_kokoro_model.py b/test/modular/test_kokoro_model.py index 823de2a3ca..3e64dbd37a 100644 --- a/test/modular/test_kokoro_model.py +++ b/test/modular/test_kokoro_model.py @@ -106,6 +106,11 @@ def test_masked_bilstm_matches_packed_bidirectional(): out = ours(x, lengths) assert torch.allclose(out, expected, atol=1e-6) assert out[1, 4:].abs().sum() == 0 and out[2, 1:].abs().sum() == 0 + # both directions as one block-diagonal LSTM: the same function + ours.fuse() + fused = ours(x, lengths) + assert torch.allclose(fused, expected, atol=1e-6) + assert fused[1, 4:].abs().sum() == 0 and fused[2, 1:].abs().sum() == 0 def test_real_stft_inverse_matches_torch_istft_per_row(): diff --git a/test/modular/test_moe_router.py b/test/modular/test_moe_router.py new file mode 100644 index 0000000000..292064be55 --- /dev/null +++ b/test/modular/test_moe_router.py @@ -0,0 +1,88 @@ +"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_softmax_topk``) +against the op sequence it replaces, and where the swept MoE tiles apply.""" +from __future__ import annotations + +import pytest +import torch + +needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@needs_cuda +@pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) +@pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) +def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): + from mstar.utils.fused_moe.router import softmax_topk + + torch.manual_seed(1) + logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 + probs = torch.softmax(logits, dim=-1, dtype=torch.float) + ref_w, ref_ids = torch.topk(probs, top_k, dim=-1) + if norm: + ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) + + w, ids = softmax_topk(logits, top_k, norm) + assert ids.dtype == torch.int64 and w.dtype == torch.float32 + # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the + # selected probabilities, and that every pick is a distinct expert holding them + torch.testing.assert_close(w, ref_w, rtol=1e-6, atol=1e-7) + assert (ids.sort(dim=-1).values.diff(dim=-1) > 0).all() + picked = probs.gather(1, ids) + if norm: + picked = picked / picked.sum(dim=-1, keepdim=True) + torch.testing.assert_close(picked, w, rtol=1e-6, atol=1e-7) + # with no ties the choice and its order are torch's + distinct = (ref_w[:, :-1] != ref_w[:, 1:]).all(dim=-1) + assert torch.equal(ids[distinct], ref_ids[distinct]) + + +@needs_cuda +def test_router_compiles_with_the_fused_topk(): + from mstar.model.components.moe import TopKRouter + + router = TopKRouter(64, 128, 8, fused_softmax_topk=True).cuda().to(torch.bfloat16).requires_grad_(False) + torch.nn.init.normal_(router.weight) + x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + eager = router(x) + compiled = torch.compile(router, fullgraph=True)(x) + assert torch.equal(eager[1], compiled[1]) + torch.testing.assert_close(eager[0], compiled[0], rtol=0, atol=0) + + +@needs_cuda +def test_router_defaults_to_the_torch_ops(): + from mstar.model.components.moe import TopKRouter + + torch.manual_seed(2) + router = TopKRouter(64, 128, 8).cuda().to(torch.bfloat16).requires_grad_(False) + torch.nn.init.normal_(router.weight) + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + w, ids, _ = router(x) + probs = torch.softmax(torch.nn.functional.linear(x, router.weight), dim=-1, dtype=torch.float) + ref_w, ref_ids = torch.topk(probs, 8, dim=-1) + assert torch.equal(ids, ref_ids) + torch.testing.assert_close(w, ref_w / ref_w.sum(dim=-1, keepdim=True), rtol=0, atol=0) + + +@pytest.mark.parametrize("platform,applies", [ + ("cuda_h", True), ("cuda_b", False), ("rocm", False), +]) +def test_tuned_moe_tiles_apply_only_on_hopper(monkeypatch, platform, applies): + pytest.importorskip("triton") + from mstar import platforms + from mstar.platforms.interface import CudaPlatform, RocmPlatform + from mstar.utils.fused_moe import kernels + + fake = { + "cuda_h": type("H", (CudaPlatform,), {"cuda_capability": lambda self, device=None: (9, 0)})(), + "cuda_b": type("B", (CudaPlatform,), {"cuda_capability": lambda self, device=None: (10, 0)})(), + "rocm": RocmPlatform(), + }[platform] + monkeypatch.setattr(platforms, "resolve_platform", lambda: fake) + kernels._tuned_tiles_apply.cache_clear() + try: + assert (kernels.get_tuned_configs(1, 128, 1536, 2048) is not None) is applies + finally: + kernels._tuned_tiles_apply.cache_clear() diff --git a/test/modular/test_omnilingual_asr_model.py b/test/modular/test_omnilingual_asr_model.py index 4fe75965fe..924c5d0dc1 100644 --- a/test/modular/test_omnilingual_asr_model.py +++ b/test/modular/test_omnilingual_asr_model.py @@ -484,6 +484,22 @@ def test_frame_lengths_follow_the_conv_arithmetic(): assert real.frame_lengths([16000, 640000, 400, 399]) == [49, 1999, 1, 0] +def test_unfolded_positional_conv_matches_the_grouped_conv(): + """A wide-group deployment runs the positional conv as per-group GEMMs (in + chunks of frames): the same values as the grouped conv, padding masked.""" + from mstar.model.omnilingual_asr.components.wav2vec2 import PositionEncoder + + torch.manual_seed(0) + pos = PositionEncoder(model_dim=32, kernel_size=8, num_groups=4).double() + x = torch.randn(3, 37, 32, dtype=torch.float64) + mask = torch.arange(37)[None, :] < torch.tensor([37, 20, 5])[:, None] + ref = pos(x, mask) + pos.unfolded = True + torch.testing.assert_close(pos(x, mask), ref) + pos.UNFOLD_CHUNK_ELEMS = 3 * 32 * 8 * 5 # five frames a chunk + torch.testing.assert_close(pos(x, mask), ref) + + def test_padded_batch_matches_single_runs_and_needs_the_mask(): model = _random_model(_tiny_config()) torch.manual_seed(1) @@ -679,6 +695,48 @@ def test_node_batches_and_matches_single_requests(tiny_model): assert sub.max_batch_size(TRANSCRIBE_WALK) == 4 +def test_a_step_over_the_audio_budget_is_not_encoded_packed(tiny_model, monkeypatch): + """With the packed encoder bound, a step whose frames fit `max_batch_seconds` + encodes in one packed call (leased or planned whole); one over it plans and + leases nothing and encodes padded per sub-batch.""" + from mstar.model.omnilingual_asr.config import ENCODER_ATTN + from mstar.model.omnilingual_asr.submodules import ENCODER_LOOP, OmniASRCTCSubmodule + + sub = tiny_model.get_submodule(ENCODER_NODE, "cpu") + small = OmniASRCTCSubmodule(sub.model, max_batch_size=4, max_batch_seconds=0.5) # 8000 samples + small.node_resources = {ENCODER_ATTN: object()} + packed_calls = [] + monkeypatch.setattr(small, "_transcribe_packed", lambda waves, attn, ei: packed_calls.append(len(waves)) or [ + torch.zeros(n, dtype=torch.long) for n in small.model.frame_lengths([w.shape[0] for w in waves]) + ]) + torch.manual_seed(6) + lease = small.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[ENCODER_LOOP].can_lease + + def step(lengths): + waves = [audio.preprocess(torch.randn(n)) for n in lengths] + rids = [f"r{i}" for i in range(len(waves))] + inputs = [ + small.prepare_inputs(TRANSCRIBE_WALK, types.SimpleNamespace(request_id=r), {WAVEFORM: [w]}) + for r, w in zip(rids, waves, strict=True) + ] + declared = small.declare_step(TRANSCRIBE_WALK, rids, inputs) + engine_inputs = types.SimpleNamespace(request_ids=rids, resources={ENCODER_ATTN: object()}) + pre = small.preprocess(TRANSCRIBE_WALK, engine_inputs, inputs) + with torch.no_grad(): + out = small.forward_batched(TRANSCRIBE_WALK, engine_inputs, **pre) + return waves, inputs, declared, out + + _, inputs, declared, _ = step((1000, 2000, 3000)) + assert lease(inputs) and declared is not None and len(declared.segments) == 3 + assert packed_calls == [3] + + waves, inputs, declared, out = step((6000, 1000, 5000, 2000)) + assert not lease(inputs) and declared is None + assert packed_calls == [3] + for i, w in enumerate(waves): + assert torch.equal(out[f"r{i}"][TOKEN_IDS][0], _run_node(sub, [w])["r0"]) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_node_step_never_waits_on_the_gpu(tiny_model): """A padded step cut into sub-batches launches without a host sync.""" diff --git a/test/modular/test_orpheus_audio_vocab.py b/test/modular/test_orpheus_audio_vocab.py new file mode 100644 index 0000000000..43f3082dfa --- /dev/null +++ b/test/modular/test_orpheus_audio_vocab.py @@ -0,0 +1,97 @@ +"""Orpheus's LM head computes only the custom-token rows, and its stop check +runs batched off the step's host rows. + +CPU-only: the sliced head equals the full head with the text rows masked and +hands the sampler full-width logits, so token ids stay in full-vocab space; +``check_stop_batched`` agrees with ``check_stop``. +""" +from types import SimpleNamespace + +import torch +from torch import nn + +from mstar.model.orpheus.config import ATTN, SAMPLER, OrpheusModelConfig +from mstar.model.orpheus.submodules import OrpheusLLMSubmodule +from mstar.model.submodule_base import HostRows + +V, H, LO = 96, 16, 60 + + +def _config() -> OrpheusModelConfig: + return OrpheusModelConfig(vocab_size=V, custom_token_base_id=LO, stop_token_id=LO + 2) + + +class _FakeLM(nn.Module): + def __init__(self, tp_size: int = 1): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(V, H) + self.lm_head = nn.Linear(H, V, bias=False) + self.lm_head.tp_size = tp_size + + def forward(self, emb, *, label): + return emb + + +class _GreedySampler: + def sample(self, request_ids, logits): + self.logits = logits + return logits.argmax(-1) + + +def _forward(tp_size: int = 1): + torch.manual_seed(3) + lm = _FakeLM(tp_size) + sub = OrpheusLLMSubmodule(lm, _config()) + sampler = _GreedySampler() + ids = torch.tensor([1, 7, LO + 4, V - 1]) + engine_inputs = SimpleNamespace(resources={SAMPLER: sampler, ATTN: None}, request_ids=list("abcd")) + with torch.no_grad(): + out = sub.forward("decode", engine_inputs, text_inputs=ids)["new_token"] + full = lm.lm_head(lm.model.embed_tokens(ids)) + return sub, sampler.logits, out, full + + +def test_real_config_ids_fall_in_the_custom_rows(): + cfg = OrpheusModelConfig() + assert cfg.stop_token_id >= cfg.custom_token_base_id + # every SNAC code position (7 per frame, 4096 codes each) after the 10 markers + assert cfg.custom_token_base_id + 10 + cfg.tokens_per_frame * 4096 <= cfg.vocab_size + + +def test_head_logits_are_the_full_head_with_text_rows_masked(): + sub, logits, out, full = _forward() + assert sub._audio_vocab_start == LO and logits.shape == (4, V) + masked = full.clone() + masked[:, :LO] = float("-inf") + torch.testing.assert_close(logits, masked) + assert torch.equal(out, masked.argmax(-1)) and (out >= LO).all() + + +def test_full_head_under_tensor_parallel(): + sub, logits, _, full = _forward(tp_size=2) + assert sub._audio_vocab_start is None + torch.testing.assert_close(logits, full) + + +def test_batched_stop_check_agrees_with_the_per_request_one(): + sub = OrpheusLLMSubmodule(_FakeLM(), _config()) + eos = LO + 2 + + def info(ignore_eos: bool, done: int, max_tokens: int = 100): + return SimpleNamespace( + resource_configs={SAMPLER: SimpleNamespace(ignore_eos=ignore_eos)}, + dynamic_loop_iter_counts={"decode_loop": done}, max_tokens=max_tokens, + ) + + cases = {"eos": (eos, info(False, 3)), "ignored": (eos, info(True, 3)), + "budget": (LO + 9, info(False, 9, max_tokens=10)), "running": (LO + 9, info(False, 3))} + infos = {rid: i for rid, (_, i) in cases.items()} + expected = { + rid: stop for rid, (tok, i) in cases.items() + if (stop := sub.check_stop(rid, i, {"new_token": [torch.tensor(tok)]})) + } + assert expected == {"eos": {"decode_loop"}, "budget": {"decode_loop"}} + order = list(reversed(cases)) + rows = HostRows(tuple(order), {"new_token": torch.tensor([cases[r][0] for r in order] + [eos])}) + assert sub.check_stop_batched(list(cases), infos, rows, "decode") == expected diff --git a/test/modular/test_piecewise_can_lease.py b/test/modular/test_piecewise_can_lease.py index 14e545341d..7a8b9f4ea2 100644 --- a/test/modular/test_piecewise_can_lease.py +++ b/test/modular/test_piecewise_can_lease.py @@ -13,6 +13,8 @@ class _Runner: + one_row_per_step = False + def __init__(self, lease_before_step=True, accepts=True): self.lease_before_step = lease_before_step self._accepts = accepts @@ -55,3 +57,23 @@ def runner(predicate): assert PiecewiseAcceleratorGraphRunner.can_lease(runner(None), []) assert not PiecewiseAcceleratorGraphRunner.can_lease(runner(lambda inputs: len(inputs) < 2), [1, 2]) + + +def test_a_one_row_per_step_region_leases_the_whole_step_as_one_row(): + """A region that replays the whole step as one row (a packed encoder) is + leased at batch size 1 for every token of the step, through the real + runner; a per-row region beside it is leased at the step's batch size.""" + def real_runner(one_row): + runner = PiecewiseAcceleratorGraphRunner.__new__(PiecewiseAcceleratorGraphRunner) + runner._config = SimpleNamespace(lease_before_step=True, can_lease=None, one_row_per_step=one_row) + runner._current_slot = 0 + runner.resolved = [] + runner._resolve = lambda bs, total: runner.resolved.append((bs, total)) or SimpleNamespace(bucket=(bs, total)) + return runner + + packed, per_row = real_runner(True), real_runner(False) + inputs = [SimpleNamespace(input_seq_len=n) for n in (100, 30, 7)] + leases = _lease({"packed": packed, "per_row": per_row}, inputs) + assert packed.resolved == [(1, 137)] and per_row.resolved == [(3, 137)] + assert leases["packed"].bucket == (1, 137) and leases["packed"].slot == 0 + assert leases["per_row"].bucket == (3, 137) diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py index 8c01e78467..322cfc61de 100644 --- a/test/modular/test_qwen3_5_device_loopback.py +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -126,7 +126,7 @@ def test_the_loop_back_decode_row_is_uniform(monkeypatch): assert LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=True) is None -def test_in_graph_rope_knob_drops_only_the_staged_tables(monkeypatch): +def test_captured_decode_drops_only_the_staged_tables(): sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4])) sub = _sub(sampler) rows = [_prepare(sub, "decode", {}) for _ in range(2)] @@ -139,10 +139,6 @@ def test_in_graph_rope_knob_drops_only_the_staged_tables(monkeypatch): embed_tokens=SimpleNamespace(weight=torch.zeros(1)), )) sub._position_ids_3d = lambda inputs: torch.zeros(3, len(inputs)) - monkeypatch.delenv("MSTAR_INGRAPH_DECODE_ROPE", raising=False) - out = LLMSubmodule.preprocess(sub, "decode", leased, rows) - assert out["input_ids"].tolist() == [3, 4] and "cos_3d" in out - monkeypatch.setenv("MSTAR_INGRAPH_DECODE_ROPE", "1") out = LLMSubmodule.preprocess(sub, "decode", leased, rows) assert out["input_ids"].tolist() == [3, 4] and "cos_3d" not in out eager = SimpleNamespace( @@ -153,11 +149,11 @@ def test_in_graph_rope_knob_drops_only_the_staged_tables(monkeypatch): assert out["input_ids"].tolist() == [3, 4] and "cos_3d" in out -def test_forward_in_graph_rope_reads_planned_positions_at_capture(monkeypatch): +def test_forward_in_graph_rope_reads_planned_positions_at_capture(): """A capture hands the region the slot's static input_ids and rotary - tables. The ids are always the static ones; with the rope knob on, the - forward builds the tables from the planned positions instead. The - capture's own context carries the piecewise walk, not "decode".""" + tables. The ids are always the static ones; the forward builds the tables + from the planned positions instead. The capture's own context carries the + piecewise walk, not "decode".""" from mstar.model.qwen3_5.config import ATTN, ROPE calls = {} @@ -199,13 +195,9 @@ def run(lease): assert out == ("sampled", ["r1", "r2"]) return dict(calls) - monkeypatch.delenv("MSTAR_INGRAPH_DECODE_ROPE", raising=False) - c = run(object()) - assert c["ids"] == [9, 9] and c["cos"] is static_cos and "built" not in c - monkeypatch.setenv("MSTAR_INGRAPH_DECODE_ROPE", "1") c = run(object()) assert c["ids"] == [9, 9] and c["built"].tolist() == [[10, 20]] * 3 and c["cos"] is not static_cos - # no lease (an eager step): the knob does not apply, the inputs are used + # no lease (an eager step): the inputs are used c = run(None) assert c["ids"] == [9, 9] and c["cos"] is static_cos and "built" not in c diff --git a/test/modular/test_qwen3_5_prompt_interleave.py b/test/modular/test_qwen3_5_prompt_interleave.py index b19ab238b2..de4d8b8e95 100644 --- a/test/modular/test_qwen3_5_prompt_interleave.py +++ b/test/modular/test_qwen3_5_prompt_interleave.py @@ -23,7 +23,7 @@ parts_from_modalities, prefill_plan, ) -from mstar.model.qwen3_5.qwen3_5_model import IMAGE_PART, TEXT_PART +from mstar.model.qwen3_5.qwen3_5_model import IMAGE_GRIDS_KEY, IMAGE_PART, TEXT_PART CKPT = os.environ.get("QWEN3_5_CKPT") @@ -62,7 +62,9 @@ def test_segments_rebuild_the_prompt(model, modalities): tensors={"image_inputs": [an_image()] * n_images}, prompt_parts=None, ) + grids = out.metadata[IMAGE_GRIDS_KEY] out = out.new_input_tensors + assert grids == [tuple(g.tolist()) for g in out["image_grid_thw"]] segments = out["text_inputs"] assert len(out["image_grid_thw"]) == n_images assert len(out["pixel_values"]) == n_images @@ -113,6 +115,7 @@ def test_schedule_matches_the_split(model): tensors={"image_inputs": [an_image(), an_image()]}, prompt_parts=parts, ) + host_grids = out.metadata[IMAGE_GRIDS_KEY] out = out.new_input_tensors # the conductor passes tensor handles; the identity is all that matters signals = { @@ -120,7 +123,7 @@ def test_schedule_matches_the_split(model): "pixel_values": list(out["pixel_values"]), "image_grid_thw": list(out["image_grid_thw"]), } - schedule = model._prefill_schedule(modalities, signals) + schedule = model._prefill_schedule(modalities, signals, host_grids={IMAGE_GRIDS_KEY: host_grids}) # An image prompt is one walk, whatever its layout: the interleaving # travels as `order` rather than as a walk per span. @@ -132,6 +135,8 @@ def test_schedule_matches_the_split(model): for s in prefill_plan(parts_from_modalities(modalities)) ] assert list(step.order) == planned + # the host copy of the grids rides with the step, in the same order + assert list(step.image_grids) == host_grids # the nth tag of a kind takes the nth tensor of that kind, in order for name in ("text_inputs", "pixel_values", "image_grid_thw"): @@ -237,3 +242,23 @@ def test_a_chat_the_template_refuses_is_a_bad_request(model, tmp_path): text, _, in_mods, parts = flatten_messages(messages, tmp_path) with pytest.raises(ValueError, match="System message must be at the beginning"): model.process_prompt(text, in_mods, ["text"], prompt_parts=parts) + + +@pytest.mark.parametrize("size", [(224, 224), (500, 377), (64, 900)]) +def test_byte_patches_normalize_to_the_processors_pixels(model, size): + """`process_prompt` ships un-normalized byte patches; the encoder's + normalization must give the processor's own pixels exactly, resized or not.""" + from mstar.model.qwen3_5.qwen3_5_model import _as_hwc_uint8 + from mstar.model.qwen3_5.submodules import VisionEncoderSubmodule + + torch.manual_seed(0) + image = torch.randint(0, 256, (3, size[1], size[0]), dtype=torch.uint8) + out = model.process_prompt( + "Describe.", ["image", "text"], ["text"], tensors={"image_inputs": [image]}, + ).new_input_tensors + patches = out["pixel_values"][0] + assert patches.dtype == torch.uint8 + encoder = VisionEncoderSubmodule.__new__(VisionEncoderSubmodule) + encoder._pixel_normalizer = model.pixel_normalizer() + reference = model.image_processor(images=[_as_hwc_uint8(image)], return_tensors="pt") + assert torch.equal(VisionEncoderSubmodule._normalize(encoder, patches), reference["pixel_values"]) diff --git a/test/modular/test_qwen3_5_video.py b/test/modular/test_qwen3_5_video.py index 6a67b41c3b..1f6a9e5967 100644 --- a/test/modular/test_qwen3_5_video.py +++ b/test/modular/test_qwen3_5_video.py @@ -341,6 +341,17 @@ def test_video_plan_puts_text_around_every_group(): # --------------------------------------------------------------------------- +def normalized(model, patches: torch.Tensor) -> torch.Tensor: + """`process_prompt` ships byte patches; the vision encoder normalizes them + on the device into the processor's pixels.""" + from mstar.model.qwen3_5.submodules import VisionEncoderSubmodule + + encoder = VisionEncoderSubmodule.__new__(VisionEncoderSubmodule) + encoder._pixel_normalizer = model.pixel_normalizer() + assert patches.dtype == torch.uint8 + return VisionEncoderSubmodule._normalize(encoder, patches) + + @needs_new_layout @pytest.mark.parametrize("layout", list(LAYOUTS)) def test_prompt_ids_match_processor(model, processor, layout): @@ -355,10 +366,10 @@ def test_patches_and_grids_match_processor(model, processor, layout): _, out, ref = run_both(model, processor, layout) tensors = out.new_input_tensors assert torch.equal(torch.stack(tensors["video_grid_thw"]), ref["video_grid_thw"]) - assert torch.equal(torch.cat(tensors["pixel_values_videos"]), ref["pixel_values_videos"]) + assert torch.equal(normalized(model, torch.cat(tensors["pixel_values_videos"])), ref["pixel_values_videos"]) if "image_grid_thw" in tensors: assert torch.equal(torch.stack(tensors["image_grid_thw"]), ref["image_grid_thw"]) - assert torch.equal(torch.cat(tensors["pixel_values"]), ref["pixel_values"]) + assert torch.equal(normalized(model, torch.cat(tensors["pixel_values"])), ref["pixel_values"]) # the raw frames are not kept past preprocessing assert tensors["video_inputs"] == [] @@ -462,7 +473,7 @@ def test_mp4_matches_processor(model, processor, tmp_path): ).new_input_tensors ref = processor.video_processor(videos=[str(path)], return_tensors="pt") assert torch.equal(out["video_grid_thw"][0], ref["video_grid_thw"][0]) - assert torch.equal(out["pixel_values_videos"][0], ref["pixel_values_videos"]) + assert torch.equal(normalized(model, out["pixel_values_videos"][0]), ref["pixel_values_videos"]) # --------------------------------------------------------------------------- @@ -522,6 +533,10 @@ def request(segments: int) -> NodeInputs: assert config.lease_before_step assert config.can_lease([request(4)]) assert not config.can_lease([request(5)]) + # a batched step replays as one row, so its requests share the ceiling + assert config.one_row_per_step + assert config.can_lease([request(2), request(2)]) + assert not config.can_lease([request(2), request(3)]) def test_vision_block_loop_replays_exactly_when_leased(): diff --git a/test/modular/test_qwen3_omni_code2wav_stream.py b/test/modular/test_qwen3_omni_code2wav_stream.py index 4f535e9df0..45e433301b 100644 --- a/test/modular/test_qwen3_omni_code2wav_stream.py +++ b/test/modular/test_qwen3_omni_code2wav_stream.py @@ -26,7 +26,7 @@ from mstar.model.qwen3_omni.config import Code2WavConfig, Qwen3OmniModelConfig from mstar.model.qwen3_omni.submodules import Code2WavSubmodule from mstar.model.submodule_base import InputMetadata, ModelInputsFromEngine -from mstar.streaming.chunk_policy import LeftContextChunkPolicy +from mstar.streaming.chunk_policy import LeftContextChunkPolicy, ScheduledLeftContextChunkPolicy from mstar.streaming.stream_buffer import StreamBuffer, StreamChunkInfo UPSAMPLE = 1920 @@ -37,18 +37,26 @@ def _to_pcm(wav: torch.Tensor) -> torch.Tensor: return (wav.clamp(-1, 1) * 32767).to(torch.int16) -def _stream_batched(submodule: Code2WavSubmodule, frames: dict[str, torch.Tensor]) -> tuple[dict, list]: +def _stream_batched( + submodule: Code2WavSubmodule, frames: dict[str, torch.Tensor], schedule: tuple[int, ...] | None = None, +) -> tuple[dict, list]: """Stream each request's ``frames`` ``[T, 16]`` through its own buffer - with the edge's policy. Every step batches the requests that still have a + with fixed 25-frame chunks, or with ``schedule`` ahead of them as the + served edge does. Every step batches the requests that still have a chunk, the way the scheduler batches Code2Wav. Returns, per request, its ``(chunk_info, audio, decoded)`` per step (``decoded``: whether the step ran the vocoder), and each step's ``cg_key_info``.""" cfg = submodule.config.code2wav + + def policy(): + if schedule is None: + return LeftContextChunkPolicy(chunk=cfg.codec_chunk_frames, left_context=cfg.codec_left_context_frames) + return ScheduledLeftContextChunkPolicy(schedule, cfg.codec_chunk_frames, cfg.codec_left_context_frames) + buffers = {} for rid, request_frames in frames.items(): buffers[rid] = StreamBuffer( - request_id=rid, edge_name="codec_tokens", from_partition="Talker", - policy=LeftContextChunkPolicy(chunk=cfg.codec_chunk_frames, left_context=cfg.codec_left_context_frames), + request_id=rid, edge_name="codec_tokens", from_partition="Talker", policy=policy(), ) for i, frame in enumerate(request_frames): buffers[rid].pre_read_register(f"t{i}") @@ -90,9 +98,9 @@ def _stream_batched(submodule: Code2WavSubmodule, frames: dict[str, torch.Tensor return steps, keys -def _stream(submodule: Code2WavSubmodule, frames: torch.Tensor) -> list[tuple]: +def _stream(submodule: Code2WavSubmodule, frames: torch.Tensor, schedule: tuple[int, ...] | None = None) -> list[tuple]: """``_stream_batched`` for a single request.""" - return _stream_batched(submodule, {"r": frames})[0]["r"] + return _stream_batched(submodule, {"r": frames}, schedule)[0]["r"] def _expected_samples_per_step(steps: list[tuple], tail: int) -> list[int]: @@ -106,6 +114,17 @@ def _expected_samples_per_step(steps: list[tuple], tail: int) -> list[int]: return expected +def _with_live_eos(steps: list[tuple], num_frames: int, full: torch.Tensor) -> torch.Tensor: + """``full`` plus the frame of silence a codec_eos frame emits when it was a + new item of a window that was not the stream's last.""" + live = any( + not info.is_final and info.start_offset + info.num_items == num_frames + 1 + and info.num_items > info.context_items + for info, _, _ in steps + ) + return torch.cat([full, torch.zeros(UPSAMPLE, dtype=full.dtype)]) if live else full + + # --------------------------------------------------------------------------- # A fake decoder with the real output length: exact trim arithmetic # --------------------------------------------------------------------------- @@ -160,13 +179,51 @@ def test_streamed_chunks_stitch_into_the_full_decode(num_frames): assert torch.equal(streamed, full) +@pytest.mark.parametrize("schedule", [(8,), (1, 3), (25,)]) +@pytest.mark.parametrize("num_frames", [1, 7, 8, 9, 32, 33, 34, 57, 58, 120, 125]) +@pytest.mark.parametrize("eos", [False, True]) +def test_a_small_first_chunk_stitches_into_the_full_decode(schedule, num_frames, eos): + """The served edge's first chunks are smaller than the 25-frame steady + chunk, so the second window carries fewer than 25 context frames; the + trim follows the window's real context and the stream is unchanged.""" + decoder = _LookaheadDecoder() + config = Qwen3OmniModelConfig() + submodule = Code2WavSubmodule(decoder, config) + codes = _codes(num_frames) + frames = torch.cat([codes, torch.full((1, 16), config.talker.codec_eos_token_id)]) if eos else codes + + steps = _stream(submodule, frames, schedule) + full = _to_pcm(decoder(codes.t()[None], torch.arange(num_frames)[None]))[0, 0] + + expected = _with_live_eos(steps, num_frames, full) if eos else full + assert torch.equal(torch.cat([audio for _, audio, _ in steps]), expected) + first_info, first_audio, _ = steps[0] + assert first_info.context_items == 0 and first_info.num_items == min(schedule[0], frames.shape[0]) + + +def test_the_served_edge_uses_the_configured_schedule(): + """The model wires ``codec_chunk_schedule`` into the Talker -> Code2Wav edge.""" + from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel + + model = Qwen3OmniModel.__new__(Qwen3OmniModel) + model.config = Qwen3OmniModelConfig() + (connection,) = [ + c for c in model.get_partition_topology().connections if c.edge_name == "codec_tokens" + ] + policy = connection.chunk_policy_factory() + assert isinstance(policy, ScheduledLeftContextChunkPolicy) + assert policy.window_size() == model.config.code2wav.codec_chunk_schedule[0] + + @pytest.mark.parametrize("num_frames", [1, 24, 25, 26, 49, 50, 51, 125]) def test_a_trailing_codec_eos_frame_adds_no_audio_and_no_decode(num_frames): """Served streams end with the Talker's codec_eos frame. A reply of k * 25 frames then ends with a flush of 25 context frames and the EOS frame, which emits nothing and decodes nothing; for k * 25 - 1 frames the - EOS frame is the last new item of a window and then a context item of the - flush.""" + EOS frame is the last new item of a live window and then a context item + of the flush. A live window decodes it as padding without reading it on + the host and emits one frame of silence for it (the last real frame's + look-ahead samples included, which a dropped frame never emits).""" decoder = _LookaheadDecoder() config = Qwen3OmniModelConfig() submodule = Code2WavSubmodule(decoder, config) @@ -177,7 +234,7 @@ def test_a_trailing_codec_eos_frame_adds_no_audio_and_no_decode(num_frames): assert decoder.calls == math.ceil(num_frames / 25) full = _to_pcm(decoder(codes.t()[None], torch.arange(num_frames)[None]))[0, 0] - assert torch.equal(torch.cat([audio for _, audio, _ in steps]), full) + assert torch.equal(torch.cat([audio for _, audio, _ in steps]), _with_live_eos(steps, num_frames, full)) flush_info, flush_audio, flush_decoded = steps[-1] assert flush_info.is_final if num_frames % 25 == 0: @@ -245,7 +302,6 @@ def test_mixed_batch_of_a_terminal_flush_and_a_live_window(): full = _to_pcm(decoder(request_codes.t()[None], torch.arange(request_codes.shape[0])[None]))[0, 0] streamed = torch.cat([audio for _, audio, _ in steps[rid]]) assert torch.equal(streamed, full), rid - assert not submodule._first_chunk_emitted # --------------------------------------------------------------------------- diff --git a/test/modular/test_qwen3_omni_fused_kernels.py b/test/modular/test_qwen3_omni_fused_kernels.py new file mode 100644 index 0000000000..522d25a9c4 --- /dev/null +++ b/test/modular/test_qwen3_omni_fused_kernels.py @@ -0,0 +1,124 @@ +"""Qwen3-Omni's fused kernels against the op sequences they replace (GPU only). + +* ``cp_qk_norm_rope_kv_write``: the code predictor's q/k norm + RoPE and its + dense-KV write, bit-for-bit what two ``fused_qk_norm_rope`` calls and two + strided copies produce. +* ``softmax_topk``: ``TopKRouter``'s softmax -> topk -> renormalize. +""" +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.mark.parametrize("bs,seq_len,cache_pos", [(1, 1, 0), (5, 1, 7), (32, 1, 15), (3, 2, 4)]) +def test_cp_prologue_matches_the_unfused_ops(bs, seq_len, cache_pos): + from mstar.model.qwen3_omni.components.cp_kernels import cp_qk_norm_rope_kv_write + from mstar.utils.attention import fused_qk_norm_rope + + torch.manual_seed(0) + hq, hkv, d, layers, max_bs, max_s = 16, 8, 128, 5, 32, 16 + dev, dt = "cuda", torch.bfloat16 + qkv = torch.randn(bs, seq_len, (hq + 2 * hkv) * d, device=dev, dtype=dt) + wq = torch.randn(d, device=dev, dtype=dt) + wk = torch.randn(d, device=dev, dtype=dt) + pos = torch.randint(0, 16, (bs, seq_len), device=dev) + cache = torch.randn(layers, max_bs, 2, max_s, hkv, d, device=dev, dtype=dt) + ref_cache, ref_qkv = cache.clone(), qkv.clone() + flat_pos = pos.reshape(-1) + + q, k, v = ref_qkv.split((hq * d, hkv * d, hkv * d), dim=-1) + q = fused_qk_norm_rope(q.view(bs * seq_len, hq, d), wq, pos=flat_pos, eps=1e-6, rope_theta=1e6) + k = fused_qk_norm_rope(k.view(bs * seq_len, hkv, d), wk, pos=flat_pos, eps=1e-5, rope_theta=1e6) + view = ref_cache[:, :bs] + view[2, :, 0, cache_pos:cache_pos + seq_len] = k.view(bs, seq_len, hkv, d) + view[2, :, 1, cache_pos:cache_pos + seq_len] = v.view(bs, seq_len, hkv, d) + + cp_qk_norm_rope_kv_write( + qkv.view(bs * seq_len, -1), wq, wk, flat_pos, cache[:, :bs][2], cache_pos, seq_len, + 1e-6, 1e-5, 1e6, hq, hkv, d, + ) + torch.testing.assert_close(qkv[..., :hq * d], ref_qkv[..., :hq * d], rtol=0, atol=0) + torch.testing.assert_close(cache, ref_cache, rtol=0, atol=0) + + +@pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) +@pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) +def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): + from mstar.utils.fused_moe.router import softmax_topk + + torch.manual_seed(1) + logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 + probs = torch.softmax(logits, dim=-1, dtype=torch.float) + ref_w, ref_ids = torch.topk(probs, top_k, dim=-1) + if norm: + ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) + + w, ids = softmax_topk(logits, top_k, norm) + assert ids.dtype == torch.int64 and w.dtype == torch.float32 + # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the + # selected probabilities, and that every pick is a distinct expert holding them + torch.testing.assert_close(w, ref_w, rtol=1e-6, atol=1e-7) + assert (ids.sort(dim=-1).values.diff(dim=-1) > 0).all() + picked = probs.gather(1, ids) + if norm: + picked = picked / picked.sum(dim=-1, keepdim=True) + torch.testing.assert_close(picked, w, rtol=1e-6, atol=1e-7) + # with no ties the choice and its order are torch's + distinct = (ref_w[:, :-1] != ref_w[:, 1:]).all(dim=-1) + assert torch.equal(ids[distinct], ref_ids[distinct]) + + +def test_router_compiles_with_the_fused_topk(): + from mstar.model.components.moe import TopKRouter + + router = TopKRouter(64, 128, 8, fused_softmax_topk=True).cuda().to(torch.bfloat16).requires_grad_(False) + torch.nn.init.normal_(router.weight) + x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + eager = router(x) + compiled = torch.compile(router, fullgraph=True)(x) + assert torch.equal(eager[1], compiled[1]) + torch.testing.assert_close(eager[0], compiled[0], rtol=0, atol=0) + + +@pytest.mark.parametrize("bs", [1, 8, 32]) +def test_cp_embed_accumulate_matches_the_unfused_ops(bs): + from mstar.model.qwen3_omni.components.cp_kernels import cp_embed_accumulate + + torch.manual_seed(2) + weight = torch.randn(2048, 1024, device="cuda", dtype=torch.bfloat16) + tokens = torch.randint(0, 2048, (bs,), device="cuda") + emb_sum = torch.randn(bs, 1024, device="cuda") + codes = torch.randint(0, 2048, (bs, 16), device="cuda") + ref_sum, ref_codes = emb_sum.clone(), codes.clone() + + ref_codes[:, 5] = tokens + ref_emb = torch.nn.functional.embedding(tokens, weight) + ref_sum.add_(ref_emb) + + emb = cp_embed_accumulate(tokens, weight, emb_sum, codes, 5) + assert torch.equal(emb, ref_emb) + assert torch.equal(emb_sum, ref_sum) + assert torch.equal(codes, ref_codes) + + + +@pytest.mark.parametrize("bs,seq,cache_pos", [(1, 1, 0), (4, 2, 0), (32, 1, 7), (3, 1, 15), (5, 2, 3)]) +def test_cp_attention_is_causal_attention_over_the_depth_cache(bs, seq, cache_pos): + from mstar.model.qwen3_omni.components.cp_kernels import cp_attention + + torch.manual_seed(4) + hq, hkv, d, depth = 16, 8, 128, 16 + cache = torch.randn(bs, 2, depth, hkv, d, device="cuda", dtype=torch.bfloat16) + q = torch.randn(bs, seq, hq, d, device="cuda", dtype=torch.bfloat16) + out = cp_attention(q, cache[:, 0], cache[:, 1], cache_pos) + for s in range(seq): + n = cache_pos + s + 1 + k = cache[:, 0, :n].float().repeat_interleave(hq // hkv, dim=2) # (bs, n, hq, d) + v = cache[:, 1, :n].float().repeat_interleave(hq // hkv, dim=2) + w = torch.softmax(torch.einsum("bhd,bnhd->bhn", q[:, s].float(), k) / d ** 0.5, -1) + ref = torch.einsum("bhn,bnhd->bhd", w, v) + torch.testing.assert_close(out[:, s].float(), ref, rtol=2e-2, atol=2e-2) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 4090cf6180..e1c3a97be1 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -41,7 +41,13 @@ from mstar.model.qwen3_tts.qwen3_tts_model import Qwen3TTSModel, clip_digest from mstar.model.qwen3_tts.submodules import CodecSubmodule, TalkerSubmodule from mstar.model.registry import HF_MODELS, get_model_class -from mstar.model.submodule_base import ARNodeInputs, InputMetadata, ModelInputsFromEngine +from mstar.model.submodule_base import ( + ARNodeInputs, + BatchedModelOutput, + HostRows, + InputMetadata, + ModelInputsFromEngine, +) from mstar.streaming.chunk_policy import ScheduledLeftContextChunkPolicy from mstar.streaming.stream_buffer import StreamBuffer, StreamChunkInfo @@ -843,7 +849,9 @@ def test_qwen3_tts_talker_builds_official_non_streaming_prefill(): SimpleNamespace(request_id="request", rid_handle="request"), {"talker_input_embeds": [torch.zeros(1, 16)]}, ) - assert torch.equal(prepared.input_embeds[0], state["tts_pad_embed"]) + # The condition rides the row; preprocess adds it for the batch. + assert torch.equal(prepared.kwargs["text_condition"], state["tts_pad_embed"]) + assert torch.equal(prepared.input_embeds[0] + prepared.kwargs["text_condition"], state["tts_pad_embed"]) def test_qwen3_tts_talker_prefill_prepends_instruction_without_speaker(): @@ -944,12 +952,13 @@ def test_qwen3_tts_clone_prefill_streams_reference_frames_first(): # (all a window can use as context) ahead of the first generated frame. reference = torch.arange(12).view(3, 4) submodule.request_state("clone").add("reference_frames", reference) - frame = torch.tensor([9, 9, 9, 9]) + frame = torch.tensor([[9, 9, 9, 9]]) # a [1, groups] row of the step's output items = submodule._codec_stream_items("talker_prefill_clone", "clone", frame) - assert [item.tolist() for item in items] == [[4, 5, 6, 7], [8, 9, 10, 11], [9, 9, 9, 9]] + # every item has the decode rows' [1, groups] shape, so a stream's items stack + assert [item.tolist() for item in items] == [[[4, 5, 6, 7]], [[8, 9, 10, 11]], [[9, 9, 9, 9]]] submodule.request_state("short").add("reference_frames", reference[:1]) assert [item.tolist() for item in submodule._codec_stream_items("talker_prefill_clone", "short", frame)] == [ - [0, 1, 2, 3], [9, 9, 9, 9], + [[0, 1, 2, 3]], [[9, 9, 9, 9]], ] # Only the clone prefill leads with the reference; decode never does. assert submodule._codec_stream_items("talker_decode", "clone", frame) == [frame] @@ -966,6 +975,7 @@ def test_qwen3_tts_prefill_frame_counts_toward_generation_limit(): submodule.request_state("request").add("generated_frames", 0) outputs = {"new_token": [torch.tensor(1)]} request_info = SimpleNamespace( + graph_walk="talker_decode", step_metadata={"talker_max_tokens": 1}, resource_configs={}, max_tokens=8192, @@ -990,6 +1000,7 @@ def test_qwen3_tts_stops_on_eos_from_routed_codec_tokens(): config.talker.codec_eos_token_id, 1, 2, 3 ])]} request_info = SimpleNamespace( + graph_walk="talker_decode", step_metadata={"talker_max_tokens": 100}, resource_configs={TALKER_SAMPLER: SimpleNamespace(ignore_eos=False)}, max_tokens=8192, @@ -1015,6 +1026,7 @@ def test_qwen3_tts_honors_ignore_eos_for_fixed_length_benchmarks(): config.talker.codec_eos_token_id )]} request_info = SimpleNamespace( + graph_walk="talker_decode", step_metadata={"talker_max_tokens": 100}, resource_configs={TALKER_SAMPLER: SimpleNamespace(ignore_eos=True)}, max_tokens=8192, @@ -1025,6 +1037,81 @@ def test_qwen3_tts_honors_ignore_eos_for_fixed_length_benchmarks(): assert submodule.check_stop("request", request_info, outputs) == set() +def test_qwen3_tts_batched_stop_check_agrees_with_the_per_request_one(): + """``check_stop_batched`` reads one host copy of the group-0 codes and + must stop exactly the requests ``check_stop`` stops: EOS (unless + ``ignore_eos``) or the frame limit, rows matched by request id.""" + config = _tiny_model_config() + submodule = TalkerSubmodule(Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config) + eos = config.talker.codec_eos_token_id + + def info(ignore_eos: bool, limit: int) -> SimpleNamespace: + return SimpleNamespace( + graph_walk="talker_decode", + step_metadata={"talker_max_tokens": limit}, + resource_configs={TALKER_SAMPLER: SimpleNamespace(ignore_eos=ignore_eos)}, + max_tokens=8192, + ) + + # rid: (code, ignore_eos, limit, frames already generated) + cases = {"eos": (eos, False, 100, 3), "ignored": (eos, True, 100, 3), + "limit": (5, False, 4, 3), "running": (5, False, 100, 3)} + infos = {rid: info(ignore, limit) for rid, (_, ignore, limit, _) in cases.items()} + for rid, (code, _, _, frames) in cases.items(): + submodule.request_state(rid).add("generated_frames", frames) + submodule.postprocess(rid, infos[rid], {"new_token": [torch.tensor(code)]}) + expected = { + rid: stop for rid in cases + if (stop := submodule.check_stop( + rid, infos[rid], {"layer0_codes": [torch.tensor([cases[rid][0]])]}, + )) + } + assert expected == {"eos": {"talker_decode_loop"}, "limit": {"talker_decode_loop"}} + # Rows in forward order (reversed here), plus a padding row past the real ones. + order = list(reversed(cases)) + codes = torch.tensor([cases[rid][0] for rid in order] + [eos]) + host_rows = HostRows(tuple(order), {"layer0_codes": codes}) + assert submodule.check_stop_batched(list(cases), infos, host_rows, "talker_decode") == expected + + +def test_qwen3_tts_talker_returns_row_outputs_except_for_the_clone_prefill(): + """A decode or plain prefill step is one frame per row, handed back as + batch tensors; the codec reads a ``[1, groups]`` row view as one frame.""" + config = _tiny_model_config() + submodule = TalkerSubmodule(Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config) + groups = config.talker.num_code_groups + frame = { + "talker_input_embeds": torch.zeros(2, 16), + "codec_tokens": torch.arange(2 * groups).view(2, groups), + "new_token": torch.tensor([3, 4]), + } + submodule._run_frame = lambda *args: frame + engine_inputs = ModelInputsFromEngine(request_ids=["a", "b"], per_request_info={}) + tensors = (torch.zeros(2, 16), torch.arange(2), torch.zeros(2, dtype=torch.bool)) + out = submodule.forward_batched("talker_decode", engine_inputs, *tensors) + assert isinstance(out, BatchedModelOutput) + assert out.row_outputs is frame + assert torch.equal(out.check_stop_buffers["layer0_codes"], frame["new_token"]) + # The clone prefill returns rows too; postprocess leads its codec items + # with the reference tail, so the captured forward serves it as well. + assert isinstance(submodule.forward_batched("talker_prefill_clone", engine_inputs, *tensors), BatchedModelOutput) + reference = torch.arange(3 * groups).view(3, groups) + submodule.request_state("a").add_all(reference_frames=reference, generated_frames=0) + row = {"codec_tokens": [frame["codec_tokens"][0:1]], "new_token": [frame["new_token"][0:1]]} + submodule.postprocess("a", SimpleNamespace(graph_walk="talker_prefill_clone"), row) + left = config.codec.left_context_frames + assert [item.shape for item in row["codec_tokens"]] == [(1, groups)] * (min(left, 3) + 1) + assert torch.equal(row["codec_tokens"][-1], frame["codec_tokens"][0:1]) + assert torch.equal(torch.cat(row["codec_tokens"][:-1]), reference[-left:]) + + codec = CodecSubmodule(_FakeCodecDecoder(4), config) + window = torch.arange(3 * groups).view(3, 1, groups) % 7 # three [1, groups] row views + prepared = codec.prepare_inputs( + "codec_chunk", SimpleNamespace(request_id="r", rid_handle="r"), {"codec_tokens": [window]}, + ) + assert prepared.kwargs["frames"] == 3 + + def test_qwen3_tts_suppresses_eos_for_official_minimum_frames(): config = _tiny_model_config() submodule = TalkerSubmodule( @@ -1106,7 +1193,7 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): ARNodeInputs( input_embeds=torch.zeros(1, 16), input_seq_len=1, - kwargs={"suppress_eos": True}, + kwargs={"suppress_eos": True, "text_condition": torch.zeros(16)}, ) for _ in range(2) ] @@ -1127,8 +1214,13 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): assert packed["suppress_eos"].tolist() == [True, True] configs = {c.capture_graph_walk: c for c in submodule.get_accelerator_graph_configs(torch.device("cpu"))} graph_config = configs["talker_decode"] - assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 32] - assert graph_config.single_request_inputs.kwargs == {"suppress_eos": True} + # 64 rows a step: the top bucket is also the Talker's batch cap + assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 24, 32, 48, 64] + assert submodule.max_batch_size("talker_decode") == 64 + template_kwargs = graph_config.single_request_inputs.kwargs + assert set(template_kwargs) == {"suppress_eos", "text_condition"} + assert template_kwargs["suppress_eos"] is True + assert torch.equal(template_kwargs["text_condition"], torch.zeros(16)) assert not graph_config.single_request_inputs.tensor_inputs # Prefill replays a packed capture: token buckets, padding rows shaped like # a prepared prefill (embeds + the dynamic EOS-suppression key). @@ -1138,7 +1230,8 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): padding = prefill_config.make_node_input(7) assert padding.input_embeds.shape == (7, 16) and padding.input_seq_len == 7 assert padding.kwargs == {"suppress_eos": True} and not padding.tensor_inputs - assert prefill_config.replay_graph_walks == ["talker_prefill"] # the clone prefill stays eager + # the clone prefill is the same forward: it replays the plain prefill's captures + assert prefill_config.replay_graph_walks == ["talker_prefill", "talker_prefill_clone"] prefill_batch = ExecutingBatch( node_name="Talker", step_context=_step_context("talker_prefill", ["a", "b"]), per_request_input_tensors={}, per_request_info=info, @@ -1148,7 +1241,7 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): node_name="Talker", step_context=_step_context("talker_prefill_clone", ["a", "b"]), per_request_input_tensors={}, per_request_info=info, ) - assert not submodule.can_use_accelerator_graphs(clone_batch, model_inputs) + assert submodule.can_use_accelerator_graphs(clone_batch, model_inputs) # Residual sampling params live in per-request sampler buffers, so requests # that disagree about them still batch AND still replay the decode graph. # (They used to fall out of both.) @@ -1201,7 +1294,11 @@ def decode_row(rid: str, generated_frames: int) -> ARNodeInputs: packed = submodule.preprocess("talker_decode", engine_inputs, rows) _assert_same_tensor(packed["suppress_eos"], _old_flags(flags)) _assert_same_tensor(packed["last_token_indices"], torch.arange(len(rows))) - _assert_same_tensor(packed["input_embeds"], torch.cat([row.input_embeds for row in rows])) + _assert_same_tensor( + packed["input_embeds"], + torch.cat([row.input_embeds for row in rows]) + + torch.stack([row.kwargs["text_condition"] for row in rows]), + ) def prefill_row(rid, walk, text_ids, layout, speaker_id, language_id, **inputs) -> ARNodeInputs: metadata = {PROMPT_LAYOUT_KEY: layout, SPEAKER_ID_KEY: speaker_id, LANGUAGE_ID_KEY: language_id} @@ -1257,6 +1354,17 @@ def prefill_row(rid, walk, text_ids, layout, speaker_id, language_id, **inputs) _assert_same_tensor(capture["suppress_eos"], torch.ones(3, dtype=torch.bool)) +class _IdentityNorm(torch.nn.Module): + """A CPU stand-in for the FlashInfer RMSNorm, fused residual add included.""" + + def forward(self, hidden_states): + return hidden_states + + def forward_residual(self, hidden_states, residual): + summed = hidden_states + residual + return summed, summed + + def test_qwen3_tts_code_predictor_projects_wider_talker_inputs(): """1.7B: Talker width 2048 vs predictor width 1024 -> biased projection on every depth input; 0.6B (equal widths) -> identity, no extra parameters.""" @@ -1280,11 +1388,11 @@ def test_qwen3_tts_code_predictor_projects_wider_talker_inputs(): ) for layer in predictor.model.layers: - layer.input_layernorm = torch.nn.Identity() - layer.post_attention_layernorm = torch.nn.Identity() + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() layer.self_attn.q_norm = torch.nn.Identity() layer.self_attn.k_norm = torch.nn.Identity() - predictor.model.norm = torch.nn.Identity() + predictor.model.norm = _IdentityNorm() import mstar.model.qwen3_tts.components.talker as talker_module original_rope = talker_module.apply_rope_pos_ids original_attn = talker_module.decode_attn_nhd @@ -1319,11 +1427,11 @@ def test_qwen3_tts_code_predictor_uses_decode_attn_nhd(monkeypatch): # RMSNorm and the Triton attention kernel are CUDA-only, so replace them # without changing attention geometry. for layer in predictor.model.layers: - layer.input_layernorm = torch.nn.Identity() - layer.post_attention_layernorm = torch.nn.Identity() + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() layer.self_attn.q_norm = torch.nn.Identity() layer.self_attn.k_norm = torch.nn.Identity() - predictor.model.norm = torch.nn.Identity() + predictor.model.norm = _IdentityNorm() rope_calls = [] def capture_rope(q, k, position_ids, rope_theta): @@ -1616,7 +1724,7 @@ def test_qwen3_tts_codec_declares_a_watermark_capture_per_schedule_length(): configs = submodule.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32, 1) frames = sorted({*config.codec.chunk_schedule, config.codec.chunk_frames}) assert sorted(configs) == sorted(f"watermark_{f * 1920}" for f in frames) - assert all(c.capture_batch_sizes == CodecSubmodule.CAPTURE_BATCH_SIZES for c in configs.values()) + assert all(c.capture_batch_sizes == CodecSubmodule.WATERMARK_BATCH_SIZES for c in configs.values()) def test_qwen3_tts_clone_reference_tail_is_the_first_windows_context(): @@ -1950,7 +2058,7 @@ def test_qwen3_tts_codec_batches_and_declares_accelerator_graphs(): for graph_config in graph_configs: assert graph_config.capture_graph_walk == "codec_chunk" assert set(graph_config.replay_graph_walks) == {"codec_chunk"} - assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 32] + assert graph_config.capture_batch_sizes == [1, 2, 3, 4, 5, 6, 8, 10, 12, 16, 20, 24, 28, 32] assert graph_config.single_request_inputs.tensor_inputs["codec_tokens"].shape == ( 4, graph_config.additional_key_info, ) diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index c6f876271e..078ed02085 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -322,6 +322,51 @@ class _Group: assert mgr._kwargs["head_dim"] == HEAD_DIM +def test_opaque_callable_stays_in_the_compiled_graph(): + """``opaque=True`` runs the attention as a custom op: a compiled region + keeps one graph around it (no break) and computes what the eager call does.""" + from mstar.engine.resources.convenience import RaggedAttentionCallable + + mgr = manager() + spans = [128, 96] + mgr.plan(step(*spans), ctx(["r0", "r1"])) + q, k, v = qkv(sum(spans)) + fn = RaggedAttentionCallable(mgr, "main", opaque=True) + + def region(q, k, v): + return fn(q * 1.0, k, v) + 1.0 + + torch._dynamo.reset() + explain = torch._dynamo.explain(region)(q, k, v) + assert explain.graph_break_count == 0 + compiled = torch.compile(region, fullgraph=True) + close(compiled(q, k, v), RaggedAttentionCallable(mgr, "main")(q, k, v) + 1.0) + + +def test_opaque_callable_built_inside_a_compiled_region(): + """A node that builds its opaque callables lazily on first forward does so + inside the compiled region; the handle must still name a registered callable.""" + from mstar.engine.resources.convenience import RaggedAttentionCallable + + mgr = manager() + spans = [128, 96] + mgr.plan(step(*spans), ctx(["r0", "r1"])) + q, k, v = qkv(sum(spans)) + fns: dict[int, RaggedAttentionCallable] = {} + + def region(q, k, v): + fn = fns.get(id(mgr)) + if fn is None: + fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main", opaque=True) + return fn(q, k, v) + 1.0 + + torch._dynamo.reset() + compiled = torch.compile(region) + expected = RaggedAttentionCallable(mgr, "main")(q, k, v) + 1.0 + for _ in range(2): + close(compiled(q, k, v), expected) + + def test_plan_then_run_matches_reference(): mgr = manager() spans = [128, 96] diff --git a/test/modular/test_ragged_attention_cpu.py b/test/modular/test_ragged_attention_cpu.py index d56d1dde89..10aeee7274 100644 --- a/test/modular/test_ragged_attention_cpu.py +++ b/test/modular/test_ragged_attention_cpu.py @@ -1,4 +1,5 @@ -"""Ragged attention head-dimension padding checks that need no GPU.""" +"""Ragged attention checks that need no GPU: head-dimension padding and head +sharding.""" import sys from types import SimpleNamespace @@ -114,3 +115,23 @@ def test_mixed_wrapper_builds_both_halves_in_graph_mode(monkeypatch): assert mixed._prefill.batch_size == 3 assert decode.call_args.kwargs["use_cuda_graph"] is True assert prefill.call_args.kwargs["use_cuda_graph"] is True + + +@pytest.mark.parametrize("shard_heads,expected", [(True, 2), (False, 4)]) +def test_shard_heads_false_keeps_every_head(shard_heads, expected): + """A node whose ranks split something other than heads (guidance branches) + keeps the full head count on every rank.""" + from mstar.engine.resources import RaggedAttentionConfig, RaggedAttentionSpec + from mstar.engine.resources.base import EngineResourceInfo, build_resource + + spec = RaggedAttentionSpec( + resource_key="ragged", nodes={"encoder"}, + config=RaggedAttentionConfig(num_qo_heads=4, num_kv_heads=4, head_dim=72, shard_heads=shard_heads), + ) + mgr = build_resource( + spec, + EngineResourceInfo( + device=torch.device("cpu"), kv_dtype=torch.bfloat16, joint_comm_group=SimpleNamespace(world_size=2), + ), + ) + assert mgr._kwargs["num_qo_heads"] == expected diff --git a/test/modular/test_serving_metrics.py b/test/modular/test_serving_metrics.py index 29a670c2c4..389a843df3 100644 --- a/test/modular/test_serving_metrics.py +++ b/test/modular/test_serving_metrics.py @@ -1349,7 +1349,7 @@ def run(self): monkeypatch.setattr(entrypoint, "APIServer", _Server) monkeypatch.setattr(entrypoint.mp, "get_context", lambda method: types.SimpleNamespace(Process=lambda **kw: _Conductor())) - monkeypatch.setattr(entrypoint, "_run_warmup_requests", lambda server, specs: events.append("warmup")) + monkeypatch.setattr(entrypoint, "_run_warmup_requests", lambda server, specs, copies=1: events.append("warmup")) monkeypatch.setattr(entrypoint.uvicorn, "Server", _Uvicorn) monkeypatch.setattr(entrypoint, "_shutdown_conductor_process", lambda proc: None) diff --git a/test/modular/test_shm_tensor_comm.py b/test/modular/test_shm_tensor_comm.py index db53d6a7fa..cc39668f3f 100644 --- a/test/modular/test_shm_tensor_comm.py +++ b/test/modular/test_shm_tensor_comm.py @@ -907,3 +907,31 @@ class _DeviceTransport(SharedMemoryCommunicationManager): ) for uuid in [infos["tok"][0].uuid, *stored.flat_uuids]: assert mgr.tensor_store.get_cpu_tensor(uuid) is None + + +def test_a_held_output_is_sent_from_its_own_tensor_not_the_reused_host_copy(): + """A step's host copies are views of the worker's pinned buffers, which the + next step overwrites. An output sent in a later step (a loop's accumulated + output) must not be read from them: once the step releases them, the send + falls back to the stored tensor.""" + with tempfile.TemporaryDirectory() as tmpdir: + sender = _make_manager(tmpdir, entity_id="worker_0", request_id="req1") + receiver = _make_manager(tmpdir, entity_id="worker_1", request_id="req1") + produced = torch.randn(4, 8) + pinned = produced.clone() # the stop check's host copy of this step + stored = sender.store_and_return_tensor_info_batch( + ["req1"], {"req1": {"out": [produced]}}, ["out"], + cpu_tensors={"req1": {"out": [pinned]}}, + ) + uuid = stored.flat_uuids[0] + assert sender.tensor_store.get_cpu_tensor(uuid) is pinned + sender.release_host_copies(stored.flat_uuids) + pinned.fill_(0) # the next step reuses the buffer + + edge = GraphEdge(next_node="client", name="out") + edge.tensor_info = [sender.tensor_store.get_info(uuid)] + sender.increment_ref(uuid, 1) + sender.register_for_send("req1", edge.tensor_info) + receiver.start_read_tensors("req1", [edge], graph_walk="w") + assert "req1" in receiver.get_ready_tensors(graph_walk="w") + assert torch.equal(receiver.get_tensor(uuid), produced) diff --git a/test/modular/test_startup_warmup.py b/test/modular/test_startup_warmup.py index 146df94cb6..b948b84db0 100644 --- a/test/modular/test_startup_warmup.py +++ b/test/modular/test_startup_warmup.py @@ -147,3 +147,50 @@ def fake_download(repo_id, filename): # a malformed reference is a failed warmup, not a crash entrypoint._run_warmup_requests(fake, [{"files": ["hf://nvidia.jpg"], "output_modalities": ["video"]}]) assert len(fake.submitted) == 1 + + +def test_warmup_reaches_every_data_parallel_replica() -> None: + """One concurrent copy of each warmup per replica, so the conductor's + least-loaded placement warms every replica, not a random one.""" + config = {"node_groups": [ + {"node_names": ["dit", "vae_decoder"], "ranks": [0, 1]}, + {"node_names": ["text_encoder"], "ranks": [0]}, + ]} + assert entrypoint._data_parallel_replicas(config) == 2 + # a lockstep instance is one replica + assert entrypoint._data_parallel_replicas({"node_groups": [{"ranks": [0, 1], "sp_size": 2}]}) == 1 + assert entrypoint._data_parallel_replicas({"node_groups": [{"ranks": [0, 1, 2, 3], "tp_size": 2}]}) == 2 + fake = _FakeServer() + entrypoint._run_warmup_requests(fake, [{"text": "a"}, {"text": "b"}], copies=2) + assert [s["request_id"] for s in fake.submitted] == ["warmup-0", "warmup-0-1", "warmup-1", "warmup-1-1"] + + +def test_concurrency_runs_copies_together_for_batched_shapes() -> None: + """``concurrency: N`` submits N copies at once (times the replica count), so + a step compiled per batch size sees that batch before the bind.""" + fake = _FakeServer() + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 3}], copies=2) + assert [s["request_id"] for s in fake.submitted] == [f"warmup-0{s}" for s in ("", "-1", "-2", "-3", "-4", "-5")] + fake = _FakeServer() + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}]) + assert fake.submitted == [] + + +def test_a_failing_copy_lets_its_siblings_finish(caplog) -> None: + """One copy failing early is reported after the others end, rather than + cancelling them mid-request.""" + import asyncio + + class _SlowServer(_FakeServer): + async def collect_results(self, request_id, raw_request=None): + if request_id == "warmup-0": + raise RuntimeError("compile blew up") + await asyncio.sleep(0.01) + self.collected.append(request_id) + return [] + + fake = _SlowServer() + with caplog.at_level(logging.WARNING): + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 3}]) + assert sorted(fake.collected) == ["warmup-0-1", "warmup-0-2"] + assert "warmup request 1/1 failed" in caplog.text diff --git a/test/modular/test_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index 990b0fb03f..53de2d96cd 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -148,7 +148,7 @@ def _omnivoice_inputs(cfg, k, target_len=6, clone=False): inputs = { "prefix_ids": [torch.arange(c * 5).view(c, 5)], "prefix_audio_mask": [torch.zeros(5, dtype=torch.bool)], - "target_len": [torch.tensor([target_len])], + "initial_canvas": [torch.full((c, target_len), cfg.audio_mask_id)], "audio_tokens": [tokens], "step_index": [torch.tensor([k], dtype=torch.int64)], } @@ -172,20 +172,23 @@ def prepare(fwd_info, inp): assert prepare(fwd(UNMASK_LOOP_NAME, steps), _omnivoice_inputs(cfg, steps)) is None -def test_omnivoice_first_step_prep_unchanged(): - # iteration 0 still reads target_len once to size the canvas +def test_omnivoice_first_step_prep_reads_no_tensor(): + # iteration 0 starts from the data worker's all-MASK canvas: its shape is + # the target length, so nothing is read on the host sub, cfg = _omnivoice() inputs = _omnivoice_inputs(cfg, 0) del inputs["audio_tokens"], inputs["step_index"] - ni = sub.prepare_inputs("speech_gen", _fwd("r", "speech_gen", _omnivoice_meta(4)), inputs) + with _HostReads() as reads: + ni = sub.prepare_inputs("speech_gen", _fwd("r", "speech_gen", _omnivoice_meta(4)), inputs) + assert reads.reads == [], reads.reads assert ni.tensor_inputs["audio_tokens"].shape == (cfg.num_audio_codebook, 6) assert (ni.tensor_inputs["audio_tokens"] == cfg.audio_mask_id).all() assert ni.input_seq_len == 5 + 2 * 6 def test_omnivoice_postprocess_reads_no_tensor(): - """The unmask step's postprocess takes the iteration from the loop counter: - same canvas as reading the step edge, and no host read.""" + """The unmask step's postprocess takes the iteration prepare_inputs read off + the loop counter: same canvas as reading the step edge, and no host read.""" from mstar.model.omnivoice.components.unmask import build_reveal_schedule sub, cfg = _omnivoice() @@ -199,19 +202,18 @@ def test_omnivoice_postprocess_reads_no_tensor(): for k in range(steps): inputs = _omnivoice_inputs(cfg, k, target_len) inputs["audio_tokens"] = [torch.full((cfg.num_audio_codebook, target_len), cfg.audio_mask_id)] - node_inputs = sub.prepare_inputs("speech_gen", _fwd("r", "speech_gen", _omnivoice_meta(steps)), inputs) logits = torch.randn(shape, generator=torch.Generator().manual_seed(100 + k)) results = {} for name, info in ( ("counter", _fwd("r", "speech_gen", _omnivoice_meta(steps), UNMASK_LOOP_NAME, k)), ("edge", _fwd("r", "speech_gen", _omnivoice_meta(steps))), ): + node_inputs = sub.prepare_inputs("speech_gen", info, inputs) outputs = {"c_logits": [logits.clone()], "u_logits": [logits.flip(-1)]} with _HostReads() as reads: sub.postprocess("r", info, outputs, inputs=node_inputs) results[name] = (outputs, reads.reads) - assert results["counter"][1] == [], f"step {k}: host reads {results['counter'][1]}" - assert results["edge"][1], "the edge path should read the step on the host" + assert reads.reads == [], f"step {k} ({name}): host reads {reads.reads}" assert torch.equal(results["counter"][0]["audio_tokens"][0], results["edge"][0]["audio_tokens"][0]) assert torch.equal(results["counter"][0]["step_index"][0], results["edge"][0]["step_index"][0]) diff --git a/test/modular/test_wan22_model.py b/test/modular/test_wan22_model.py index e6c580ee90..11d4b59ed9 100644 --- a/test/modular/test_wan22_model.py +++ b/test/modular/test_wan22_model.py @@ -108,9 +108,15 @@ def test_wan22_declares_expected_nodes(): assert walk_nodes == {"text_encoder", "vae_encoder", "dit", "vae_decoder"} -def test_wan22_declares_no_engine_resources(): - """Nothing here is cached across steps, so the engine builds nothing.""" - model = _make_model() +def test_wan22_declares_only_dit_attention_resources(): + """Nothing here is cached across steps: the only resource is the dit's + cacheless self-attention, and the sdpa backend declares none.""" + model = _make_model() + specs = model.get_node_resources() + assert [spec.resource_key for spec in specs] == [Wan22DitSubmodule.SELF_ATTN] + assert all(spec.nodes == {"dit"} for spec in specs) + assert all(not spec.config.shard_heads for spec in specs) + model.config = Wan22Config(attention_backend="sdpa") assert model.get_node_resources() == [] assert model.get_request_resource_configs(partition_fwd_args={}) == {} @@ -967,3 +973,51 @@ def test_wan22_config_echo_normalizes_bogus_decode_override(monkeypatch, caplog) line = next(r.getMessage() for r in caplog.records if r.getMessage().startswith("Wan2.2 request:")) assert "decode=auto" in line assert "banana" not in line + + +# ---------------------------------------------------------------------- +# dit attention layout and guidance-branch split +# ---------------------------------------------------------------------- + + +class _FakeSPGroup: + def __init__(self, world_size: int, rank: int = 0): + self.world_size = world_size + self.rank = rank + + +def _dit_info(guidance: float) -> CurrentForwardPassInfo: + info = _fwd_info(requested_steps=2, iter_count=1) + info.step_metadata["guidance_scale"] = guidance + return info + + +@pytest.mark.parametrize("world_size,guidance,branches", [ + (None, 5.0, 2), # one GPU: both guidance branches in one batch-2 forward + (None, 1.0, 1), # guidance off: the conditional branch alone + (2, 5.0, 1), # two ranks split the branches + (2, 1.0, 1), # nothing to split: both ranks run the one branch + (4, 5.0, 2), # more ranks than branches: every rank runs both +]) +def test_wan22_dit_branches_per_rank(world_size, guidance, branches): + group = _FakeSPGroup(world_size) if world_size else None + submodule = Wan22DitSubmodule(transformer=None, config=Wan22Config(), sp_group=group) + node_inputs = submodule.prepare_inputs( + graph_walk=Wan22Model.VIDEO_GEN_WALK, fwd_info=_dit_info(guidance), inputs=_dit_loop_inputs(k=1), + ) + layout = node_inputs.resource_step_info + # latents [1, 4, 2, 2, 2] patchify (1, 2, 2) to 2 * 1 * 1 tokens + assert layout.video_tokens == 2 and node_inputs.input_seq_len == 2 + assert layout.branches == branches + + +def test_wan22_dit_declares_one_segment_per_branch(): + submodule = Wan22DitSubmodule(transformer=None, config=Wan22Config()) + node_inputs = submodule.prepare_inputs( + graph_walk=Wan22Model.VIDEO_GEN_WALK, fwd_info=_dit_info(5.0), inputs=_dit_loop_inputs(k=1), + ) + step = submodule.declare_step(Wan22Model.VIDEO_GEN_WALK, ["r0"], [node_inputs]) + self_step = step.get(Wan22DitSubmodule.SELF_ATTN) + assert [(s.label, s.span) for s in self_step.segments] == [("video", 2), ("video", 2)] + sdpa = Wan22DitSubmodule(transformer=None, config=Wan22Config(attention_backend="sdpa")) + assert sdpa.declare_step(Wan22Model.VIDEO_GEN_WALK, ["r0"], [node_inputs]) is None diff --git a/test/modular/test_wan_vae_compile_options.py b/test/modular/test_wan_vae_compile_options.py index d5727ee118..28fd42b0bd 100644 --- a/test/modular/test_wan_vae_compile_options.py +++ b/test/modular/test_wan_vae_compile_options.py @@ -23,8 +23,7 @@ ) -@pytest.fixture(scope="module") -def vae(): +def _tiny_vae(): return AutoencoderKLWan( base_dim=8, decoder_base_dim=16, z_dim=8, dim_mult=[1, 2, 2, 2], num_res_blocks=1, is_residual=True, in_channels=12, out_channels=12, patch_size=2, scale_factor_spatial=16, @@ -32,6 +31,11 @@ def vae(): ).eval() +@pytest.fixture(scope="module") +def vae(): + return _tiny_vae() + + @pytest.fixture def compiles(monkeypatch): seen = [] @@ -64,6 +68,7 @@ def test_eager_decoder_step_compiles_nothing(vae, compiles): def test_cosmos3_vae_decode_nodes_compile_without_comprehensive_padding(vae, compiles, monkeypatch, warm): monkeypatch.delenv("COSMOS3_COMPILE_VAE", raising=False) cfg = Cosmos3Config() + cfg.vae_decode_fast = False # the stock compiled decode; the fast path is tested below cfg.windowed_decode_warm_cache = warm cfg.vae_decode_channels_last = False Cosmos3VAEDecoderSubmodule(vae, cfg) @@ -71,3 +76,11 @@ def test_cosmos3_vae_decode_nodes_compile_without_comprehensive_padding(vae, com # the stateless node, then the AR node's stateless decode and, when warm, its DecoderStep assert len(compiles) == (3 if warm else 2) assert all(_padding_off(kw) for kw in compiles) + + +def test_cosmos3_fast_decode_compiles_only_its_decoder_step(compiles, monkeypatch): + monkeypatch.delenv("COSMOS3_COMPILE_VAE", raising=False) + # Its own VAE: the fast path casts and converts the node's instance in place. + sub = Cosmos3VAEDecoderSubmodule(_tiny_vae(), Cosmos3Config(vae_decode_fast=True)) + assert sub._fast_step is not None + assert len(compiles) == 1 and _padding_off(compiles[0]) diff --git a/test/modular/test_worker_drain.py b/test/modular/test_worker_drain.py index 3dd33e38a7..e831659aee 100644 --- a/test/modular/test_worker_drain.py +++ b/test/modular/test_worker_drain.py @@ -192,6 +192,7 @@ def test_a_walk_the_cache_served_whole_is_still_routed(): ) tensors.cleanup_collectable = lambda *freed: None tensors.increment_ref_batch_uniform = tensors.dereference_batch_uniform = lambda *args: None + tensors.release_host_copies = lambda uuids: None # counts each rid's outputs and stores none of them tensors.store_and_return_tensor_info_batch = lambda rids, outputs, signals, **kwargs: StoredOutputs( flat_uuids=[], flat_rids=[], signal_idxs=[], diff --git a/test/modular/test_z_image_batch_invariance.py b/test/modular/test_z_image_batch_invariance.py index a13cd23017..872b21c88c 100644 --- a/test/modular/test_z_image_batch_invariance.py +++ b/test/modular/test_z_image_batch_invariance.py @@ -122,7 +122,9 @@ def _encode(subs, prompts, batched: bool) -> list[torch.Tensor]: rids = [f"t{i}" for i in range(len(prompts))] with torch.inference_mode(): if batched: - out = text.forward_batched(IMAGE_GEN_WALK, _eng(rids), **text.preprocess(IMAGE_GEN_WALK, _eng(rids), rows)) + kwargs = text.preprocess(IMAGE_GEN_WALK, _eng(rids), rows) + packed = text.forward_batched(IMAGE_GEN_WALK, _eng(rids), **kwargs) + out = text.unpack_packed_outputs(packed, rids, [], rows, {}) return [out[rid][TEXT_EMBEDS][0] for rid in rids] return [text.forward(IMAGE_GEN_WALK, _eng([rid]), **text.preprocess(IMAGE_GEN_WALK, _eng([rid]), [row])) [TEXT_EMBEDS][0] for rid, row in zip(rids, rows, strict=True)] diff --git a/test/omnivoice/test_omnivoice_backbone_graphs.py b/test/omnivoice/test_omnivoice_backbone_graphs.py new file mode 100644 index 0000000000..ef0d12f452 --- /dev/null +++ b/test/omnivoice/test_omnivoice_backbone_graphs.py @@ -0,0 +1,101 @@ +"""The captured backbone path: its attention layout and its numerics. + +``test_document_*`` / ``test_attention_step_*`` / ``test_node_resources_*`` + No GPU, no weights. Each request owns two documents (a zero-length + unconditional one with CFG off), and the step spreads them over its ids. + +``test_encode_head_matches_forward`` + GPU + checkpoint. ``encode`` + ``head`` -- what the captured region and the + eager fallback run -- against the reference's packed ``forward`` through + the same planned kernel: the same logits up to half-precision rounding. +""" + +import os + +import pytest +import torch + +from mstar.model.omnivoice.components.backbone import CanvasItem, build_packed_canvas +from mstar.model.omnivoice.config import OmniVoiceConfig +from mstar.model.omnivoice.submodules import ( + BACKBONE_ATTN, + OmniVoiceBackboneSubmodule, + document_pairs, +) + +MODEL_PATH = os.environ.get("OMNIVOICE_PATH", "k2-fsa/OmniVoice") +MASK_ID = 1024 +NUM_CODEBOOK = 8 + + +def test_document_pairs_keep_a_slot_per_unguided_request(): + assert document_pairs([(5, 3), (4, 0)]) == [5, 3, 4, 0] + + +def _spans(step): + return [(seg.request_id, seg.span) for seg in step.segments] + + +def test_attention_step_two_documents_per_request(): + step = OmniVoiceBackboneSubmodule._attention_step(["a", "b"], [5, 3, 4, 0]) + assert _spans(step) == [("a", 5), ("a", 3), ("b", 4), ("b", 0)] + assert not step.steps[BACKBONE_ATTN].causal + + +def test_attention_step_padded_replay_and_capture_layouts(): + # replay: two documents per real request, one zero per padding row + replay = OmniVoiceBackboneSubmodule._attention_step(["a", "b", "p0", "p1"], [5, 3, 4, 0, 0, 0]) + assert [n for _, n in _spans(replay)] == [5, 3, 4, 0, 0, 0] + assert [rid for rid, _ in _spans(replay)][:4] == ["a", "a", "b", "b"] + # capture: one length per row + capture = OmniVoiceBackboneSubmodule._attention_step(["p0", "p1"], [8, 8]) + assert _spans(capture) == [("p0", 8), ("p1", 8)] + + +def test_node_resources_follow_backbone_graphs(): + from mstar.model.omnivoice.omnivoice_model import OmniVoiceModel + + on = OmniVoiceModel("k2-fsa/OmniVoice", skip_weight_loading=True) + off = OmniVoiceModel("k2-fsa/OmniVoice", skip_weight_loading=True, backbone_graphs=False) + (spec,) = on.get_node_resources() + assert spec.resource_key == BACKBONE_ATTN and spec.config.max_segments_per_request == 2 + assert off.get_node_resources() == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_encode_head_matches_forward(): + pytest.importorskip("omnivoice.models.omnivoice") + from omnivoice.models.omnivoice import OmniVoice + from omnivoice.models.omnivoice_flashinfer import apply_flashinfer + + from mstar.model.omnivoice.components.backbone import OmniVoiceBackbone + + model = OmniVoice.from_pretrained(MODEL_PATH, dtype=torch.float16).eval().cuda() + apply_flashinfer(model, enable_cuda_graph=False) + backbone = OmniVoiceBackbone(model, torch.float16).eval() + + gen = torch.Generator().manual_seed(0) + items = [] + for rid, (prefix_len, target_len, guidance) in enumerate([(23, 40, 2.0), (31, 25, 0.0), (17, 52, 2.0)]): + tokens = torch.full((1, NUM_CODEBOOK, target_len), MASK_ID, dtype=torch.long) + tokens[..., ::3] = torch.randint(0, MASK_ID, tokens[..., ::3].shape, generator=gen) + items.append(CanvasItem( + request_id=str(rid), + prefix_ids=torch.randint(0, 1000, (NUM_CODEBOOK, prefix_len), generator=gen), + prefix_audio_mask=torch.zeros(prefix_len, dtype=torch.bool), + tokens=tokens, + guidance_scale=guidance, + )) + canvas = build_packed_canvas(items, OmniVoiceConfig().audio_mask_id, torch.device("cuda")) + + with torch.inference_mode(): + want = backbone(canvas) + hidden = backbone.encode( + canvas.packed_ids[0].t(), canvas.audio_mask[0], canvas.position_ids[0].to(torch.int32), + attn=model._fi_runner.wrapper, # planned for this canvas by forward + ) + got = backbone.head(hidden, canvas.tgt_index) + # the residual adds fold into the next RMSNorm (a float sum, not a rounded + # half-precision one), so the logits agree to rounding, not bit for bit + assert torch.allclose(got.float(), want.float(), atol=0.25, rtol=0.02) + assert (got.argmax(-1) == want.argmax(-1)).float().mean() > 0.95 diff --git a/test/omnivoice/test_omnivoice_fused_weights.py b/test/omnivoice/test_omnivoice_fused_weights.py new file mode 100644 index 0000000000..27eac84a99 --- /dev/null +++ b/test/omnivoice/test_omnivoice_fused_weights.py @@ -0,0 +1,31 @@ +"""The boot check for the fused weights ``apply_flashinfer`` leaves on each +layer, which ``OmniVoiceBackbone.encode`` reads.""" + +from types import SimpleNamespace + +import pytest + +from mstar.model.omnivoice.components.backbone import assert_flashinfer_patched + + +def _layer(**drop): + attn = {"_fi_w_qkv": 0, "_fi_qkv_split": (1, 1, 1), "_fi_rope_theta": 1e6} + mlp = {"_fi_w_gate_up": 0} + for name in drop: + attn.pop(name, None) + mlp.pop(name, None) + return SimpleNamespace(self_attn=SimpleNamespace(**attn), mlp=SimpleNamespace(**mlp)) + + +def _model(*layers): + return SimpleNamespace(llm=SimpleNamespace(layers=list(layers))) + + +def test_a_patched_model_passes(): + assert_flashinfer_patched(_model(_layer(), _layer())) + + +@pytest.mark.parametrize("name", ["_fi_w_qkv", "_fi_qkv_split", "_fi_rope_theta", "_fi_w_gate_up"]) +def test_a_missing_fused_weight_fails_at_boot(name): + with pytest.raises(RuntimeError, match=name): + assert_flashinfer_patched(_model(_layer(), _layer(**{name: True}))) diff --git a/test/omnivoice/test_omnivoice_parity.py b/test/omnivoice/test_omnivoice_parity.py index ff3020802d..f608db3bfd 100644 --- a/test/omnivoice/test_omnivoice_parity.py +++ b/test/omnivoice/test_omnivoice_parity.py @@ -612,7 +612,7 @@ def test_postprocess_reveals_without_touching_the_loop_edge(): inputs = NodeInputs(tensor_inputs={ "audio_tokens": edge, "step_index": torch.zeros(1, dtype=torch.int64), - }) + }, kwargs={"step": 0}) sub.postprocess("a", _fwd_info("a", 0), outputs, inputs) diff --git a/test/omnivoice/test_omnivoice_silence.py b/test/omnivoice/test_omnivoice_silence.py new file mode 100644 index 0000000000..bb5d768244 --- /dev/null +++ b/test/omnivoice/test_omnivoice_silence.py @@ -0,0 +1,46 @@ +"""The vectorized silence trim against the reference's pydub one: sample-exact.""" + +import numpy as np +import pytest + +from mstar.model.omnivoice.components.silence import remove_silence + +reference = pytest.importorskip("omnivoice.utils.audio") + +RATE = 24000 + + +def _speech(rng, seconds): + """Bursts of tone and noise separated by silences of varied length.""" + out = [] + while sum(len(x) for x in out) < seconds * RATE: + n = int(rng.uniform(0.05, 0.9) * RATE) + if rng.random() < 0.5: + t = np.arange(n) / RATE + out.append((rng.uniform(0.05, 0.6) * np.sin(2 * np.pi * rng.uniform(80, 400) * t)).astype(np.float32)) + else: + level = rng.choice([0.0, 1e-5, 2e-3, 3.2e-3, 1e-2]) + out.append((level * rng.standard_normal(n)).astype(np.float32)) + return np.concatenate(out)[None, : int(seconds * RATE) + int(rng.integers(0, 24))] + + +@pytest.mark.parametrize("seed", range(40)) +def test_matches_pydub(seed): + rng = np.random.default_rng(seed) + audio = _speech(rng, rng.uniform(0.05, 12.0)) + for mid, lead, trail in ((500, 100, 100), (300, 100, 300), (0, 100, 300)): + want = reference.remove_silence(audio, RATE, mid_sil=mid, lead_sil=lead, trail_sil=trail) + got = remove_silence(audio, RATE, mid_sil=mid, lead_sil=lead, trail_sil=trail) + assert got.dtype == want.dtype and np.array_equal(got, want), (seed, mid, got.shape, want.shape) + + +@pytest.mark.parametrize("audio", [ + np.zeros((1, 0), np.float32), + np.zeros((1, 7), np.float32), + np.zeros((1, RATE), np.float32), + np.full((1, 13), 0.5, np.float32), +]) +def test_edge_cases_match_pydub(audio): + want = reference.remove_silence(audio, RATE, mid_sil=500, lead_sil=100, trail_sil=100) + got = remove_silence(audio, RATE, mid_sil=500, lead_sil=100, trail_sil=100) + assert np.array_equal(got, want) diff --git a/test/waypoint/serve_rollout.py b/test/waypoint/serve_rollout.py index f0df330ff5..9e2e2b102b 100644 --- a/test/waypoint/serve_rollout.py +++ b/test/waypoint/serve_rollout.py @@ -347,7 +347,8 @@ def _dit_schedule(log_text: str, request_ids: set[str]) -> list[tuple[str, ...]] continue if not isinstance(batch, (list, tuple)): continue - filtered = tuple(rid for rid in batch if rid in request_ids) + base_ids = tuple(str(rid).split("#", 1)[0] for rid in batch) + filtered = tuple(rid for rid in base_ids if rid in request_ids) if filtered: scheduled.append(filtered) return scheduled @@ -385,7 +386,7 @@ def _vetoed_overshoots(log_text: str, request_ids: set[str]) -> dict[str, int]: if _OVERSHOOT_MARKER not in line: continue for rid in request_ids: - if f"(request {rid} runs" in line: + if f"(request {rid} runs" in line or f"(request {rid}#" in line: counts[rid] += 1 return counts @@ -415,7 +416,11 @@ def _execution_count_failure( def _cleaned_request_ids(log_text: str) -> set[str]: - return {line.split(_CLEANUP_MARKER, 1)[1].strip() for line in log_text.splitlines() if _CLEANUP_MARKER in line} + # The server logs the internal id, the client's id plus a "#" suffix. + return { + line.split(_CLEANUP_MARKER, 1)[1].strip().split("#", 1)[0] + for line in log_text.splitlines() if _CLEANUP_MARKER in line + } def _read_log_since(log_path: Path, offset: int) -> str: From 6b7dc9527a8a5a1a576bdaf4ff793e81c9dd531f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:15:13 -0700 Subject: [PATCH 02/59] higgs_audio: lay out the merged prefill as the one step it runs Main's KV layouts (#808) gave Higgs three spans for its old prefill_text/prefill_audio/prefill_text walks; this branch prefills the text and audio in one prefill_audio step, so a keyed request dropped its layout and debug asserts failed. Co-Authored-By: Claude Opus 5.5 --- mstar/model/higgs_audio/higgs_audio_model.py | 9 ++++----- test/modular/test_model_kv_layouts.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/mstar/model/higgs_audio/higgs_audio_model.py b/mstar/model/higgs_audio/higgs_audio_model.py index c1d463deea..8e50413968 100644 --- a/mstar/model/higgs_audio/higgs_audio_model.py +++ b/mstar/model/higgs_audio/higgs_audio_model.py @@ -431,12 +431,11 @@ def enc(s: str) -> list[int]: post_ids += enc("\n\n\n\n") audio = sum(self.config.encoder_output_length(int(n)) for n in audio_feature_lens.tolist()) + # One step writes the text around the audio and the audio together, so + # one span; it mixes text and audio, so nothing keys it. + prompt = len(pre_ids) + audio + len(post_ids) layout = LabelLayout( - spans=( - Span("ids", "prefill_text", len(pre_ids), len(pre_ids), "text_inputs"), - Span("digest", "prefill_audio", audio, audio, ("audio", 0)), - Span("ids", "prefill_text", len(post_ids), len(post_ids), "text_inputs"), - ), + spans=(Span("rows", "prefill_audio", prompt, prompt, "audio_features"),), growth=(("decode_loop", 1),), ) return ProcessPromptOutput({ diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index 5bf4d2b0b1..4e3f896b63 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -439,6 +439,23 @@ def test_higgs_caps_decode_where_its_stop_check_ends_it(max_tokens): ), "the cap and the stop check end decode on different iterations" +def test_higgs_lays_out_the_prefill_steps_it_schedules(): + from mstar.model.higgs_audio.config import KV_CACHE + + model = _higgs() + output = model.process_prompt( + "transcribe", ["audio"], ["text"], tensors={"audio_inputs": [torch.zeros(16000 * 35)]}, + ) + [layout] = output.metadata["kv_layout"][KV_CACHE].values() + tensors = output.new_input_tensors + schedule = model._build_prefill_schedule({name: tensors[name] for name in tensors}) + + assert [span.walk for span in layout.spans] == [walk for walk, _ in schedule] + audio = sum(model.config.encoder_output_length(int(n)) for n in tensors["audio_feature_lens"][0].tolist()) + text = sum(ids.numel() for ids in tensors["text_inputs"]) + assert sum(span.length for span in layout.spans) == text + audio + + @pytest.mark.parametrize("max_tokens", [1, 64]) def test_qwen3_5_caps_decode_where_its_stop_check_ends_it(max_tokens): from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel From 76a03ede7195fc899b6c67cdd0678ae211a4c732 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:29:28 -0700 Subject: [PATCH 03/59] configs: leave TP async scheduling off for Qwen3.5 TP2/TP4 The flip rested on measurements from before main's TP teardown-order and TP prefix-caching fixes; it moves to its own change after a GPU soak. Co-Authored-By: Claude Opus 5.5 --- configs/qwen3_5_27b_tp4.yaml | 4 ---- configs/qwen3_5_9b_tp2.yaml | 4 ---- 2 files changed, 8 deletions(-) diff --git a/configs/qwen3_5_27b_tp4.yaml b/configs/qwen3_5_27b_tp4.yaml index 06c4470a99..0f3bfb516f 100644 --- a/configs/qwen3_5_27b_tp4.yaml +++ b/configs/qwen3_5_27b_tp4.yaml @@ -3,10 +3,6 @@ max_seq_len: 32768 # GIL switch interval of the API server process (see docs/serving.rst): # token streaming keeps three of its threads busy api_switch_interval_s: 0.001 -# Speculate and pre-plan step N+1 of the TP LLM during step N, as single-GPU -# nodes do (see MSTAR_TP_ASYNC_SCHED); without it each step's schedule and -# plan run serially between the ranks' forwards. -tp_async_sched: true node_groups: - node_names: [vision_encoder] ranks: [0] diff --git a/configs/qwen3_5_9b_tp2.yaml b/configs/qwen3_5_9b_tp2.yaml index 0949a72d89..526d2eb82b 100644 --- a/configs/qwen3_5_9b_tp2.yaml +++ b/configs/qwen3_5_9b_tp2.yaml @@ -3,10 +3,6 @@ max_seq_len: 32768 # GIL switch interval of the API server process (see docs/serving.rst): # token streaming keeps three of its threads busy api_switch_interval_s: 0.001 -# Speculate and pre-plan step N+1 of the TP LLM during step N, as single-GPU -# nodes do (see MSTAR_TP_ASYNC_SCHED); without it each step's schedule and -# plan run serially between the ranks' forwards. -tp_async_sched: true node_groups: # Text, images and video. Colocated: the # encoder runs once per request, then idles. From f424c98a51d9d212247c5e4e31c12ea738d9a9b6 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:32:22 -0700 Subject: [PATCH 04/59] Share one pinned non-blocking H2D helper mstar.utils.device.to_device_nonblocking replaces the copies in Qwen3.5, the diffusion denoise loop and the Higgs audio encoder. The Omnilingual ASR wav2vec2 module keeps its inline copy: it is kept free of M* imports so the parity harness can run it in the reference's environment. Co-Authored-By: Claude Opus 5.5 --- .../model/components/diffusion/denoise_loop.py | 14 ++------------ mstar/model/higgs_audio/submodules.py | 8 ++++---- mstar/model/qwen3_5/submodules.py | 18 +++++------------- mstar/utils/device.py | 10 ++++++++++ 4 files changed, 21 insertions(+), 29 deletions(-) diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index d85944114f..9ddcc26e48 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -58,21 +58,11 @@ from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule -from mstar.utils.device import DeviceCapability +from mstar.utils.device import to_device_nonblocking logger = logging.getLogger(__name__) -def _to_device_no_sync(tensor: torch.Tensor, device: torch.device) -> torch.Tensor: - """A small host tensor onto ``device`` without a stream sync. A pageable ``.to(device)`` - waits for everything already queued on the stream (other requests' denoise steps), which - stalls the GPU thread under concurrency; a pinned non-blocking copy does not.""" - device = torch.device(device) - if tensor.device.type != "cpu" or not DeviceCapability.of(device).is_accelerator: - return tensor.to(device) - return tensor.pin_memory().to(device, non_blocking=True) - - class StepScalar(NamedTuple): """One per-step per-request scalar, sliced out of the request's schedule. @@ -382,7 +372,7 @@ def _step_scalar_sources( f"entries for {schedule.num_steps} steps, but " f"{scalar.source} has {tensor.shape[0]}." ) - staged.setdefault(scalar.source, _to_device_no_sync(tensor, device)) + staged.setdefault(scalar.source, to_device_nonblocking(tensor, device)) return staged def _carried_inputs( diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 1505f3f3b9..6723e9c73a 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -43,6 +43,7 @@ NodeSubmodule, device_loopback_enabled, ) +from mstar.utils.device import to_device_nonblocking logger = logging.getLogger(__name__) @@ -150,10 +151,9 @@ def preprocess( "valid_lengths": [inp.kwargs["valid_lengths"] for inp in inputs], } if padded: - frames = torch.tensor(widths, dtype=torch.long) - if out["audio_features"].is_cuda: - frames = frames.pin_memory() - out["valid_frames"] = frames.to(self.get_device(), non_blocking=True) + out["valid_frames"] = to_device_nonblocking( + torch.tensor(widths, dtype=torch.long), self.get_device(), + ) return out def _encode(self, features: torch.Tensor, valid_frames: torch.Tensor | None = None) -> torch.Tensor: diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 0792181286..96e742660b 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -76,7 +76,7 @@ device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides -from mstar.utils.device import DeviceCapability +from mstar.utils.device import DeviceCapability, to_device_nonblocking from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) @@ -97,14 +97,6 @@ def _host_grids( ] -def _to_device(t: torch.Tensor, device: torch.device) -> torch.Tensor: - """A host-built tensor onto the device without blocking the host: a copy - from pageable memory waits for the stream to drain.""" - if device.type == "cpu": - return t - return t.pin_memory().to(device, non_blocking=True) - - class LLMSubmodule(ARNodeSubmodule): # the decode walk reads its input ids off the sampler's slot masters reads_device_loopback = True @@ -321,7 +313,7 @@ def _vision_inputs( return ARNodeInputs( input_seq_len=input_embeds.shape[0], input_embeds=input_embeds, - custom_pos_ids=_to_device(torch.cat(pos_ids, dim=1), device), + custom_pos_ids=to_device_nonblocking(torch.cat(pos_ids, dim=1), device), tensor_inputs={ "mrope_advance": pos - start_pos, # Ids cannot ride `input_ids` — that is what `preprocess` keys @@ -822,9 +814,9 @@ def prepare_inputs( input_seq_len=num_patches, tensor_inputs={ "pixel_values": torch.cat(inputs["pixel_values"], dim=0), - "indices": _to_device(indices, device), - "weights": _to_device(weights, device), - "position_ids": _to_device( + "indices": to_device_nonblocking(indices, device), + "weights": to_device_nonblocking(weights, device), + "position_ids": to_device_nonblocking( vision_grid_position_ids(grid, merge, torch.device("cpu")), device, ), # not forward args: `declare_step` makes segments of the diff --git a/mstar/utils/device.py b/mstar/utils/device.py index c7d26b8776..efb3cd8546 100644 --- a/mstar/utils/device.py +++ b/mstar/utils/device.py @@ -54,6 +54,16 @@ def is_capturing(self) -> bool: ) +def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: + """A host tensor onto ``device`` without blocking the host. A copy from + pageable memory waits for everything already queued on the stream; a + pinned non-blocking one does not.""" + device = torch.device(device) + if t.device.type != "cpu" or not DeviceCapability.of(device).pin_memory: + return t.to(device) + return t.pin_memory().to(device, non_blocking=True) + + @lru_cache(maxsize=None) def _device_capability(device: torch.device) -> DeviceCapability: return DeviceCapability(device) From 50000661c66609b65d6664aa420296abbed53fc0 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:33:28 -0700 Subject: [PATCH 05/59] Share the device loop-back decode input ids on ARNodeSubmodule ARNodeSubmodule.loopback_input_ids reads a decode step's tokens off the sampler and overlays the rows that brought their own, with the row index uploaded through pinned memory. It replaces the copies in Qwen3.5 and Higgs Audio; Higgs's copy uploaded the index with a synchronous pageable copy on every request's first decode step. Co-Authored-By: Claude Opus 5.5 --- mstar/model/higgs_audio/submodules.py | 12 +++--------- mstar/model/qwen3_5/submodules.py | 24 ++++-------------------- mstar/model/submodule_base.py | 23 +++++++++++++++++++++++ 3 files changed, 30 insertions(+), 29 deletions(-) diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 6723e9c73a..672774bd0f 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -390,15 +390,9 @@ def preprocess( return {"input_ids": torch.cat([inp.input_ids for inp in inputs])} # device loop-back: each row's last sampled token, read by slot; # rows whose previous step was the prefill brought theirs - sampler: SamplerResource = engine_inputs.resources[SAMPLER] - ids = sampler.loopback_tokens(engine_inputs.request_ids) - routed = [i for i, inp in enumerate(inputs) if inp.input_ids is not None] - if routed: - ids.index_copy_( - 0, torch.tensor(routed, device=ids.device), - torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), - ) - return {"input_ids": ids} + return {"input_ids": self.loopback_input_ids( + engine_inputs.resources[SAMPLER], engine_inputs.request_ids, inputs, + )} return { "input_embeds": torch.cat( [inp.input_embeds for inp in inputs], dim=0 diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 96e742660b..c0c7598071 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -76,7 +76,7 @@ device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides -from mstar.utils.device import DeviceCapability, to_device_nonblocking +from mstar.utils.device import to_device_nonblocking from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) @@ -437,33 +437,17 @@ def preprocess( # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] t0 = perf_counter() if PHASE_PERIOD else 0.0 - ids = sampler.loopback_tokens(engine_inputs.request_ids) - # rows whose last step ran elsewhere brought their token - routed = [ - i for i, inp in enumerate(inputs) - if inp.input_ids is not None and inp.input_seq_len - ] if graph_walk != "decode": # a mixed step packs its decode rows' single tokens beside the # prompts: each row takes its own ids, or its loop-back token + ids = sampler.loopback_tokens(engine_inputs.request_ids) ids = torch.cat([ inp.input_ids.to(ids.dtype) if inp.input_ids is not None else ids[i:i + 1] for i, inp in enumerate(inputs) ]) - elif routed: - # Through pinned memory, non-blocking: a pageable - # ``torch.tensor(..., device=cuda)`` is a synchronous copy - # that waits for the queued step, and every request's first - # decode step after its prefill takes this path. - index = torch.tensor( - routed, dtype=torch.long, - pin_memory=DeviceCapability.of(ids.device).pin_memory, - ).to(ids.device, non_blocking=True) - ids.index_copy_( - 0, index, - torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), - ) + else: + ids = self.loopback_input_ids(sampler, engine_inputs.request_ids, inputs) out["input_ids"] = ids if PHASE_PERIOD: phase_record("engine.preprocess.loopback", perf_counter() - t0) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 2196695fa5..be9de759ab 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -13,6 +13,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep +from mstar.utils.device import to_device_nonblocking from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) @@ -31,6 +32,7 @@ class HostRows(NamedTuple): from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner from mstar.engine.engine import ExecutingBatch + from mstar.engine.resources.sampler.resource import SamplerResource from mstar.streaming.stream_buffer import StreamChunkInfo @@ -1041,3 +1043,24 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: # input name to tensor pass + + @staticmethod + def loopback_input_ids( + sampler: SamplerResource, request_ids: list[str], inputs: Sequence[ARNodeInputs], + ) -> torch.Tensor: + """A device loop-back decode step's input ids: each row's last sampled + token, read off the sampler on the device, except where a row brought + its own (its previous step ran elsewhere, e.g. the prefill).""" + ids = sampler.loopback_tokens(request_ids) + routed = [ + i for i, inp in enumerate(inputs) + if inp.input_ids is not None and inp.input_seq_len + ] + if routed: + # every request's first decode step after its prefill comes here, + # so the row index goes up without a blocking copy + ids.index_copy_( + 0, to_device_nonblocking(torch.tensor(routed, dtype=torch.long), ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + return ids From 1b27803f361451ce7c67c7a8f491a1c72bc0fb80 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:34:01 -0700 Subject: [PATCH 06/59] bagel: combine_cfg never reads a branch that sat the step out Each parallel-CFG node counts the step's branches from its own worker's loop counter, and a branch that sat out sent zeros_like(latents). Had a branch node counted fewer branches than combine_cfg (the 2-GPU config puts LLM_cfg_* on another worker), the combine would have used those zeros as an unconditional velocity. A node that sits out now sends a zero-row velocity, which the combine tells apart on the host: it drops the branches it does not read and fails the request if one it reads sat out. Co-Authored-By: Claude Opus 5.5 --- mstar/model/bagel/submodules.py | 46 ++++++++++++++++++-------- test/modular/test_bagel_cfg_prune.py | 49 ++++++++++++++++++++++++++-- 2 files changed, 79 insertions(+), 16 deletions(-) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index e91f98ea6c..c9fb9981c6 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -196,6 +196,16 @@ def cfg_branch_skipped(label: str, n_branches: int | None) -> bool: return n_branches is not None and CFG_BRANCH_LABELS.index(label) >= n_branches +def sat_out_velocity(latents: torch.Tensor) -> torch.Tensor: + """What a parallel-CFG node sends for a step it sits out: a zero-row velocity. + + Each node counts the step's branches from its own worker's loop counter, so + ``combine_cfg`` checks this on the host rather than trusting the counts agree; + a placeholder of the velocity's shape would pass for a computed one. + """ + return latents.new_empty((0, latents.shape[-1])) + + def combine_cfg_velocity( v_main: torch.Tensor, v_cfg_text: torch.Tensor, @@ -1719,9 +1729,8 @@ def _forward_image_gen_single_branch( "LLM_cfg_img": "v_cfg_img", }.get(self.node_name, "v_main") if cfg_branch_skipped(self._NODE_TO_CFG_LABEL.get(self.node_name, "main"), n_branches): - # sits this step out: combine_cfg never reads the velocity (see `cfg_branches`), - # but the edge still has to carry a tensor of the right shape - return {output_name: [torch.zeros_like(latents)]} + # sits this step out (`cfg_branches`); the edge still carries a tensor + return {output_name: [sat_out_velocity(latents)]} # See _forward_image_gen: tie the separately-passed sequence dims so # torch.compile's dynamic shapes can prove the adds below line up. @@ -2007,12 +2016,21 @@ def prepare_inputs( # worker device without loading an otherwise-unused model component. device = inputs["v_main"][0].device - result = { - "v_main": inputs["v_main"][0], - "v_cfg_text": inputs["v_cfg_text"][0], - "v_cfg_img": inputs["v_cfg_img"][0], - } - if "latents" not in inputs or len(inputs["latents"]) == 0: + first_iter = "latents" not in inputs or not inputs["latents"] + n_branches = self._cfg_branches(fwd_info, first_iter) + # The branches this step's combine reads; None reads all three (device-side interval test) + reads = CFG_BRANCH_LABELS[:n_branches] + result = {"v_main": inputs["v_main"][0]} + for label in CFG_BRANCH_LABELS[1:]: + v = inputs[f"v_{label}"][0] + if v.shape[0] > 0: + result[f"v_{label}"] = v + elif label in reads: + raise ValueError( + f"BAGEL parallel CFG: LLM_{label} sat out a step whose combine reads " + f"{list(reads)}; the nodes' loop counters disagree" + ) + if first_iter: H = fwd_info.step_metadata.get("height", 1024) W = fwd_info.step_metadata.get("width", 1024) result["latents"], result["time_index"] = _init_latents_and_time_index( @@ -2025,7 +2043,6 @@ def prepare_inputs( **result, } kwargs = dict(fwd_info.step_metadata) - n_branches = self._cfg_branches(fwd_info, first_iter="latents" not in inputs or not inputs["latents"]) if n_branches is not None: kwargs[CFG_BRANCHES] = n_branches return NodeInputs(tensor_inputs=result, kwargs=kwargs) @@ -2053,10 +2070,10 @@ def forward( graph_walk: str, engine_inputs: ModelInputsFromEngine, v_main: torch.Tensor, - v_cfg_text: torch.Tensor, - v_cfg_img: torch.Tensor, latents: torch.Tensor, time_index: torch.Tensor, + v_cfg_text: torch.Tensor | None = None, + v_cfg_img: torch.Tensor | None = None, cfg_text_scale: float = None, cfg_img_scale: float = None, cfg_renorm_type: str = None, @@ -2085,8 +2102,9 @@ def forward( # below line up under dynamic shapes. n_latent = latents.shape[0] torch._check(v_main.shape[0] == n_latent) - torch._check(v_cfg_text.shape[0] == n_latent) - torch._check(v_cfg_img.shape[0] == n_latent) + for v in (v_cfg_text, v_cfg_img): # None: that branch sat the step out + if v is not None: + torch._check(v.shape[0] == n_latent) torch._check(time_index.shape[0] == n_latent) # Compute timestep and step size diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py index 421e3207dc..e26a7cd7ea 100644 --- a/test/modular/test_bagel_cfg_prune.py +++ b/test/modular/test_bagel_cfg_prune.py @@ -295,11 +295,56 @@ def test_parallel_combine_reduces_to_the_old_combine(): for k in (0, N - 2): time_index = torch.full((16,), float(k)) n = _branches(k, **T2I) - old = node.forward("image_gen_cfg", None, *v, latents, time_index)["latents"][0] - new = node.forward("image_gen_cfg", None, *v, latents, time_index, cfg_branches=n)["latents"][0] + tensors = dict( + v_main=v[0], v_cfg_text=v[1], v_cfg_img=v[2], latents=latents, time_index=time_index, + ) + old = node.forward("image_gen_cfg", None, **tensors)["latents"][0] + new = node.forward("image_gen_cfg", None, **tensors, cfg_branches=n)["latents"][0] torch.testing.assert_close(new, old, rtol=0, atol=1e-6) +def _combine_inputs(k: int, sat_out: tuple[str, ...]): + """combine_cfg's prepare_inputs at step ``k`` of a t2i request, with the named + branch nodes having sat the step out.""" + from mstar.model.bagel.submodules import CombineCFGSubmodule, sat_out_velocity + + cfg = types.SimpleNamespace(cfg_prune=True, **{k_: _DEFAULTS[k_] for k_ in ( + "num_timesteps", "timestep_shift", "cfg_text_scale", "cfg_img_scale", + "cfg_interval", "cfg_renorm_type", "cfg_renorm_min", + )}) + latents = torch.randn(16, 8, generator=torch.Generator().manual_seed(0)) + inputs = { + name: [sat_out_velocity(latents) if name[2:] in sat_out else torch.randn(16, 8)] + for name in ("v_main", "v_cfg_text", "v_cfg_img") + } + inputs.update(latents=[latents], time_index=[torch.full((16,), float(k))]) + info = types.SimpleNamespace( + step_metadata={"requires_cfg": True}, + dynamic_loop_iter_counts={IMAGE_GEN_LOOP: k}, + ) + node = CombineCFGSubmodule(cfg) + return node, node.prepare_inputs("image_gen_cfg", info, inputs) + + +def test_combine_never_reads_a_branch_that_sat_the_step_out(): + """A branch node counts from its own worker's loop counter: if it sat out a step + whose combine reads its velocity, the request fails rather than combining a + placeholder as the unconditional velocity.""" + k_out = N - 2 + assert _branches(k_out, **T2I) == 1 + node, prepared = _combine_inputs(k_out, sat_out=("cfg_text", "cfg_img")) + assert set(prepared.tensor_inputs) == {"v_main", "latents", "time_index"} + out = node.forward("image_gen_cfg", None, **prepared.tensor_inputs, **prepared.kwargs) + assert out["latents"][0].shape == (16, 8) + + k_in = next(k for k in STEPS if k > 0 and _branches(k, **T2I) == 2) + # t2i never reads cfg_img, so it sitting out is consistent + _, prepared = _combine_inputs(k_in, sat_out=("cfg_img",)) + assert "v_cfg_img" not in prepared.tensor_inputs + with pytest.raises(ValueError, match="cfg_text sat out"): + _combine_inputs(k_in, sat_out=("cfg_text", "cfg_img")) + + def test_image_gen_cfg_walk_names_the_loop_too(): from mstar.model.bagel.bagel_model import BagelModel From eab410d16cf8154df6d0fad97d1f1cec0ecbb400 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:34:49 -0700 Subject: [PATCH 07/59] omnivoice: stage the packed canvas indices through to_device_nonblocking A shared mstar.utils.device helper replaces the backbone's local copy. Co-Authored-By: Claude Opus 5.5 --- mstar/model/omnivoice/components/backbone.py | 11 ++++------- mstar/utils/device.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index db80b0be2c..ed1907cf06 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -34,6 +34,8 @@ import torch.nn.functional as F from torch import nn +from mstar.utils.device import to_device_nonblocking + logger = logging.getLogger(__name__) # The private surface this integration rides on. Checked once at load: the @@ -232,17 +234,12 @@ def build_packed_canvas( uncond_rows.append(torch.arange(cursor, cursor + t_len)) cursor += t_len - def to_device(t: torch.Tensor) -> torch.Tensor: - if device.type == "cuda": - return t.pin_memory().to(device, non_blocking=True) - return t.to(device) - return PackedCanvas( packed_ids=torch.cat(id_parts, dim=-1).unsqueeze(0), audio_mask=torch.cat(mask_parts).unsqueeze(0), - position_ids=to_device(torch.cat(positions).unsqueeze(0)), + position_ids=to_device_nonblocking(torch.cat(positions).unsqueeze(0), device), doc_lens=doc_lens, - tgt_index=to_device(torch.cat(cond_rows + uncond_rows)), + tgt_index=to_device_nonblocking(torch.cat(cond_rows + uncond_rows), device), flat_target_total=flat_cursor, flat_uncond_total=flat_u_cursor, items=items, diff --git a/mstar/utils/device.py b/mstar/utils/device.py index c7d26b8776..871a4f0e24 100644 --- a/mstar/utils/device.py +++ b/mstar/utils/device.py @@ -57,3 +57,13 @@ def is_capturing(self) -> bool: @lru_cache(maxsize=None) def _device_capability(device: torch.device) -> DeviceCapability: return DeviceCapability(device) + + +def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: + """``t`` on ``device`` without a stream sync: a host tensor is pinned, then + copied asynchronously. A tensor already on ``device`` comes back as is.""" + if t.device == torch.device(device): + return t + if t.device.type == "cpu" and not t.is_pinned(): + t = t.pin_memory() + return t.to(device, non_blocking=True) From ea74a68c77ef5e1bc8728749e7aa33dcefeb2183 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:35:37 -0700 Subject: [PATCH 08/59] bagel: bound the device caches keyed by request shapes The image_gen shape inputs become a small LRU keyed by output size, and the ViT's cu_seqlens, a two-element tensor, is staged through to_device_nonblocking per image instead of cached per token count. Co-Authored-By: Claude Opus 5.5 --- mstar/model/bagel/submodules.py | 24 ++++++++++++++---------- test/modular/test_bagel_cfg_prune.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 10 deletions(-) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index c9fb9981c6..0a610705b5 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -47,6 +47,7 @@ NodeSubmodule, StackingMethod, ) +from mstar.utils.device import to_device_nonblocking logger = logging.getLogger(__name__) @@ -137,6 +138,9 @@ def vit_patches( # Absent means the old behaviour, all three whenever guidance is on. CFG_BRANCHES = "cfg_branches" +# Output sizes whose image_gen shape inputs an LLM node keeps on the device +_IMAGE_GEN_SHAPES_CACHED = 8 + def _f32(x: float) -> float: """``x`` rounded to float32, as a device compare against a Python scalar sees it.""" @@ -316,7 +320,6 @@ def __init__( # request, so the segment ceiling is the batch size. Raise only if that # changes — and raise the resource's `max_segments_per_request` with it. self._max_images_per_request = 1 - self._cu_seqlens: dict[tuple, torch.Tensor] = {} def prepare_inputs( self, @@ -360,13 +363,9 @@ def prepare_inputs( "not the count `vit_patches` gives a layout to key it by" ) - # built once per (length, device): `torch.tensor(..., device=)` syncs the stream, which - # held this request's prepare until the in-flight decode steps of others had finished - cu_seqlens = self._cu_seqlens.get((num_tokens, device)) - if cu_seqlens is None: - cu_seqlens = self._cu_seqlens[(num_tokens, device)] = torch.tensor( - [0, num_tokens], dtype=torch.int32, device=device - ) + # not `torch.tensor(..., device=)`: that syncs the stream, which held this request's + # prepare until the in-flight decode steps of others had finished + cu_seqlens = to_device_nonblocking(torch.tensor([0, num_tokens], dtype=torch.int32), device) return NodeInputs( tensor_inputs={ @@ -888,6 +887,7 @@ def __init__( self.eos_token_id = eos_token_id # image_gen runs only the CFG branches a step needs (`cfg_branches`) self._cfg_prune = config.cfg_prune + # least recently used first; bounded, since the output size is the client's self._image_gen_shape_cache: dict[tuple, dict[str, torch.Tensor]] = {} self._sentinel_ids: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} self._warned_loop_counter = False @@ -1126,15 +1126,19 @@ def _image_gen_shape_inputs(self, H: int, W: int, device) -> dict[str, torch.Ten step's ``prepare_inputs`` until the current step had finished. Read-only downstream. """ key = (H, W, str(device)) - cached = self._image_gen_shape_cache.get(key) + cache = self._image_gen_shape_cache + cached = cache.pop(key, None) if cached is None: pos = get_flattened_position_ids_extrapolate( H, W, self.config.latent_downsample, max_num_patches_per_side=self.config.max_latent_size, ).to(device) - cached = self._image_gen_shape_cache[key] = { + cached = { "vae_position_ids": pos, **self._get_text_vae_idxs(pos.shape[0] + IMAGE_SENTINELS, device), } + cache[key] = cached + if len(cache) > _IMAGE_GEN_SHAPES_CACHED: + del cache[next(iter(cache))] return cached def declare_step( diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py index e26a7cd7ea..6597393bc5 100644 --- a/test/modular/test_bagel_cfg_prune.py +++ b/test/modular/test_bagel_cfg_prune.py @@ -217,6 +217,24 @@ def test_preprocess_hands_the_forward_the_declared_count(n_branches): assert out["requires_cfg"] is True +def test_image_gen_shape_inputs_stay_bounded_across_output_sizes(): + """The output size is the client's, so the per-size device inputs are an LRU, not a dict + that grows with every size requested.""" + from mstar.model.bagel.submodules import _IMAGE_GEN_SHAPES_CACHED + + stub = _stub( + config=types.SimpleNamespace(latent_downsample=16, max_latent_size=64), + _image_gen_shape_cache={}, + _get_text_vae_idxs=lambda n, device: LLMSubmodule._get_text_vae_idxs(None, n, device), + ) + first = LLMSubmodule._image_gen_shape_inputs(stub, 256, 256, "cpu") + for i in range(3 * _IMAGE_GEN_SHAPES_CACHED): + LLMSubmodule._image_gen_shape_inputs(stub, 256, 256 + 16 * (i + 1), "cpu") + # the size in use stays cached while others come and go + assert LLMSubmodule._image_gen_shape_inputs(stub, 256, 256, "cpu") is first + assert len(stub._image_gen_shape_cache) == _IMAGE_GEN_SHAPES_CACHED + + def _request(k, first_iter, prune=True, **meta): cfg = types.SimpleNamespace( num_timesteps=N, timestep_shift=SHIFT, From 781c5a7f83bfc252e630935b6b880bce9497458c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:36:01 -0700 Subject: [PATCH 09/59] One token stop check for AR submodules, declared as TokenStop Qwen3.5, Higgs Audio, Orpheus and the Qwen3-TTS Talker each carried a check_stop and a check_stop_batched that had drifted apart (only Qwen3.5 had the in-order fast path). ARNodeSubmodule now implements both off a TokenStop declaration (loop, token buffer, stop ids, sampler, optional walk and generated-count key), and the four copies are gone. The tests that compare the batched and per-request checks now run the shared path. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 5 +- mstar/model/higgs_audio/submodules.py | 53 +---------- mstar/model/orpheus/submodules.py | 48 +--------- mstar/model/qwen3_5/submodules.py | 70 ++------------ mstar/model/qwen3_tts/submodules.py | 62 ++----------- mstar/model/submodule_base.py | 91 +++++++++++++++++++ test/modular/test_model_kv_layouts.py | 7 +- .../test_qwen3_5_check_stop_batched.py | 5 +- 8 files changed, 122 insertions(+), 219 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 4ac6a7b8af..212f22a688 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -952,7 +952,10 @@ methods are: ``check_stop(...) -> set[str]`` (optional) Runs off the GPU thread, and may read tensor values. Return the names of the ``Loop`` objects to stop, for example after seeing the EOS token. This is how a decode loop - terminates. + terminates. An AR submodule whose loop ends on a stop token or the token budget sets + ``self.token_stop = TokenStop(...)`` in ``__init__`` instead of writing one; that also + gives it ``check_stop_batched``, which checks a whole step off one host copy of its + tokens. Qwen3.5's ``LLMSubmodule`` is an example. ``cleanup_request(self, request_id)`` (optional) Free per-request state held inside the submodule when a request finishes: buffers, diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 672774bd0f..7c2e5caebe 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -37,10 +37,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, device_loopback_enabled, ) from mstar.utils.device import to_device_nonblocking @@ -269,6 +269,10 @@ def __init__(self, llm: nn.Module, config: HiggsAudioModelConfig): super().__init__() self.model = llm self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=frozenset(config.stop_token_ids), + sampler=SAMPLER, walk="decode", + ) def get_accelerator_graph_configs( self, device: torch.device, tp_world_size: int = 1, @@ -483,22 +487,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs or request_info.graph_walk != "decode": - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - if (not ignore_eos and token in self.config.stop_token_ids) or \ - decoded_tokens >= request_info.max_tokens: - return {"decode_loop"} - return set() - def step_is_input_free(self, graph_walk: str) -> bool: # a decode row is one token, like the capture template's return graph_walk == "decode" @@ -507,34 +495,3 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # the decode loop's client edge carries the token the stop check # already has on the host; prefill's client edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for the whole batch off one ``tolist``.""" - if step_walk is not None and step_walk != "decode": - return {} - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - stop_ids = self.config.stop_token_ids - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - if step_walk is None and info.graph_walk != "decode": - continue - if ( - values[i] in stop_ids - and not info.resource_configs[SAMPLER].ignore_eos - ) or info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens: - stops[rid] = {"decode_loop"} - return stops diff --git a/mstar/model/orpheus/submodules.py b/mstar/model/orpheus/submodules.py index 3194142765..75a10dfd3d 100644 --- a/mstar/model/orpheus/submodules.py +++ b/mstar/model/orpheus/submodules.py @@ -23,11 +23,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo @@ -53,6 +53,10 @@ def __init__( self.embed_tokens = language_model.model.embed_tokens self.lm_head = language_model.lm_head self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", + stop_ids=frozenset({config.stop_token_id}), sampler=SAMPLER, + ) # Orpheus only ever emits ```` ids (SNAC codes and the # speech/turn markers): the LM head computes the logits of rows # ``[custom_token_base_id, vocab)`` alone, about 18% of the head, and @@ -222,48 +226,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - eos_token_id = self.config.stop_token_id - ignore_eos= request_info.resource_configs[SAMPLER].ignore_eos - if (not ignore_eos and eos_token_id == token) or \ - (request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= request_info.max_tokens): - return {"decode_loop"} - return set() - - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every row off one host copy of the step's tokens.""" - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(-1).tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - eos = self.config.stop_token_id - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - if (values[i] == eos and not info.resource_configs[SAMPLER].ignore_eos) or ( - info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens - ): - stops[rid] = {"decode_loop"} - return stops - class SNACDecoderSubmodule(NodeSubmodule): """SNAC 24kHz streaming decoder submodule. diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index c0c7598071..1029481fa2 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -69,10 +69,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides @@ -154,6 +154,11 @@ def __init__( self.model = model self.config = config self.vision_config = vision_config + # only a decode row's token ends decode_loop; a prefill is not in it yet + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=config.stop_token_ids, + sampler=SAMPLER, walk="decode", + ) self._vision_sentinels: tuple[torch.Tensor, torch.Tensor] | None = None # ------------------------------------------------------------------ @@ -652,24 +657,6 @@ def _token_is_real(request_info: CurrentForwardPassInfo) -> bool: "last_prefill", False, ) - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - # only a decode row's stop can end decode_loop; a prefill is not in it yet - if request_info.graph_walk != "decode" or "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - hit_eos = not ignore_eos and token in self.config.stop_token_ids - out_of_budget = ( - request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - >= request_info.max_tokens - ) - return {"decode_loop"} if hit_eos or out_of_budget else set() - def step_is_input_free(self, graph_walk: str) -> bool: # decode rows are one token each, like the template; the prefill walks # feed `prefill_tokens` from the inputs @@ -680,51 +667,6 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # check already has on the host; prefill's new_token edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every request off one ``tolist``.""" - if step_walk is not None and step_walk != "decode": - return {} # as check_stop: only a decode row's stop ends decode_loop - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - rows = host_rows.request_ids - if len(rows) >= len(request_ids) and all( - a == b for a, b in zip(request_ids, rows, strict=False) - ): - # the common case: the batch in forward order, no row map needed - pairs = zip(request_ids, values, strict=False) - else: - row_of = {rid: i for i, rid in enumerate(rows)} - pairs = ( - (rid, values[i]) for rid in request_ids - if (i := row_of.get(rid)) is not None and i < len(values) - ) # no row, no output: the per-request path skips it too - stop_ids = self.config.stop_token_ids - stops: dict[str, set[str]] = {} - mixed = step_walk is None - for rid, value in pairs: - info = request_infos[rid] - # a mixed step reads each row's walk; a decode step need not - if mixed and info.graph_walk != "decode": - continue - # the set test first: a stop id is rare, the config read behind it - # goes through the wrapper - if ( - value in stop_ids - and not info.resource_configs[SAMPLER].ignore_eos - ) or ( - info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - >= info.max_tokens - ): - stops[rid] = {"decode_loop"} - return stops # The vision tower's block loop, as a piecewise capture region diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 35478346f1..c8f66ceb8e 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -85,11 +85,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo from mstar.utils.h2d import PinnedStager @@ -162,6 +162,12 @@ def __init__( self.talker_config = config.talker self.cp_config = config.talker.code_predictor self.num_codes = config.talker.num_code_groups + # ``postprocess`` counts the frames, the prefill's included + self.token_stop = TokenStop( + loop="talker_decode_loop", buffer="layer0_codes", + stop_ids=frozenset({config.talker.codec_eos_token_id}), + sampler=TALKER_SAMPLER, generated_key="generated_frames", + ) self._suppress_mask: torch.Tensor | None = None self._decode_last_token_indices: torch.Tensor | None = None self._flag_stager = PinnedStager(torch.bool, numel=self.MAX_BATCH_SIZE) @@ -796,60 +802,6 @@ def postprocess( "generated_frames", int(state.get("generated_frames", 0)) + 1 ) - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - """Stop the graph loop on codec EOS or the request frame limit. - - This callback runs off the GPU execution thread, so reading the sampled - token with ``item`` is allowed here rather than in ``postprocess``. - """ - if "layer0_codes" not in outputs: - return set() - token = int(outputs["layer0_codes"][0].item()) - generated = int( - self.request_state(request_id).get("generated_frames", 0) - ) - max_tokens = request_info.max_tokens - sampling_config = request_info.resource_configs.get(TALKER_SAMPLER) - ignore_eos = bool(getattr(sampling_config, "ignore_eos", False)) - reached_eos = ( - not ignore_eos and token == self.talker_config.codec_eos_token_id - ) - if reached_eos or generated >= max_tokens: - logger.debug("talker %s: stop after %d frames (eos=%s)", request_id, generated, reached_eos) - return {"talker_decode_loop"} - return set() - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every row off one host copy of the group-0 codes.""" - codes = host_rows.buffers.get("layer0_codes") - if not torch.is_tensor(codes) or codes.dim() == 0: - return None - values = codes.reshape(-1).tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - eos = self.talker_config.codec_eos_token_id - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - generated = int(self.request_state(rid).get("generated_frames", 0)) - ignore_eos = bool(getattr(info.resource_configs.get(TALKER_SAMPLER), "ignore_eos", False)) - if (values[i] == eos and not ignore_eos) or generated >= info.max_tokens: - stops[rid] = {"talker_decode_loop"} - return stops - def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Admit compatible prefill/decode requests to continuous batching. diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index be9de759ab..f46fad97f0 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -28,6 +28,28 @@ class HostRows(NamedTuple): buffers: dict[str, torch.Tensor] +class TokenStop(NamedTuple): + """When an AR submodule's sampled token ends its decode loop: on a stop + id (unless the request's sampler config sets ``ignore_eos``), or once + the request has generated ``max_tokens`` tokens. + + Declared as ``ARNodeSubmodule.token_stop``; ``check_stop`` and + ``check_stop_batched`` follow from it. + """ + + loop: str + # the output (and ``check_stop_buffers`` entry) holding the step's token + buffer: str + stop_ids: frozenset[int] + # resource key of the sampler whose request config carries ``ignore_eos`` + sampler: str + # the only walk whose rows can stop the loop; None for every walk + walk: str | None = None + # request-state key counting the tokens generated so far; None counts + # the loop's iterations, plus the token this step sampled + generated_key: str | None = None + + if TYPE_CHECKING: from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner @@ -980,6 +1002,75 @@ def cleanup_request(self, request_id: str): class ARNodeSubmodule(NodeSubmodule): + # Set (in ``__init__``, where the stop ids are known) by a submodule whose + # sampled token ends a dynamic loop; the stop checks below follow from it. + token_stop: TokenStop | None = None + + def check_stop( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + stop = self.token_stop + if stop is None: + return super().check_stop(request_id, request_info, outputs) + if (stop.walk is not None and request_info.graph_walk != stop.walk) or stop.buffer not in outputs: + return set() + token = outputs[stop.buffer][0].item() + return self._token_stops(((request_id, token),), {request_id: request_info}, False).get(request_id, set()) + + def check_stop_batched( + self, request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + stop = self.token_stop + if stop is None: + return super().check_stop_batched(request_ids, request_infos, host_rows, step_walk) + if stop.walk is not None and step_walk is not None and step_walk != stop.walk: + return {} + tokens = host_rows.buffers.get(stop.buffer) + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() + rows = host_rows.request_ids + if len(rows) >= len(request_ids) and all( + a == b for a, b in zip(request_ids, rows, strict=False) + ): + # the common case: the batch in forward order, no row map needed + pairs = zip(request_ids, values, strict=False) + else: + row_of = {rid: i for i, rid in enumerate(rows)} + pairs = ( + (rid, values[i]) for rid in request_ids + if (i := row_of.get(rid)) is not None and i < len(values) + ) # no row, no output: the per-request path skips it too + # a mixed step reads each row's walk; a step of one walk need not + return self._token_stops(pairs, request_infos, stop.walk is not None and step_walk is None) + + def _token_stops(self, pairs, request_infos, check_walk: bool) -> dict[str, set[str]]: + """Both stop checks' verdict on ``(rid, token)`` pairs.""" + stop = self.token_stop + stop_ids, loop, key = stop.stop_ids, stop.loop, stop.generated_key + stops: dict[str, set[str]] = {} + for rid, token in pairs: + info = request_infos[rid] + if check_walk and info.graph_walk != stop.walk: + continue + # the set test first: a stop id is rare, and the config read + # behind it goes through the info wrapper + if token in stop_ids and not info.resource_configs[stop.sampler].ignore_eos: + stops[rid] = {loop} + continue + generated = ( + info.dynamic_loop_iter_counts.get(loop, 0) + 1 if key is None + else int(self.request_state(rid).get(key, 0)) + ) + if generated >= info.max_tokens: + stops[rid] = {loop} + return stops + def split_inputs( self, graph_walk: str, diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index 4e3f896b63..11fd31c45f 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -461,15 +461,12 @@ def test_qwen3_5_caps_decode_where_its_stop_check_ends_it(max_tokens): from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel from mstar.model.qwen3_5.submodules import LLMSubmodule - llm = SimpleNamespace(config=SimpleNamespace(stop_token_ids=())) - - def check_stop(request_id, request_info, outputs): - return LLMSubmodule.check_stop(llm, request_id, request_info, outputs) + llm = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset())) cap = Qwen3_5DenseModel.get_loop_caps(object.__new__(Qwen3_5DenseModel), max_output_tokens=max_tokens) assert cap["decode_loop"] == _stops_after( - check_stop, "decode_loop", "decode", 2048, + llm.check_stop, "decode_loop", "decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, ), "the cap and the stop check end decode on different iterations" diff --git a/test/modular/test_qwen3_5_check_stop_batched.py b/test/modular/test_qwen3_5_check_stop_batched.py index 0441960842..f50b2d65bc 100644 --- a/test/modular/test_qwen3_5_check_stop_batched.py +++ b/test/modular/test_qwen3_5_check_stop_batched.py @@ -23,7 +23,7 @@ def _info(ignore_eos=False, done=0, max_tokens=100, walk="decode", last_prefill= def _sub(): - return SimpleNamespace(config=SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) + return LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) @pytest.mark.parametrize("shape", ["flat", "column"]) @@ -118,8 +118,7 @@ def test_rows_out_of_batch_order_still_map_by_request(): from mstar.model.qwen3_5.submodules import LLMSubmodule from mstar.model.submodule_base import HostRows - sub = LLMSubmodule.__new__(LLMSubmodule) - sub.config = SimpleNamespace(stop_token_ids={7}) + sub = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({7}))) rows = HostRows(("a", "b", "c"), {"new_token": torch.tensor([1, 7, 1])}) def info(max_tokens, iters=0): From b143d72ccdd720378e8dac318bac1dd297b71409 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:36:40 -0700 Subject: [PATCH 10/59] bagel: reject an unsupported output_format at intake encode_image's format check moves into image_output_format, which BAGEL's process_prompt calls for image output, so a bad value fails the request before the denoise instead of in postprocess after it. Co-Authored-By: Claude Opus 5.5 --- docs/clients.rst | 3 ++- mstar/model/bagel/bagel_model.py | 4 +++- mstar/model/components/diffusion/image_io.py | 18 +++++++++++++----- test/modular/test_bagel_postprocess.py | 9 ++++++++- 4 files changed, 26 insertions(+), 8 deletions(-) diff --git a/docs/clients.rst b/docs/clients.rst index cb81cc44bd..88ebe28240 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -282,7 +282,8 @@ Waypoint emits live RGB frame chunks and is not routed through the encoded-video Per-model notes: - **BAGEL** — chat returns text only; use ``/v1/images/generations`` and - ``/v1/images/edits`` for image output. + ``/v1/images/edits`` for image output, with ``output_format`` and ``output_compression`` + as for FLUX.2 klein below. - **FLUX.2 klein / Z-Image** — ``/v1/images/generations`` (and ``/v1/images/edits`` for klein) only: ``size`` as ``WxH`` (multiples of 16); ``seed``; ``num_inference_steps``, ``output_format`` (``png`` default, ``jpeg``, ``webp``), ``output_compression`` (JPEG / diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index f945cfe0d8..a90d8d9b4b 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -95,7 +95,7 @@ Span, prompt_spans, ) -from mstar.model.components.diffusion.image_io import encode_image +from mstar.model.components.diffusion.image_io import encode_image, image_output_format from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -633,6 +633,8 @@ def process_prompt( Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ + if "image" in output_modalities: + image_output_format(kwargs) # fail a bad output_format now, not after the denoise result: NameToTensorList = {} if prompt is not None: diff --git a/mstar/model/components/diffusion/image_io.py b/mstar/model/components/diffusion/image_io.py index b4f541d88b..a9088373c0 100644 --- a/mstar/model/components/diffusion/image_io.py +++ b/mstar/model/components/diffusion/image_io.py @@ -146,6 +146,18 @@ def deflate(index: int, part: bytes) -> bytes: OUTPUT_FORMATS = ("png", "jpeg", "webp") +def image_output_format(request_kwargs: dict | None) -> str: + """The request's ``output_format`` (default ``png``; ``jpg`` means ``jpeg``), or ValueError + for one :func:`encode_image` cannot write. A model calls it at intake too, so a bad value + fails the request before generation rather than after.""" + fmt = str((request_kwargs or {}).get("output_format") or "png").lower() + if fmt == "jpg": + fmt = "jpeg" + if fmt not in OUTPUT_FORMATS: + raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") + return fmt + + def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> bytes: """Encode a uint8 ``[3, H, W]`` image for the client per the OpenAI images knobs. @@ -155,11 +167,7 @@ def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> byt Everything else decodes to the pixels the VAE produced (PNG losslessly). """ kwargs = request_kwargs or {} - fmt = str(kwargs.get("output_format") or "png").lower() - if fmt == "jpg": - fmt = "jpeg" - if fmt not in OUTPUT_FORMATS: - raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") + fmt = image_output_format(kwargs) if fmt == "png": return uint8_to_png(image, compress_level=int(kwargs.get("png_compress_level", 1))) from PIL import Image diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py index 9bbea4ff4f..e731ad4cad 100644 --- a/test/modular/test_bagel_postprocess.py +++ b/test/modular/test_bagel_postprocess.py @@ -1,4 +1,5 @@ -"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path sent.""" +"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path +sent, and a format it cannot write fails at intake.""" import io import numpy as np @@ -19,3 +20,9 @@ def test_image_png_decodes_to_the_truncated_pixels(): # the old path: permute to HWC, scale, truncate to uint8 expected = (output[0].permute(1, 2, 0) * 255).to(torch.uint8).numpy() assert np.array_equal(decoded, expected) + + +def test_unsupported_output_format_fails_at_intake(): + model = BagelModel.__new__(BagelModel) + with pytest.raises(ValueError, match="output_format"): + model.process_prompt(None, [], ["image"], output_format="gif") From ed6b4a823eada4a8d7326f849c60a6294e24f97e Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:37:02 -0700 Subject: [PATCH 11/59] Cache NodeSubmodule.get_device in the base, cleared on a move get_device found a parameter by walking the module tree on every call, and prepare_inputs / preprocess call it per row per step. Qwen3-Omni patched that with a mixin for two submodules; the base now caches the device and _apply (.to(), to_empty()) clears it, so the mixin is gone. Co-Authored-By: Claude Opus 5.5 --- mstar/model/qwen3_omni/submodules.py | 19 ++----------------- mstar/model/submodule_base.py | 14 +++++++++++++- test/modular/test_submodule_alloc_dtype.py | 22 ++++++++++++++++++++++ 3 files changed, 37 insertions(+), 18 deletions(-) diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 38dde6e5f7..68334778c7 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -110,21 +110,6 @@ def talker_text_mask(input_ids: torch.Tensor, config: Qwen3OmniModelConfig) -> t # =================================================================== - -class _CachedDeviceMixin: - """``get_device`` once: the base walks the module tree for a parameter on - every call, and prepare_inputs / preprocess call it per row per step - (~6% of the Talker GPU thread's held-GIL samples at c=32). These - submodules never move once loaded.""" - - _device: torch.device | None = None - - def get_device(self): - if self._device is None: - self._device = super().get_device() - return self._device - - class AudioEncoderSubmodule(NodeSubmodule): """Thin wrapper around the HF Whisper-style audio encoder. @@ -341,7 +326,7 @@ def forward( # =================================================================== -class ThinkerSubmodule(_CachedDeviceMixin, ARNodeSubmodule): +class ThinkerSubmodule(ARNodeSubmodule): """Wraps the FlashInfer-based Thinker MoE transformer. Dispatches on graph_walk: @@ -1178,7 +1163,7 @@ def check_stop( # 4. TalkerSubmodule (ar engine) -- SECOND MOST COMPLEX # =================================================================== -class TalkerSubmodule(_CachedDeviceMixin, ARNodeSubmodule): +class TalkerSubmodule(ARNodeSubmodule): MAX_BATCH_SIZE = 32 def __init__( diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index f46fad97f0..3ea2d7f1d1 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -607,8 +607,20 @@ def request_state(self, request_id: str) -> PerRequestState: state = self.request_states[request_id] = PerRequestState() return state + # Cached: prepare_inputs and preprocess ask per row per step, and finding + # a parameter walks the module tree. ``_apply`` (``.to()``, ``.cuda()``, + # ``to_empty()``) clears it. + _device_cache: torch.device | None = None + def get_device(self): - return next(self.parameters()).device + device = self._device_cache + if device is None: + device = self._device_cache = next(self.parameters()).device + return device + + def _apply(self, fn, *args, **kwargs): + self._device_cache = None + return super()._apply(fn, *args, **kwargs) @abstractmethod def prepare_inputs( diff --git a/test/modular/test_submodule_alloc_dtype.py b/test/modular/test_submodule_alloc_dtype.py index 6e0294afdb..9fad26d508 100644 --- a/test/modular/test_submodule_alloc_dtype.py +++ b/test/modular/test_submodule_alloc_dtype.py @@ -53,3 +53,25 @@ def test_to_on_meta_is_in_place_and_zero_alloc(): assert ret is m assert all(p.is_meta for p in m.parameters()) assert all(p.dtype == torch.bfloat16 for p in m.parameters()) + + +def test_submodule_device_follows_materialisation(): + """``NodeSubmodule.get_device`` is cached; moving the module clears it.""" + from mstar.model.submodule_base import NodeSubmodule + + class _Sub(NodeSubmodule): + def __init__(self): + super().__init__() + with torch.device("meta"): + self.layer = nn.Linear(2, 2) + + def prepare_inputs(self, *args, **kwargs): + raise NotImplementedError + + def forward(self, *args, **kwargs): + raise NotImplementedError + + sub = _Sub() + assert sub.get_device().type == "meta" + sub.to_empty(device="cpu") + assert sub.get_device().type == "cpu" From c20f448bad72a2a88f4015bc4d350b004e1cd316 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:37:45 -0700 Subject: [PATCH 12/59] Refuse one_row_per_step on a piecewise config it means nothing for The flag only changes how a lease_before_step packed region captured at batch size 1 is leased; anywhere else it was silently ignored. The config now raises ValueError at construction. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/accelerator_graph_config.py | 12 +++++++++++- test/modular/test_piecewise_can_lease.py | 22 ++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/mstar/engine/accelerator_graph_config.py b/mstar/engine/accelerator_graph_config.py index 82bd359941..1134da1e89 100644 --- a/mstar/engine/accelerator_graph_config.py +++ b/mstar/engine/accelerator_graph_config.py @@ -301,7 +301,8 @@ class PiecewiseAcceleratorGraphConfig(ABC): can_lease: Callable[[list[NodeInputs]], bool] | None = None # The region replays a whole step as ONE row, however many requests the # step carries (a vision tower packing every request's images as segments - # of one ragged batch), so its slot is leased at batch size 1. + # of one ragged batch), so its slot is leased at batch size 1. Only for a + # packed `lease_before_step` region captured at batch size 1 alone. one_row_per_step: bool = False # static kwargs threaded into capture_fn (e.g. cond_tokens, is_causal) forward_kwargs: dict[str, Any] = field(default_factory=dict) @@ -316,6 +317,15 @@ class PiecewiseAcceleratorGraphConfig(ABC): def __post_init__(self): check_compile_mode(self.compile_mode) + if self.one_row_per_step and ( + self.get_config_type() is not PiecewiseConfigType.PACKED + or not self.lease_before_step + or self.capture_batch_sizes != [1] + ): + raise ValueError( + "one_row_per_step needs a packed lease_before_step region " + "with capture_batch_sizes=[1]" + ) @abstractmethod def get_config_type(self) -> PiecewiseConfigType: diff --git a/test/modular/test_piecewise_can_lease.py b/test/modular/test_piecewise_can_lease.py index 7a8b9f4ea2..d45207c80d 100644 --- a/test/modular/test_piecewise_can_lease.py +++ b/test/modular/test_piecewise_can_lease.py @@ -8,6 +8,8 @@ from types import SimpleNamespace +import pytest + from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner from mstar.engine.engine import Engine @@ -77,3 +79,23 @@ def real_runner(one_row): assert packed.resolved == [(1, 137)] and per_row.resolved == [(3, 137)] assert leases["packed"].bucket == (1, 137) and leases["packed"].slot == 0 assert leases["per_row"].bucket == (3, 137) + + +@pytest.mark.parametrize("batched, overrides", [ + (False, {"capture_batch_sizes": [1, 2]}), + (False, {"lease_before_step": False}), + (True, {}), +]) +def test_one_row_per_step_is_refused_where_it_means_nothing(batched, overrides): + from mstar.engine.accelerator_graph_config import PiecewiseBatchedConfig, PiecewisePackedConfig + + kwargs = { + "capture_fn": lambda inp: {}, "make_static_inputs": lambda shape: {}, + "lease_before_step": True, "one_row_per_step": True, "capture_batch_sizes": [1], + } + PiecewisePackedConfig(**kwargs, total_tokens=[8]) # the shape it is for + with pytest.raises(ValueError, match="one_row_per_step"): + if batched: + PiecewiseBatchedConfig(**kwargs, seq_len=8) + else: + PiecewisePackedConfig(**{**kwargs, **overrides}, total_tokens=[8]) From a9a8465f0f1c97a947b19dde344e38339191cd98 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:37:45 -0700 Subject: [PATCH 13/59] wan22: call the two-GPU dit CFG-parallel, which is what it is configs/wan22_sp2.yaml becomes wan22_cfg_parallel_2gpu.yaml, after bagel_cfg_parallel_2gpu.yaml: its two ranks each run one guidance branch, and nothing splits the sequence. The mechanism still rides the node's sp group; the docs, comments and the shard_heads description now say so. Co-Authored-By: Claude Opus 5.5 --- ..._sp2.yaml => wan22_cfg_parallel_2gpu.yaml} | 11 ++++++---- docs/models.rst | 9 +++++---- mstar/engine/resources/attn/ragged/config.py | 10 +++++----- mstar/model/wan22/submodules.py | 20 ++++++++++--------- mstar/model/wan22/wan22_model.py | 9 +++++---- 5 files changed, 33 insertions(+), 26 deletions(-) rename configs/{wan22_sp2.yaml => wan22_cfg_parallel_2gpu.yaml} (50%) diff --git a/configs/wan22_sp2.yaml b/configs/wan22_cfg_parallel_2gpu.yaml similarity index 50% rename from configs/wan22_sp2.yaml rename to configs/wan22_cfg_parallel_2gpu.yaml index d247cad0e5..e2b422bc86 100644 --- a/configs/wan22_sp2.yaml +++ b/configs/wan22_cfg_parallel_2gpu.yaml @@ -1,9 +1,10 @@ model: "wan22" -# See wan22.yaml. Two GPUs, latency-oriented: the dit runs on both, each rank -# computing one classifier-free-guidance branch at batch 1, and the two noise -# predictions are all-gathered every step, so a guided denoise step costs one -# branch's forward. The text encoder and both VAE passes run on rank 0. +# See wan22.yaml. CFG-parallel on two GPUs, latency-oriented: the dit runs on +# both, each rank computing one classifier-free-guidance branch at batch 1, and +# the two noise predictions are all-gathered every step, so a guided denoise +# step costs one branch's forward. The text encoder and both VAE passes run on +# rank 0. # For throughput under concurrency, wan22_dp2.yaml (one full replica per GPU) # avoids rank 1 idling through rank 0's VAE decode. max_seq_len: 512 @@ -26,4 +27,6 @@ node_groups: - node_names: - dit ranks: [0, 1] + # a two-rank lockstep instance; Wan2.2's dit uses it for the CFG branches, + # not to split the sequence sp_size: 2 diff --git a/docs/models.rst b/docs/models.rst index 4114265d2a..3f047b0a4e 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -1067,15 +1067,16 @@ tiles only when free VRAM at decode time is short of the untiled estimate 32 GiB card. No CUDA-graph capture, no continuous batching (one request per step) and no component offload. -**Two GPUs.** ``configs/wan22_sp2.yaml`` runs the dit as a two-rank lockstep -instance (``sp_size: 2``) that splits the classifier-free guidance batch: each -rank runs one branch at batch 1 and the two noise predictions are all-gathered +**Two GPUs.** ``configs/wan22_cfg_parallel_2gpu.yaml`` is CFG-parallel: it runs +the dit as a two-rank lockstep instance (``sp_size: 2``, though nothing splits +the sequence) that splits the classifier-free guidance batch: each rank runs one +branch at batch 1 and the two noise predictions are all-gathered every step, so a guided step costs one branch's forward (502 vs 1000 ms at 1280x704x81 on H200s; one request 13.4 s vs 23.3 s end to end at 20 steps). The text encoder and VAE run on rank 0. Without guidance both ranks compute the same step, correct but no faster. For throughput under concurrency, ``configs/wan22_dp2.yaml`` serves a whole replica per GPU instead: 5.2 videos/min -at 4 concurrent vs 4.7 for sp2 (sp2's rank 1 idles through rank 0's VAE +at 4 concurrent vs 4.7 CFG-parallel (whose rank 1 idles through rank 0's VAE decode), while one request takes the 1-GPU time. LTX-2.5 (``ltx2_5``) diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 22793d47b8..1fa1fdeb10 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -73,11 +73,11 @@ class RaggedAttentionConfig: backend: RaggedAttnBackend = RaggedAttnBackend.AUTO flashinfer_backend: str = "auto" - # Whether the node's lockstep instance splits the attention heads across its - # ranks (tensor parallelism, Ulysses sequence parallelism), so each rank's - # resource attends its slice. A node that parallelizes over something else -- - # each rank attending whole requests or guidance branches -- sets False and - # keeps every head on every rank. + # Whether the node's ranks split the attention heads. True (tensor + # parallelism, Ulysses sequence parallelism): each rank's resource is sized + # for its slice of the heads. False: every rank attends with all the heads, + # for a node whose ranks divide other work -- Wan2.2's CFG parallelism, one + # guidance branch per rank, is the one user. shard_heads: bool = True # Activation dtype of the q / k / v the node hands the kernel (and of its diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 505215edbf..13122031f3 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -340,8 +340,8 @@ class Wan22DitSubmodule(_SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule): """Dense 5B video DiT with the UniPC step run inline. One loop iteration is one batch-2 transformer forward (the positive and - negative prompt embeddings stacked on the batch axis; on a two-rank sp - instance, one branch per rank) plus one UniPC predictor/corrector update + negative prompt embeddings stacked on the batch axis; CFG-parallel on a + two-rank instance, one branch per rank) plus one UniPC predictor/corrector update from ``components/unipc.py``. Loop-carried edges, all float32 over the latent grid: @@ -385,16 +385,18 @@ def __init__(self, transformer: nn.Module, config: Wan22Config, sp_group=None): super().__init__() self.transformer = transformer self.config = config - # Two ranks split the CFG batch: each runs one branch at batch 1 and the - # predictions are all-gathered, so the step costs one branch's forward plus a - # latent-sized gather. Every rank then runs the same combine and UniPC step on - # replicated latents. Any other shape (guidance off, more ranks) computes the - # whole batch on every rank, which is correct but no faster. + # CFG parallelism: the node's lockstep group arrives as its sp group, but + # nothing splits the sequence. Two ranks split the CFG batch, each running + # one branch at batch 1, and the predictions are all-gathered, so the step + # costs one branch's forward plus a latent-sized gather. Every rank then runs + # the same combine and UniPC step on replicated latents. Any other shape + # (guidance off, more ranks) computes the whole batch on every rank, which is + # correct but no faster. self.sp_group = sp_group if sp_group is not None and sp_group.world_size > 1 else None if self.sp_group is not None and self.sp_group.world_size != 2: logger.warning( - "Wan2.2 dit: sp_size=%d; only 2 ranks split the CFG batch, so every rank " - "computes the whole step.", self.sp_group.world_size, + "Wan2.2 dit: sp_size=%d; CFG parallelism needs exactly 2 ranks, so every " + "rank computes the whole step.", self.sp_group.world_size, ) self._attn_fns: Wan22AttentionFns | None = None self._record_fp32_islands() diff --git a/mstar/model/wan22/wan22_model.py b/mstar/model/wan22/wan22_model.py index 58616b981c..9e09814e9e 100644 --- a/mstar/model/wan22/wan22_model.py +++ b/mstar/model/wan22/wan22_model.py @@ -144,7 +144,7 @@ def get_node_resources(self) -> list[NodeResourceSpec]: head_dim=cfg.attention_head_dim, # one segment per guidance branch max_segments_per_request=2, - # the dit's sp ranks split guidance branches, not heads + # CFG-parallel: the dit's two ranks split guidance branches, not heads shard_heads=False, dtype=torch.bfloat16, ), @@ -152,9 +152,10 @@ def get_node_resources(self) -> list[NodeResourceSpec]: ] def get_default_sharding_config(self) -> ShardingConfig: - # The dit accepts an sp group: with guidance on, each of two ranks runs one - # CFG branch and they all-gather the two noise predictions (see - # Wan22DitSubmodule). Signals stay replicated, so no shard_dim entries. + # CFG-parallel, through the sp slot: the dit's lockstep group (sp_size) is + # used for guidance, each of two ranks running one CFG branch and + # all-gathering the two noise predictions (see Wan22DitSubmodule). Nothing + # splits the sequence, and signals stay replicated: no shard_dim entries. return ShardingConfig( groups=[], tp_enabled_nodes=set(), shard_dim={}, sp_enabled_nodes={"dit"}, ) From 5023afdb16589dca6b064150bcb1a1db3002207e Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:38:56 -0700 Subject: [PATCH 14/59] Qwen3 text encoder: one rotary table per device, not per prompt length The tables were cached per (length, device), so every new prompt length added a device allocation that was never freed. The encoder now builds one table per device at max_sequence_length and slices it; a longer prompt is refused with ValueError. Evicting instead would free a table a captured forward still reads. Co-Authored-By: Claude Opus 5.5 --- .../components/diffusion/qwen3/encoder.py | 21 +++++++++++-------- .../diffusion/qwen3/weight_loading.py | 1 + 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/mstar/model/components/diffusion/qwen3/encoder.py b/mstar/model/components/diffusion/qwen3/encoder.py index 47a1a9c1f4..bc20f2754c 100644 --- a/mstar/model/components/diffusion/qwen3/encoder.py +++ b/mstar/model/components/diffusion/qwen3/encoder.py @@ -203,13 +203,14 @@ def __init__( rms_norm_eps: float, rope_theta: float, tap_layers: tuple[int, ...], + max_positions: int, ): super().__init__() if min(tap_layers) < 1: raise ValueError("tap layers are 1-indexed decoder layers (0 would be the raw embeddings)") self.tap_layers = tuple(tap_layers) - self.head_dim, self.rope_theta = head_dim, rope_theta - self._rotary_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + self.head_dim, self.rope_theta, self.max_positions = head_dim, rope_theta, max_positions + self._rotary_cache: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} self.embed_tokens = nn.Embedding(vocab_size, hidden_size) self.layers = nn.ModuleList( DecoderLayer( @@ -240,14 +241,16 @@ def forward( None reads it off ``attention_mask``, a host sync (and not graph-capturable). """ seq = input_ids.shape[1] - # cached per (length, device): building them copies inv_freq from the host, which a - # captured forward cannot do (and the tables never change) - key = (seq, input_ids.device) - rotary = self._rotary_cache.get(key) - if rotary is None: - rotary = self._rotary_cache[key] = qwen3_rotary_tables( - torch.arange(seq, device=input_ids.device), self.head_dim, self.rope_theta, + if seq > self.max_positions: + raise ValueError(f"{seq} prompt tokens, past the encoder's {self.max_positions}") + # one table per device at the longest length, sliced: building it copies inv_freq from + # the host, which a captured forward cannot do, and a table a capture read must outlive it + tables = self._rotary_cache.get(input_ids.device) + if tables is None: + tables = self._rotary_cache[input_ids.device] = qwen3_rotary_tables( + torch.arange(self.max_positions, device=input_ids.device), self.head_dim, self.rope_theta, ) + rotary = (tables[0][:seq], tables[1][:seq]) # No padding anywhere in the batch: HF's create_causal_mask returns None and attention # runs its is_causal path; mirror that rather than passing an all-True mask. if padded is None: diff --git a/mstar/model/components/diffusion/qwen3/weight_loading.py b/mstar/model/components/diffusion/qwen3/weight_loading.py index ef9285a501..3eb4808f53 100644 --- a/mstar/model/components/diffusion/qwen3/weight_loading.py +++ b/mstar/model/components/diffusion/qwen3/weight_loading.py @@ -51,4 +51,5 @@ def make_text_encoder(text_config: Qwen3EncoderConfig) -> Qwen3HiddenStateEncode rms_norm_eps=text_config.rms_norm_eps, rope_theta=text_config.rope_theta, tap_layers=text_config.hidden_state_layers, + max_positions=text_config.max_sequence_length, ) From 9a539bee7db1ea3f01139983f0b492655006f893 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:40:39 -0700 Subject: [PATCH 15/59] conductor: count each worker's live requests instead of rescanning them The least-loaded DP replica pick rebuilt a load map over every live request on each admission. The conductor now keeps per-worker counters, updated where a request enters self.requests (_register_request) and on the one way out of it (_drop_request), which both teardown paths take. Co-Authored-By: Claude Opus 5.5 --- mstar/conductor/conductor.py | 32 +++++++++++-------- .../test_conductor_replica_assignment.py | 17 +++++----- test/modular/test_conductor_tick_exception.py | 3 +- test/modular/test_process_liveness.py | 2 ++ test/modular/test_teardown_barrier.py | 14 +++++--- 5 files changed, 41 insertions(+), 27 deletions(-) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 410cf2b34e..a69f88c399 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -7,7 +7,7 @@ import socket import threading import time -from collections import defaultdict, deque +from collections import Counter, defaultdict, deque from dataclasses import dataclass, field import numpy as np @@ -394,6 +394,9 @@ def __init__( status_interval_s: float | None = None, ): self.requests: dict[str, RequestData] = {} + # worker id -> requests in self.requests it serves, for the DP replica + # pick; kept by _register_request / _drop_request + self._worker_load: Counter[str] = Counter() # Requests in teardown: kept in self.requests (so they still count toward # concurrency until their GPU state is freed) with barrier state here. self.draining: dict[str, DrainingRequest] = {} @@ -857,23 +860,14 @@ def _assign_worker_graphs_to_workers(self) -> dict[int, list[str]]: """ # _group_id -> chosen DP-replica index within that group's ranks group_id_to_replica_idx: dict[int, int] = {} - # worker id -> requests it currently serves; built only when a group - # actually has more than one replica to choose from - load: dict[str, int] | None = None def pick(replicas: list[list[int]]) -> int: """The least-loaded replica (by live requests on its workers), ties at random. A random pick alone puts two of two concurrent requests on one replica half the time, leaving the other idle.""" - nonlocal load if len(replicas) == 1: return 0 - if load is None: - load = defaultdict(int) - for request_data in self.requests.values(): - for worker_id in self._request_workers(request_data): - load[worker_id] += 1 - loads = [sum(load[f"worker_{r}"] for r in ranks) for ranks in replicas] + loads = [sum(self._worker_load[f"worker_{r}"] for r in ranks) for ranks in replicas] least = min(loads) return int(np.random.choice([i for i, n in enumerate(loads) if n == least])) @@ -1118,7 +1112,7 @@ def _admit_request(self, body: NewRequestConductor): except Exception as exc: # noqa: BLE001 request_data = self.requests.get(body.request_id) if request_data is None or not request_data.dispatched: - self.requests.pop(body.request_id, None) + self._drop_request(body.request_id) if isinstance(exc, ValueError): self._reject_request(body.request_id, str(exc), status=400) return @@ -1223,7 +1217,7 @@ def _do_ingest_request( conductor_ingest_time=ingest_time, conductor_arrival_time=arrival_time, ) - self.requests[body.request_id] = request_data + self._register_request(body.request_id, request_data) # Kick off all partitions by calling get_initial_forward_pass_args per partition partition_fwd_args: dict[str, ForwardPassArgs] = {} @@ -1366,6 +1360,16 @@ def _request_workers(self, request_data: RequestData) -> set[str]: for worker_id in worker_ids } + def _register_request(self, request_id: str, request_data: RequestData) -> None: + self.requests[request_id] = request_data + self._worker_load.update(self._request_workers(request_data)) + + def _drop_request(self, request_id: str) -> None: + """The one way a request leaves ``self.requests``, so ``_worker_load`` cannot leak.""" + request_data = self.requests.pop(request_id, None) + if request_data is not None: + self._worker_load.subtract(self._request_workers(request_data)) + def _register_draining( self, request_id: str, expected_acks: set[str], participants: set[str], failure_error: str | None = None, failure_status: int = 500, @@ -1436,7 +1440,7 @@ def _finalize_draining(self, request_id: str): ), ), ) - self.requests.pop(request_id, None) + self._drop_request(request_id) logger.info("Tore down request %s; freed worker resources", request_id) self._try_admit_waiting() diff --git a/test/modular/test_conductor_replica_assignment.py b/test/modular/test_conductor_replica_assignment.py index e2effee144..83481ec856 100644 --- a/test/modular/test_conductor_replica_assignment.py +++ b/test/modular/test_conductor_replica_assignment.py @@ -1,6 +1,8 @@ """Data-parallel replica choice: a new request goes to the replica serving the fewest live requests, so concurrent requests spread over the replicas instead of -landing on one at random.""" +landing on one at random. The load is per-worker counters the conductor keeps as +requests register and drop.""" +from collections import Counter from types import SimpleNamespace import pytest @@ -18,12 +20,13 @@ def _conductor(worker_graphs): conductor = object.__new__(Conductor) conductor.worker_graphs = worker_graphs conductor.requests = {} + conductor._worker_load = Counter() return conductor def _admit(conductor, rid): assignment = conductor._assign_worker_graphs_to_workers() - conductor.requests[rid] = SimpleNamespace(worker_graph_to_workers=assignment) + conductor._register_request(rid, SimpleNamespace(worker_graph_to_workers=assignment)) return assignment @@ -35,7 +38,10 @@ def test_requests_alternate_over_dp_replicas(): # a group's wgs stay on one replica, and the second request takes the idle one assert first[0] == first[1] and second[0] == second[1] assert first[0] != second[0] - del conductor.requests["a"] + # a request leaving frees its replica's count, once + conductor._drop_request("a") + conductor._drop_request("a") + assert conductor._worker_load == Counter({first[0][0]: 0, second[0][0]: 1}) assert _admit(conductor, "c")[0] == first[0] @@ -45,8 +51,3 @@ def test_lockstep_instances_are_balanced_whole(): picks = {tuple(_admit(conductor, rid)[0]) for rid in ("a", "b")} assert picks == {("worker_0", "worker_1"), ("worker_2", "worker_3")} - -def test_single_replica_needs_no_load_scan(): - conductor = _conductor({0: _wg(0, [3])}) - conductor.requests = None # never read when there is nothing to choose - assert conductor._assign_worker_graphs_to_workers() == {0: ["worker_3"]} diff --git a/test/modular/test_conductor_tick_exception.py b/test/modular/test_conductor_tick_exception.py index a9b7608d01..709b727b0a 100644 --- a/test/modular/test_conductor_tick_exception.py +++ b/test/modular/test_conductor_tick_exception.py @@ -6,7 +6,7 @@ error, its concurrency slot is freed, and every other request in the same tick carries on. """ -from collections import deque +from collections import Counter, deque from types import SimpleNamespace import pytest @@ -105,6 +105,7 @@ def _conductor(*ticks, requests): c._startup_message_backlog = [] c.requests = dict(requests) c.draining = {} + c._worker_load = Counter() c.waiting_queue = [] c.max_concurrent_requests = None c._early_reads_done = {} diff --git a/test/modular/test_process_liveness.py b/test/modular/test_process_liveness.py index a7ab5df44e..976c142e10 100644 --- a/test/modular/test_process_liveness.py +++ b/test/modular/test_process_liveness.py @@ -12,6 +12,7 @@ import signal import threading import time +from collections import Counter from types import SimpleNamespace from unittest import mock @@ -104,6 +105,7 @@ def _conductor(procs, inbox=()): c._next_liveness_check = 0.0 c.requests = {} c.draining = {} + c._worker_load = Counter() c.waiting_queue = [] return c diff --git a/test/modular/test_teardown_barrier.py b/test/modular/test_teardown_barrier.py index b057bb6412..e91f846d26 100644 --- a/test/modular/test_teardown_barrier.py +++ b/test/modular/test_teardown_barrier.py @@ -7,7 +7,7 @@ """ import types -from collections import deque +from collections import Counter, deque import pytest import torch @@ -57,7 +57,10 @@ def _conductor(requests, drain_ttl_s=120.0): c.communicator = types.SimpleNamespace( send=lambda entity_id, msg: c.sent.append((entity_id, msg)) ) - c.requests = dict(requests) + c.requests = {} + c._worker_load = Counter() + for rid, request_data in requests.items(): + c._register_request(rid, request_data) c.draining = {} c._drain_ttl_s = drain_ttl_s c._early_reads_done = {} @@ -116,6 +119,8 @@ def test_abort_drains_all_participants_then_removes(): assert _remove_targets(c, "r1") == {"w0", "w1", PREPROCESS} assert "r1" not in c.requests and "r1" not in c.draining assert c.admits == 1 + # the DP replica pick no longer counts it + assert not any(c._worker_load.values()) def test_abort_of_unknown_request_sends_nothing(): @@ -292,11 +297,12 @@ def test_bad_value_registered_but_not_dispatched_is_a_400(): c = _conductor({}) def bad_config(body): - c.requests[body.request_id] = _request_data(workers=("w0",)) + c._register_request(body.request_id, _request_data(workers=("w0",))) raise ValueError("top_p must be a finite number in (0, 1]; got -1") c._do_ingest_request = bad_config c._ingest_request(_new_request("r1")) assert "r1" not in c.requests and c.draining == {} + assert not any(c._worker_load.values()) assert _remove_targets(c, "r1") == {PREPROCESS} [failure] = _failures(c, "r1") assert failure.status == 400 and "top_p" in failure.error_message @@ -306,7 +312,7 @@ def test_ingest_error_after_dispatch_drains_like_a_worker_failure(): c = _conductor({}) def boom(body): - c.requests[body.request_id] = _request_data(workers=("w0",)) + c._register_request(body.request_id, _request_data(workers=("w0",))) c.requests[body.request_id].dispatched = True raise ValueError("half dispatched") c._do_ingest_request = boom From e645af5b92dafa2fa1fe350cf5e5f99a9f48f8c3 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:41:03 -0700 Subject: [PATCH 16/59] Qwen3.5 vision: require the pixel normalizer and byte patches _normalize passed non-uint8 patches through unnormalized, and the None default normalizer raised a TypeError on bytes. The normalizer is now a required argument and any other dtype is refused with ValueError. Co-Authored-By: Claude Opus 5.5 --- mstar/model/qwen3_5/submodules.py | 5 +++-- test/modular/test_qwen3_5_video.py | 7 ++++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 1029481fa2..b01c06bef7 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -708,7 +708,7 @@ class VisionEncoderSubmodule(NodeSubmodule): def __init__( self, model: Qwen3_5VisionModel, config: Qwen3_5VisionConfig, - pixel_normalizer: tuple[torch.Tensor, torch.Tensor] | None = None, + pixel_normalizer: tuple[torch.Tensor, torch.Tensor], ): super().__init__() self.model = model @@ -905,7 +905,8 @@ def _normalize(self, pixel_values: torch.Tensor) -> torch.Tensor: compiled embed, so the arithmetic (a subtract, a true divide) is the processors' exactly.""" if pixel_values.dtype != torch.uint8: - return pixel_values + # `process_prompt` ships bytes; anything else would skip normalization + raise ValueError(f"vision patches must be uint8, got {pixel_values.dtype}") mean, std = self._pixel_normalizer if mean.device != pixel_values.device: mean, std = mean.to(pixel_values.device), std.to(pixel_values.device) diff --git a/test/modular/test_qwen3_5_video.py b/test/modular/test_qwen3_5_video.py index c1da5d74bd..76c6047877 100644 --- a/test/modular/test_qwen3_5_video.py +++ b/test/modular/test_qwen3_5_video.py @@ -514,7 +514,7 @@ def _encoder(max_segments: int): patch_size=16, temporal_patch_size=2, spatial_merge_size=2, num_position_embeddings=16, out_hidden_size=64, ) - encoder = VisionEncoderSubmodule(nn.Linear(1, 1), config) + encoder = VisionEncoderSubmodule(nn.Linear(1, 1), config, (torch.zeros(1), torch.ones(1))) encoder.node_resources = {VISION_ATTN: SimpleNamespace(max_segments_per_request=max_segments)} return encoder @@ -542,6 +542,11 @@ def request(segments: int) -> NodeInputs: assert not config.can_lease([request(2), request(3)]) +def test_vision_encoder_refuses_patches_it_would_not_normalize(): + with pytest.raises(ValueError, match="uint8"): + _encoder(max_segments=4)._normalize(torch.zeros(2, 3)) + + def test_vision_block_loop_replays_exactly_when_leased(): """A leased step declares nothing of its own, so the forward replays then and runs eagerly otherwise, whatever the runner's buckets allow.""" From fc883cdf5d30c4c8a9962dd23b85aeb5b80658a1 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:41:21 -0700 Subject: [PATCH 17/59] Fix stale Qwen3.5 vision docstrings and the flash-attn-4 install paragraph The block loop replays a whole step as one row, not one prompt per replay. In installation.rst, "That backend needs flash-attn-4" followed the LingBot sentence; the paragraph now reads in order. Co-Authored-By: Claude Opus 5.5 --- docs/installation.rst | 9 ++++++--- mstar/model/qwen3_5/submodules.py | 7 ++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 4f927bc53b..bfba01126e 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -362,12 +362,15 @@ flash-attn-4 (Waypoint, FA3 sm90 kernel) ---------------------------------------- Waypoint's DiT attention runs torch flex-attention with the ``FLASH`` backend by -default. LingBot-World's ``configs/lingbot_world_1_3b_fast.yaml`` (``reference_numerics: -false``) also calls it directly. Install it after ``flash-attn`` (FA2): the FA2 package -ships an older ``flash_attn/cute`` that overwrites flash-attn-4's. That backend needs the **flash-attn-4** package, which provides +default. That backend needs the **flash-attn-4** package, which provides ``flash_attn.cute`` — the CuTe DSL rewrite of flash-attn; on H100 it runs the FA3-style sm90 kernel with TMA and warpgroup specialisation. It is **not on PyPI** as of 2026-09-17, and it is **not** pulled in by any extra. +LingBot-World's ``configs/lingbot_world_1_3b_fast.yaml`` (``reference_numerics: +false``) also calls flash-attn-4 directly. + +Install it after ``flash-attn`` (FA2): the FA2 package ships an older +``flash_attn/cute`` that overwrites flash-attn-4's. FLASH needs **torch 2.11 or newer**: older torch does not pass flash-attn-4 the block-sparse block size, and capture fails with ``Block sparse tensors ... diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index b01c06bef7..fc02311cd1 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -785,11 +785,12 @@ def get_piecewise_accelerator_graph_configs( self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, **kwargs, ) -> dict[str, PiecewiseAcceleratorGraphConfig]: - """The block loop, one prompt per replay, by patch-count bucket. + """The block loop, one replay per step, by patch-count bucket. Patch embed, position resample and rope run eagerly before it, the merger after (it quarters the row count, which a region's output view - cannot express). Batch size 1: one replay is one prompt's segments. + cannot express). Batch size 1: a replay is the whole step as one row, + every request's segments in it. """ hidden, head_dim = self.config.hidden_size, self.config.head_dim @@ -850,7 +851,7 @@ def _replay_block_loop( sin: torch.Tensor, seq_lengths: tuple[int, ...], ) -> torch.Tensor | None: - """The block loop replayed once per prompt, or None if not leased. + """The block loop replayed once for the whole step, or None if not leased. A leased step declares nothing of its own (`declare_step`), so ``engine_inputs.step`` is None; a bucket that fits is not enough, as From 805bd71db088dbe755c59c0b8d0194cb07ef1310 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:41:37 -0700 Subject: [PATCH 18/59] wan22: keep the unfused FFN GELU when torch lacks _addmm_activation Co-Authored-By: Claude Opus 5.5 --- mstar/model/wan22/components/dit.py | 9 ++++++++- test/modular/test_wan22_model.py | 13 +++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/mstar/model/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index 61b317a4a6..d79d09dd7f 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -29,6 +29,7 @@ two-linear MLPs ARE the shared ``components.mlp.MLP``. """ +import logging import math from typing import NamedTuple @@ -40,6 +41,8 @@ from mstar.model.components.mlp import MLP from mstar.model.wan22.config import Wan22Config +logger = logging.getLogger(__name__) + # 3D RoPE base frequency (reference WanRotaryPosEmbed default `theta`). ROPE_THETA = 10000.0 @@ -352,7 +355,11 @@ def __init__(self, config: Wan22Config): def set_fuse_ffn_gelu(self, enabled: bool) -> None: """Run every block's FFN GELU in the GEMM epilogue (see ``_ffn_gelu_epilogue``); off by default, so the module alone stays the - reference port.""" + reference port. ``torch._addmm_activation`` is private: without it the + blocks keep the unfused GELU.""" + if enabled and not hasattr(torch, "_addmm_activation"): + logger.warning("Wan2.2 DiT: torch._addmm_activation is missing; the FFN GELU stays unfused") + enabled = False for block in self.blocks: block.fuse_ffn_gelu = enabled diff --git a/test/modular/test_wan22_model.py b/test/modular/test_wan22_model.py index 11d4b59ed9..8ea75170a3 100644 --- a/test/modular/test_wan22_model.py +++ b/test/modular/test_wan22_model.py @@ -177,6 +177,19 @@ def test_wan22_forwards_run_without_autograd(): assert not torch.is_autocast_enabled() + +def test_wan22_dit_keeps_the_unfused_gelu_without_the_private_op(monkeypatch): + """The fused FFN GELU calls ``torch._addmm_activation``, which torch does not promise.""" + from types import SimpleNamespace + + from mstar.model.wan22.components.dit import Wan22DiT + + monkeypatch.delattr(torch, "_addmm_activation", raising=False) + dit = SimpleNamespace(blocks=[SimpleNamespace(fuse_ffn_gelu=False)]) + Wan22DiT.set_fuse_ffn_gelu(dit, True) + assert dit.blocks[0].fuse_ffn_gelu is False + + def test_wan22_encoder_walks_persist_embeddings(): model = _make_model() walks = model.get_graph_walk_graphs() From 44bcfb8aec806962d02451eb63f118409163fd3c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:44:17 -0700 Subject: [PATCH 19/59] omnivoice: carry target_len as step metadata, not an all-MASK canvas edge process_prompt returns the target length as ProcessPromptOutput metadata (an internal kwarg), the conductor puts it in step metadata, and the backbone builds the iteration-0 canvas on the device. The [C, target_len] initial_canvas edge is gone. Co-Authored-By: Claude Opus 5.5 --- mstar/model/omnivoice/omnivoice_model.py | 29 ++++++++++------------- mstar/model/omnivoice/submodules.py | 18 +++++++++----- test/modular/test_step_prep_host_syncs.py | 20 +++++++++++----- 3 files changed, 39 insertions(+), 28 deletions(-) diff --git a/mstar/model/omnivoice/omnivoice_model.py b/mstar/model/omnivoice/omnivoice_model.py index c6d990c508..00a3c4ff40 100644 --- a/mstar/model/omnivoice/omnivoice_model.py +++ b/mstar/model/omnivoice/omnivoice_model.py @@ -40,7 +40,7 @@ TensorPointerInfo, ) from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, Model, ProcessPromptOutput from mstar.model.omnivoice.components.backbone import ( OmniVoiceBackbone, assert_flashinfer_api, @@ -55,6 +55,7 @@ from mstar.model.omnivoice.config import OmniVoiceConfig from mstar.model.omnivoice.submodules import ( BACKBONE_ATTN, + TARGET_LEN, UNMASK_LOOP_NAME, OmniVoiceBackboneSubmodule, OmniVoiceCode2WavSubmodule, @@ -73,6 +74,9 @@ class OmniVoiceModel(Model): SPEECH_GEN_CLONE_WALK = "speech_gen_clone" UNMASK_LOOP_NAME = UNMASK_LOOP_NAME + # The canvas length process_prompt settles (it needs the tokenizer and the + # duration estimator); it reaches the backbone as step metadata. + INTERNAL_KWARGS = (TARGET_LEN,) def __init__( self, @@ -166,7 +170,6 @@ def _build_speech_walk(self, clone: bool) -> GraphSection: backbone_inputs = [ "prefix_ids", "prefix_audio_mask", - "initial_canvas", "audio_tokens", "step_index", ] @@ -317,7 +320,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Validate, then build the canvas prefix on the data worker. Every malformed-request check lives here rather than in @@ -407,12 +410,6 @@ def process_prompt( out: NameToTensorList = { "prefix_ids": [prefix_ids], "prefix_audio_mask": [prefix_audio_mask], - # the all-MASK canvas the unmask loop starts from: its shape carries - # the target length, so the backbone sizes the request without a - # host read of a device value - "initial_canvas": [ - torch.full((self.config.num_audio_codebook, target_len), self.config.audio_mask_id, dtype=torch.long) - ], } if ref_audio: out["ref_audio_inputs"] = [ref_waveform] @@ -420,7 +417,7 @@ def process_prompt( # _prepare_reference: post_process uses it to put the clone back at # the source's loudness, so a quiet reference yields quiet speech. out["ref_rms"] = [torch.tensor([ref_rms], dtype=torch.float32)] - return out + return ProcessPromptOutput(out, {TARGET_LEN: target_len}) def _prepare_reference( self, waveform: torch.Tensor, ref_text: str | None @@ -494,8 +491,8 @@ def get_initial_forward_pass_args( # Backstops. process_prompt already rejected each of these on the data # worker where a ValueError becomes a 400; a raise here runs at the # conductor, whose loop swallows it and leaves the client hanging. - if not input_signals.get("prefix_ids"): - raise ValueError("OmniVoice requires a processed prompt (prefix_ids)") + if not input_signals.get("prefix_ids") or TARGET_LEN not in model_kwargs: + raise ValueError("OmniVoice requires a processed prompt (prefix_ids, target_len)") target_output = output_modalities[0] if output_modalities else "audio" if target_output != "audio": raise ValueError( @@ -522,9 +519,8 @@ def get_initial_forward_pass_args( kwargs = { "walk_schedule": schedule, "walk_step": 0, - # target_len is NOT here: it is computed on the data worker (it - # needs the tokenizer and the duration estimator) and rides to the - # backbone as the initial canvas's shape. + # from process_prompt, which settles it on the data worker + TARGET_LEN: int(model_kwargs[TARGET_LEN]), "num_step": num_step, "guidance_scale": float( model_kwargs.get("guidance_scale", defaults.guidance_scale) @@ -576,6 +572,7 @@ def _get_step_metadata(self, metadata: CurrentForwardConductorMetadata) -> dict: kw = metadata.kwargs return { "is_prefill": metadata.is_prefill, + TARGET_LEN: kw[TARGET_LEN], "num_step": kw["num_step"], "guidance_scale": kw["guidance_scale"], "t_shift": kw["t_shift"], @@ -596,7 +593,7 @@ def _speech_gen_inputs( step counter at iteration 0. """ inputs: list[GraphEdge] = [] - for name in ("prefix_ids", "prefix_audio_mask", "initial_canvas"): + for name in ("prefix_ids", "prefix_audio_mask"): edge = GraphEdge(next_node="backbone", name=name) edge.tensor_info = persist_signals.get(name, []) inputs.append(edge) diff --git a/mstar/model/omnivoice/submodules.py b/mstar/model/omnivoice/submodules.py index a8955ac7f4..2b6c4d8329 100644 --- a/mstar/model/omnivoice/submodules.py +++ b/mstar/model/omnivoice/submodules.py @@ -51,6 +51,8 @@ logger = logging.getLogger(__name__) UNMASK_LOOP_NAME = "unmask_loop" +# Step metadata: the canvas length, in frames +TARGET_LEN = "target_len" # The backbone's ragged attention resource, and its captured layer stack. BACKBONE_ATTN = "backbone_attn" BACKBONE_LAYERS = "backbone_layers" @@ -227,12 +229,16 @@ def prepare_inputs( if "audio_tokens" not in inputs or len(inputs["audio_tokens"]) == 0: # Iteration 0: the loop-back edges arrive empty, and the loop starts - # from the all-MASK canvas the data worker built (it knows the - # target length from the tokenizer and the duration estimator). - # There is no noise and no seed to honour: the stochastic part of - # this model is the reveal order, not the starting state. - tokens = inputs["initial_canvas"][0].to(device) - target_len = tokens.shape[-1] + # from an all-MASK canvas. There is no noise and no seed to honour: + # the stochastic part of this model is the reveal order, not the + # starting state. + target_len = int(meta[TARGET_LEN]) + tokens = torch.full( + (self.config.num_audio_codebook, target_len), + self.config.audio_mask_id, + dtype=torch.long, + device=device, + ) # 1-D, not 0-dim: the worker's output fanout reads dims[0] of # every routed edge. step_index = torch.zeros(1, dtype=torch.int64, device=device) diff --git a/test/modular/test_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index 53de2d96cd..221b675c50 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -134,9 +134,9 @@ def _omnivoice(): return sub, cfg -def _omnivoice_meta(steps): +def _omnivoice_meta(steps, target_len=6): return { - "num_step": steps, "t_shift": 0.1, "guidance_scale": 2.0, "class_temperature": 0.0, + "target_len": target_len, "num_step": steps, "t_shift": 0.1, "guidance_scale": 2.0, "class_temperature": 0.0, "layer_penalty_factor": 5.0, "position_temperature": 5.0, } @@ -148,7 +148,6 @@ def _omnivoice_inputs(cfg, k, target_len=6, clone=False): inputs = { "prefix_ids": [torch.arange(c * 5).view(c, 5)], "prefix_audio_mask": [torch.zeros(5, dtype=torch.bool)], - "initial_canvas": [torch.full((c, target_len), cfg.audio_mask_id)], "audio_tokens": [tokens], "step_index": [torch.tensor([k], dtype=torch.int64)], } @@ -173,13 +172,22 @@ def prepare(fwd_info, inp): def test_omnivoice_first_step_prep_reads_no_tensor(): - # iteration 0 starts from the data worker's all-MASK canvas: its shape is - # the target length, so nothing is read on the host + # iteration 0 builds the all-MASK canvas on the device from the target + # length process_prompt settled, carried as step metadata, so nothing is + # read on the host + from mstar.model.omnivoice.omnivoice_model import OmniVoiceModel + sub, cfg = _omnivoice() + model = OmniVoiceModel.__new__(OmniVoiceModel) + model.config = cfg + first = model.get_initial_forward_pass_args( + "default", ["text"], ["audio"], {"prefix_ids": [None]}, model_kwargs={"target_len": 6}, + ) + assert {edge.name for edge in first.inputs} == {"prefix_ids", "prefix_audio_mask", "audio_tokens", "step_index"} inputs = _omnivoice_inputs(cfg, 0) del inputs["audio_tokens"], inputs["step_index"] with _HostReads() as reads: - ni = sub.prepare_inputs("speech_gen", _fwd("r", "speech_gen", _omnivoice_meta(4)), inputs) + ni = sub.prepare_inputs("speech_gen", _fwd("r", "speech_gen", first.step_metadata), inputs) assert reads.reads == [], reads.reads assert ni.tensor_inputs["audio_tokens"].shape == (cfg.num_audio_codebook, 6) assert (ni.tensor_inputs["audio_tokens"] == cfg.audio_mask_id).all() From eb70ae905a77c9d08a875df5bfaaf312a9b733d2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:45:58 -0700 Subject: [PATCH 20/59] warmup: let a spec's concurrency be a list of batch sizes Each value runs as its own warmup, in order, so the Cosmos3 DROID configs warm batch sizes 1-8 with one entry instead of eight copies of it. Co-Authored-By: Claude Opus 5.5 --- configs/cosmos3_droid.yaml | 36 +----------------- configs/cosmos3_edge_droid.yaml | 32 +--------------- docs/serving.rst | 3 +- mstar/api_server/entrypoint.py | 57 +++++++++++++++-------------- test/modular/test_startup_warmup.py | 9 ++++- 5 files changed, 42 insertions(+), 95 deletions(-) diff --git a/configs/cosmos3_droid.yaml b/configs/cosmos3_droid.yaml index 35a5e9b5a6..0df754f65c 100644 --- a/configs/cosmos3_droid.yaml +++ b/configs/cosmos3_droid.yaml @@ -38,38 +38,4 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 2 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 3 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 4 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 5 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 6 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 7 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 8 + concurrency: [1, 2, 3, 4, 5, 6, 7, 8] diff --git a/configs/cosmos3_edge_droid.yaml b/configs/cosmos3_edge_droid.yaml index 2632aa69e2..16fcf44df2 100644 --- a/configs/cosmos3_edge_droid.yaml +++ b/configs/cosmos3_edge_droid.yaml @@ -45,34 +45,4 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 2 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 3 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 4 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 5 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 6 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 7 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 8 + concurrency: [2, 3, 4, 5, 6, 7, 8] diff --git a/docs/serving.rst b/docs/serving.rst index 38309d2cb7..1cc84e5c5b 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -262,7 +262,8 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key comes from the extension, or give ``{path: ..., modality: ...}``. ``repeat: N`` sends ``text`` N times over, a long prompt for a path only long prompts take (a long-context prefill). ``concurrency: N`` submits N - copies at once, so they batch, for a step that compiles per batch size. A + copies at once, so they batch, for a step that compiles per batch size; a + list (``concurrency: [1, 2, 4]``) runs the entry once per size, in order. A malformed or failing warmup is logged and skipped. * - ``dist_timeout_s`` - *(optional)* Timeout in seconds for the NCCL world group and its parallel diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 8ed9cfda0b..cd6dc9063d 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -1348,13 +1348,14 @@ class WarmupRequest: """One entry of the deployment's ``warmup_requests``: a ``/generate`` request in yaml form. ``files`` are the media inputs in prompt order (any modality, any number), followed by ``text`` as on ``/generate``, - ``repeat`` times over (a long prompt for a long-context path).""" + ``repeat`` times over (a long prompt for a long-context path). + ``concurrency`` holds the batch sizes it runs at, one warmup each, in order.""" text: str | None = None files: list[WarmupMedia] = field(default_factory=list) output_modalities: list[str] = field(default_factory=lambda: ["text"]) model_kwargs: dict | None = None - concurrency: int = 1 + concurrency: tuple[int, ...] = (1,) @classmethod def from_config(cls, spec: Any) -> "WarmupRequest": @@ -1372,8 +1373,9 @@ def from_config(cls, spec: Any) -> "WarmupRequest": if repeat > 1 and not text: raise ValueError("repeat needs text") concurrency = spec.get("concurrency", 1) - if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: - raise ValueError(f"concurrency must be a positive integer, got {concurrency!r}") + levels = concurrency if isinstance(concurrency, list) else [concurrency] + if not levels or any(isinstance(n, bool) or not isinstance(n, int) or n < 1 for n in levels): + raise ValueError(f"concurrency must be a positive integer or a list of them, got {concurrency!r}") out_mods = spec.get("output_modalities") or ["text"] if isinstance(out_mods, str): out_mods = [m.strip() for m in out_mods.split(",") if m.strip()] @@ -1388,7 +1390,7 @@ def from_config(cls, spec: Any) -> "WarmupRequest": files=[WarmupMedia.from_config(f) for f in files], output_modalities=list(out_mods), model_kwargs=model_kwargs, - concurrency=concurrency, + concurrency=tuple(levels), ) def submit(self, server: APIServer, request_id: str) -> str: @@ -1432,11 +1434,11 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non warmup that is malformed or fails is logged and skipped; it must not keep the server from coming up. - Each spec runs as ``copies`` x its ``concurrency`` concurrent requests: - ``copies`` warms every replica of a data-parallel deployment, since the - conductor sends concurrent requests to different replicas (the - least-loaded one each), and ``concurrency`` lets them batch, for a step - that compiles per batch size. + Each spec runs as ``copies`` x its ``concurrency`` concurrent requests, + once per value when ``concurrency`` is a list: ``copies`` warms every + replica of a data-parallel deployment, since the conductor sends + concurrent requests to different replicas (the least-loaded one each), and + ``concurrency`` lets them batch, for a step that compiles per batch size. """ if not isinstance(specs, list): logger.warning("warmup_requests must be a list; ignoring %r", type(specs).__name__) @@ -1445,9 +1447,8 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non async def _one(request: WarmupRequest, request_id: str) -> None: await server.collect_results(request.submit(server, request_id)) - async def _all(request: WarmupRequest, index: int) -> None: - n = copies * request.concurrency - ids = [f"warmup-{index}"] + [f"warmup-{index}-{c}" for c in range(1, n)] + async def _all(request: WarmupRequest, base_id: str, n: int) -> None: + ids = [base_id] + [f"{base_id}-{c}" for c in range(1, n)] # every copy runs to its end before a failure is reported, so none is # cancelled mid-request and still in the server when the next spec starts results = await asyncio.gather(*(_one(request, rid) for rid in ids), return_exceptions=True) @@ -1461,21 +1462,23 @@ async def _all(request: WarmupRequest, index: int) -> None: except Exception as exc: # noqa: BLE001 — a repeat too large to build, say logger.warning("warmup_requests[%d] skipped: %s", index, exc) continue - t0 = time.perf_counter() - try: - asyncio.run(_all(request, index)) - except Exception as exc: # noqa: BLE001 — a cold server is better than none - logger.warning( - "warmup request %d/%d failed after %.1f s: %s", - index + 1, len(specs), time.perf_counter() - t0, - getattr(exc, "detail", exc), + for level, concurrency in enumerate(request.concurrency): + t0 = time.perf_counter() + try: + base_id = f"warmup-{index}" + (f".{level}" if level else "") + asyncio.run(_all(request, base_id, copies * concurrency)) + except Exception as exc: # noqa: BLE001 — a cold server is better than none + logger.warning( + "warmup request %d/%d failed after %.1f s at concurrency %d: %s", + index + 1, len(specs), time.perf_counter() - t0, concurrency, + getattr(exc, "detail", exc), + ) + continue + logger.info( + "warmup request %d/%d (%s) done in %.1f s at concurrency %d", + index + 1, len(specs), ",".join(request.output_modalities), + time.perf_counter() - t0, concurrency, ) - continue - logger.info( - "warmup request %d/%d (%s) done in %.1f s", - index + 1, len(specs), ",".join(request.output_modalities), - time.perf_counter() - t0, - ) # ------------------------------------------------------------------ diff --git a/test/modular/test_startup_warmup.py b/test/modular/test_startup_warmup.py index b948b84db0..b2e6988b23 100644 --- a/test/modular/test_startup_warmup.py +++ b/test/modular/test_startup_warmup.py @@ -171,8 +171,15 @@ def test_concurrency_runs_copies_together_for_batched_shapes() -> None: fake = _FakeServer() entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 3}], copies=2) assert [s["request_id"] for s in fake.submitted] == [f"warmup-0{s}" for s in ("", "-1", "-2", "-3", "-4", "-5")] + # a list runs one warmup per batch size, in order, each under its own ids fake = _FakeServer() - entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}]) + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": [1, 2]}]) + assert [s["request_id"] for s in fake.submitted] == ["warmup-0", "warmup-0.1", "warmup-0.1-1"] + fake = _FakeServer() + entrypoint._run_warmup_requests(fake, [ + {"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}, + {"text": "c", "concurrency": []}, {"text": "d", "concurrency": [2, 0]}, + ]) assert fake.submitted == [] From e60401cf647110466738b58da1d17bd7162d39af Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:46:35 -0700 Subject: [PATCH 21/59] lingbot_world: say the reference mode's decode is within ~73 dB, not exact Co-Authored-By: Claude Opus 5.5 --- configs/lingbot_world_1_3b_fast.yaml | 2 +- docs/models.rst | 4 ++++ mstar/model/lingbot_world/config.py | 4 +++- mstar/model/lingbot_world/submodules.py | 4 ++-- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/configs/lingbot_world_1_3b_fast.yaml b/configs/lingbot_world_1_3b_fast.yaml index 2cf4ca4b2e..52e23921fa 100644 --- a/configs/lingbot_world_1_3b_fast.yaml +++ b/configs/lingbot_world_1_3b_fast.yaml @@ -1,6 +1,6 @@ model: "lingbot_world_1_3b" -# Speed over bit-parity with the reference (~1.5x at 480x832 on an H200): +# Speed over closeness to the reference (~1.5x at 480x832 on an H200): # flash-attn-4 attention, compiled block pointwise ops, an autotuned condition # encoder and a bf16 decode. Rollouts drift from the reference's about as far as # under any other attention kernel. Needs flash-attn-4. diff --git a/docs/models.rst b/docs/models.rst index 3f047b0a4e..8b232d1f6a 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -126,6 +126,10 @@ Registry keys live in ``mstar/model/registry.py`` (``MODEL_REGISTRY`` / ``HF_MOD A request carries the whole camera trajectory as JSON in ``model_kwargs``, which the API server caps at 1 MiB: about 3,900 frames of ``[4, 4]`` float32 poses; send ``[F, 3, 4]`` rows for up to the 4,077-frame maximum. Needs ``flash-attn``; one rollout per GPU. + The default keeps the reference's numerics except the VAE decode, which runs + channels-last and compiled: deterministic, and within ~73 dB of the reference's + decode rather than bit-exact. ``configs/lingbot_world_1_3b_fast.yaml`` trades + more of that closeness for ~1.5x the speed. The weights are CC-BY-NC-SA-4.0 (non-commercial, share-alike), not Apache-2.0; check it before deploying. The code is Apache-2.0, adapted from LingBot-World v1 (``Robbyant/lingbot-world``). * - ``flux2_klein`` diff --git a/mstar/model/lingbot_world/config.py b/mstar/model/lingbot_world/config.py index 1c40b1ddd3..949582a582 100644 --- a/mstar/model/lingbot_world/config.py +++ b/mstar/model/lingbot_world/config.py @@ -84,7 +84,9 @@ class LingBotWorldConfig: vae_decode_cudnn_benchmark: bool = False vae_decode_compile: bool = True - # False trades bit-parity with the reference for speed: flash-attn-4 + # True (the default) keeps the reference's numerics but for the decode + # above: deterministic, within ~73 dB of the reference's, not bit-exact. + # False trades closeness to the reference for speed: flash-attn-4 # attention (needs flash-attn-4, see docs/installation.rst), compiled block # pointwise ops, and a cuDNN-autotuned condition encoder and bf16 decode: # 12.6 -> 18.6 fps at 480x832 on an H200. A free-running rollout then drifts diff --git a/mstar/model/lingbot_world/submodules.py b/mstar/model/lingbot_world/submodules.py index f53a646d08..d1670fd8d9 100644 --- a/mstar/model/lingbot_world/submodules.py +++ b/mstar/model/lingbot_world/submodules.py @@ -150,8 +150,8 @@ class LingBotDitSubmodule(_SingleRequestMixin, NodeSubmodule): At the first iteration the request's state is seeded: the K/V cache, the generator (seeded with the request seed) and the full rollout's initial noise drawn from it in one call. The reference draws the initial noise the - same way, which is what makes the rollout reproducible bit for bit; it - requires the rollout length up front. + same way, which is what makes the rollout follow the reference's from the + same seed; it requires the rollout length up front. ``image`` arrives resized and normalized to [-1, 1] by ``LingBotWorldModel.process_prompt``, which runs the reference's CPU bicubic From db26cd3b56b7b7fbbcd5a27dda58e88261d21848 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:47:10 -0700 Subject: [PATCH 22/59] ragged attention: the resource owns its compile op; no per-call-site opt-in The ragged attention resource now owns the mstar::ragged_attention custom op, and its run() goes through it whenever it is traced by torch.compile. Every compiled caller keeps the attention in its graph without opting in, and every eager caller keeps the plain kernel call. Sketch: - Goal: a compiled region keeps ragged attention in its graph, so Inductor fuses the pointwise work around it, without the model choosing compile mechanics. Non-goal: changing eager or capture-only callers. - Foundation: a resource-owned custom op, as with mstar::rope_apply_qk_inplace, mstar::rms_norm and mstar::ring_scatter, routed inside the manager's run() by torch.compiler.is_compiling(). Rejected: the opaque=True call-site flag (the model picks compile mechanics, against invariant 1); routing in RaggedAttentionCallable alone (a layer calling resource.run directly would still break, and the callable is a convenience, not the owner); opaque-object op arguments (newer than the torch floor in pyproject). - New concepts: none. The op moves here from convenience.py. Its index-keyed list becomes _MANAGERS, a dict keyed by resource key (unique within the worker's engine), filled when a manager is built. A rebuild replaces its entry, so the table cannot grow and the key baked into a compiled graph does not change. - Op contract, stated with the op: output shaped like q in the manager's dtype (the fake reads the dtype from the manager), freshly allocated, and no argument mutated. - Per-model cost: none. Wan2.2 drops opaque=True. - Deleted: RaggedAttentionCallable's opaque parameter and _run, _RAGGED_CALLABLES, _register_ragged, the op in convenience.py, and the flashinfer.py docstring saying ragged attention deliberately is not behind a custom op. Co-Authored-By: Claude Opus 5.5 --- .claude/skills/diffusion-scaffold/SKILL.md | 2 +- docs/adding_models.rst | 4 +- mstar/engine/resources/attn/ragged/base.py | 4 +- .../resources/attn/ragged/flashinfer.py | 43 ++++++++++++++++--- mstar/engine/resources/convenience.py | 41 +----------------- mstar/model/wan22/submodules.py | 2 +- test/modular/test_ragged_attention.py | 31 ++++++++----- 7 files changed, 68 insertions(+), 59 deletions(-) diff --git a/.claude/skills/diffusion-scaffold/SKILL.md b/.claude/skills/diffusion-scaffold/SKILL.md index cb042bee33..b921d45181 100644 --- a/.claude/skills/diffusion-scaffold/SKILL.md +++ b/.claude/skills/diffusion-scaffold/SKILL.md @@ -50,7 +50,7 @@ Every attention in the DiT should go through a ragged resource, so the model get - The engine's blanket compile is off for the node (`disable_torch_compile = True`); the model compiles its transformer region with `compile_utils.compile_transformer_forward`, which keeps the exact-op exclusions that preserve eager rounding. Compiling the DiT mostly buys the memory-bound elementwise work (modulation, fp32 norms, RoPE, gating), which on LTX-2.5 was ~45% of an eager step. - `capture_buckets` lists `(graph_walk, bucket_key)` pairs to capture; `capture_request_inputs` supplies one placeholder row per bucket. A bucket is captured per batch size in `capture_batch_sizes`, and each one costs startup, so capture the shapes you serve, not every shape you accept. `replay_walks` lets another walk reuse a captured bucket. -- The captured graph removes the gaps the per-layer attention graph breaks leave in a compiled-but-eager step; measure both before deciding a model doesn't need capture. +- The captured graph removes the launch gaps of a compiled-but-eager step. Ragged attention does not break that compiled graph (the resource's `run` traces as the `mstar::ragged_attention` custom op), so measure both before deciding a model needs capture. ## The overshoot veto diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 4ac6a7b8af..3264c7e9cb 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -742,7 +742,9 @@ encoder node only, and names no cache: A layer binds the key as it binds any other resource, then attends through the resource's ``run(q, k, v, label=None)``. There is no KV write to pair it with, so no -``AttentionCallable`` is involved. +``AttentionCallable`` is involved. Inside a ``torch.compile`` region, ``run`` is the +``mstar::ragged_attention`` custom op, so the region keeps one graph across the attention +instead of breaking per layer; an eager call goes straight to the kernel. The step is an ordinary ``AttentionStep``, and its ``segments`` are the entire layout. There is no cache to read the layout off, so a label that carries no segment in the diff --git a/mstar/engine/resources/attn/ragged/base.py b/mstar/engine/resources/attn/ragged/base.py index 4b4a44004c..f38b0acd21 100644 --- a/mstar/engine/resources/attn/ragged/base.py +++ b/mstar/engine/resources/attn/ragged/base.py @@ -37,7 +37,9 @@ def _build_ragged(spec, info: EngineResourceInfo, manager_cls, **kwargs): "RaggedAttentionConfig.dtype on a node without a KV-backed attention " "(the engine has no KV dtype to fall back on)." ) - return manager_cls(device=info.device, dtype=dtype, config=config, **kwargs) + return manager_cls( + device=info.device, dtype=dtype, config=config, key=spec.resource_key, **kwargs, + ) class RaggedAttnManager(AttentionResource): diff --git a/mstar/engine/resources/attn/ragged/flashinfer.py b/mstar/engine/resources/attn/ragged/flashinfer.py index 4d5c630f70..0d04277195 100644 --- a/mstar/engine/resources/attn/ragged/flashinfer.py +++ b/mstar/engine/resources/attn/ragged/flashinfer.py @@ -25,6 +25,31 @@ from mstar.engine.resources.base import CGSlotKey from mstar.engine.resources.step import Segment, SlotLease, StepContext +# Managers by resource key, which is how the custom op below names one: an op +# takes tensors and scalars, not objects. The key is unique within the worker's +# engine, so a rebuilt resource replaces its entry, and a compiled graph, which +# bakes the key in as a constant, reaches the new manager without recompiling. +_MANAGERS: dict[str, "_FlashInferRaggedBase"] = {} + + +@torch.library.custom_op("mstar::ragged_attention", mutates_args=()) +def ragged_attention( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, key: str, label: str | None, +) -> torch.Tensor: + """``_MANAGERS[key].run(q, k, v, label)`` as one node of a compiled graph. + + Contract: returns a fresh tensor shaped like ``q`` in the manager's dtype, + and mutates no argument. The wrapper's output is a new allocation, or a + clone of the graph-mode buffer the next call overwrites, so it aliases + neither an input nor anything a later call writes. + """ + return _MANAGERS[key]._wrapper(label).run(q, k, v) + + +@ragged_attention.register_fake +def _(q, k, v, key, label): + return q.new_empty(q.shape, dtype=_MANAGERS[key]._dtype) + class _FlashInferRaggedBase: """What the self- and cross-attention managers share: persistent wrappers per @@ -38,7 +63,10 @@ def __init__( device: torch.device, dtype: torch.dtype, config: RaggedAttentionConfig, + key: str, ): + self._key = key + _MANAGERS[key] = self self._config = config self._device = device self._dtype = dtype @@ -218,18 +246,20 @@ def _wrapper(self, label: str | None) -> _RaggedPrefillBase: ) return wrapper - @torch.compiler.disable def run( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, label: str | None = None, ) -> torch.Tensor: """One layer's attention through the plan for ``label``. - Not behind a custom op, unlike the paged backend's ``run``: the ragged - caller (an encoder tower) has no per-layer KV write to keep in the same - graph, so the break this costs is one per layer of a region that is - CUDA-graph captured rather than compiled. + Traced by ``torch.compile``, this is the ``mstar::ragged_attention`` op: + the region keeps one graph, and Inductor fuses the pointwise work on + either side of the attention instead of breaking per layer. Called + eagerly, including in a captured region that is not compiled, it is the + plain kernel call, with no op dispatch. """ + if torch.compiler.is_compiling(): + return ragged_attention(q, k, v, self._key, label) return self._wrapper(label).run(q, k, v) @@ -283,9 +313,10 @@ def __init__( device: torch.device, dtype: torch.dtype, config: RaggedAttentionConfig, + key: str, block_size: int, ): - super().__init__(device=device, dtype=dtype, config=config) + super().__init__(device=device, dtype=dtype, config=config, key=key) self._kwargs["block_size"] = block_size def plan(self, step: AttentionStep, ctx: StepContext): diff --git a/mstar/engine/resources/convenience.py b/mstar/engine/resources/convenience.py index f855765061..ed46c83b15 100644 --- a/mstar/engine/resources/convenience.py +++ b/mstar/engine/resources/convenience.py @@ -154,55 +154,18 @@ def __init__( self, attn: RaggedAttnManager | RaggedCrossAttnManager | RaggedBlockCausalAttnManager, label: str, - opaque: bool = False, ): self.attn = attn self.label = label - # ``opaque``: call through the ``mstar::ragged_attention`` custom op, which - # a compiled region keeps in its graph as one node. Otherwise the call is a - # graph break per layer, which also stops Inductor fusing the pointwise - # work on either side of it: on Wan2.2's compiled, uncaptured DiT the - # breaks cost ~85 ms of a ~1 s step in extra elementwise kernels. - self._handle = _register_ragged(self) if opaque else None + # Not compiler-disabled: a compiled region traces through to the resource's + # ``run``, which keeps the attention in its graph as one custom op. def __call__( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: - if self._handle is not None: - return _ragged_attention_op(q, k, v, self._handle) - return self._run(q, k, v) - - @torch.compiler.disable - def _run(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return self.attn.run(q, k, v, label=self.label) -# Opaque callables by handle: a custom op takes tensors and scalars only, so it -# names its callable by index. One entry per callable built, i.e. per resource -# binding and label, which is a handful per process. -_RAGGED_CALLABLES: list[RaggedAttentionCallable] = [] - - -# Eager even when the callable is built inside a compiled region (a node that -# creates its callables lazily on first forward): traced, the append is not -# replayed but the handle Dynamo read is still baked into the graph. -@torch.compiler.disable -def _register_ragged(fn: RaggedAttentionCallable) -> int: - _RAGGED_CALLABLES.append(fn) - return len(_RAGGED_CALLABLES) - 1 - - -@torch.library.custom_op("mstar::ragged_attention", mutates_args=()) -def _ragged_attention_op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, handle: int) -> torch.Tensor: - # the wrapper returns a fresh tensor (or a clone of its graph buffer), never an alias - return _RAGGED_CALLABLES[handle]._run(q, k, v).to(q.dtype) - - -@_ragged_attention_op.register_fake -def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, handle: int) -> torch.Tensor: - return torch.empty_like(q) - - class Mamba2Callable: """``LinearAttnCallable`` for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor. diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 505215edbf..c969ce0aa5 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -502,7 +502,7 @@ def _attention_fns(self) -> Wan22AttentionFns | None: if self_attn is None: return None self._attn_fns = Wan22AttentionFns( - self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN, opaque=True), cross_attn=None, + self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN), cross_attn=None, ) logger.info("Wan2.2 DiT: self-attention through the engine's ragged attention resource") return self._attn_fns diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 078ed02085..bcf588b930 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -322,16 +322,16 @@ class _Group: assert mgr._kwargs["head_dim"] == HEAD_DIM -def test_opaque_callable_stays_in_the_compiled_graph(): - """``opaque=True`` runs the attention as a custom op: a compiled region - keeps one graph around it (no break) and computes what the eager call does.""" +def test_compiled_caller_keeps_attention_in_its_graph(): + """Traced by ``torch.compile``, the resource's ``run`` is one custom op: the + region keeps one graph around it (no break) and computes what eager does.""" from mstar.engine.resources.convenience import RaggedAttentionCallable mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fn = RaggedAttentionCallable(mgr, "main", opaque=True) + fn = RaggedAttentionCallable(mgr, "main") def region(q, k, v): return fn(q * 1.0, k, v) + 1.0 @@ -340,12 +340,12 @@ def region(q, k, v): explain = torch._dynamo.explain(region)(q, k, v) assert explain.graph_break_count == 0 compiled = torch.compile(region, fullgraph=True) - close(compiled(q, k, v), RaggedAttentionCallable(mgr, "main")(q, k, v) + 1.0) + close(compiled(q, k, v), region(q, k, v)) -def test_opaque_callable_built_inside_a_compiled_region(): - """A node that builds its opaque callables lazily on first forward does so - inside the compiled region; the handle must still name a registered callable.""" +def test_callable_built_inside_a_compiled_region(): + """A node that builds its callables lazily on first forward does so inside + the compiled region; later calls still reach the resource.""" from mstar.engine.resources.convenience import RaggedAttentionCallable mgr = manager() @@ -357,16 +357,27 @@ def test_opaque_callable_built_inside_a_compiled_region(): def region(q, k, v): fn = fns.get(id(mgr)) if fn is None: - fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main", opaque=True) + fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main") return fn(q, k, v) + 1.0 torch._dynamo.reset() compiled = torch.compile(region) - expected = RaggedAttentionCallable(mgr, "main")(q, k, v) + 1.0 + expected = mgr.run(q, k, v, label="main") + 1.0 for _ in range(2): close(compiled(q, k, v), expected) +def test_rebuilding_the_resource_replaces_its_op_entry(): + """The op names its manager by resource key, so a rebuild replaces the entry + a compiled graph reaches instead of adding one per build.""" + from mstar.engine.resources.attn.ragged.flashinfer import _MANAGERS + + manager() + count = len(_MANAGERS) + rebuilt = manager() + assert len(_MANAGERS) == count and _MANAGERS[ATTN] is rebuilt + + def test_plan_then_run_matches_reference(): mgr = manager() spans = [128, 96] From c3a3069cb3734fd9220b3776cd8f814182010613 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:47:53 -0700 Subject: [PATCH 23/59] utils: keep one to_device_nonblocking after merging the fix branches Co-Authored-By: Claude Opus 5.5 --- mstar/utils/device.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/mstar/utils/device.py b/mstar/utils/device.py index 59b6c25d36..2b8e339368 100644 --- a/mstar/utils/device.py +++ b/mstar/utils/device.py @@ -68,12 +68,3 @@ def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: def _device_capability(device: torch.device) -> DeviceCapability: return DeviceCapability(device) - -def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: - """``t`` on ``device`` without a stream sync: a host tensor is pinned, then - copied asynchronously. A tensor already on ``device`` comes back as is.""" - if t.device == torch.device(device): - return t - if t.device.type == "cpu" and not t.is_pinned(): - t = t.pin_memory() - return t.to(device, non_blocking=True) From b45e8f7169a43327cce65daf853844d865e4e76a Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 17:51:09 -0700 Subject: [PATCH 24/59] fused_moe: one router top-k op for the softmax and sigmoid routers TopKRouter's fused softmax top-k (mstar::moe_softmax_topk in fused_moe/router.py) duplicated the one-program-per-token top-k of mstar::moe_router_topk in fused_moe/decode.py. The decode kernel gains a SOFTMAX constexpr: fp32 softmax scores ranked by themselves, versus the sigmoid scores ranked by score + bias. The sigmoid branch is unchanged (same loads for fp32 partials, same launch, PDL and tuning), and the softmax branch keeps the old semantics: descending top-k, lower expert id on ties, optional renormalize (the 1e-20 guard rounds away, since a softmax top-k sum is at least 1/E), fp32 weights and int64 ids. router_topk takes bias as optional and a softmax flag, reads logits of any float dtype in fp32, and its fake now returns fp32 weights for bf16 logits. router.py is deleted, and the duplicated softmax-router tests in test_qwen3_omni_fused_kernels.py go with it; test_moe_router.py keeps them against the shared op. Co-Authored-By: Claude Opus 5.5 --- mstar/model/components/moe.py | 8 +-- mstar/utils/fused_moe/decode.py | 49 ++++++++------ mstar/utils/fused_moe/router.py | 64 ------------------- test/modular/test_fused_moe_decode.py | 15 +++-- test/modular/test_moe_router.py | 6 +- test/modular/test_qwen3_omni_fused_kernels.py | 41 ------------ 6 files changed, 47 insertions(+), 136 deletions(-) delete mode 100644 mstar/utils/fused_moe/router.py diff --git a/mstar/model/components/moe.py b/mstar/model/components/moe.py index 17d864c63b..9a9fc24f4c 100644 --- a/mstar/model/components/moe.py +++ b/mstar/model/components/moe.py @@ -97,7 +97,7 @@ class TopKRouter(nn.Module): norm_topk_prob: if True, renormalize the top-k probabilities so they sum to 1. fused_softmax_topk: on CUDA, route through one Triton kernel - (``mstar::moe_softmax_topk``) instead of softmax, ``torch.topk`` + (``mstar::moe_router_topk``) instead of softmax, ``torch.topk`` and the renormalize. Equal probabilities go to the lower expert id, where ``torch.topk`` may pick another, so a model opts in once it has checked its outputs. @@ -138,10 +138,10 @@ def forward( hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) router_logits = F.linear(hidden_states, self.weight) if self.fused_softmax_topk and router_logits.is_cuda: - from mstar.utils.fused_moe.router import softmax_topk + from mstar.utils.fused_moe.decode import router_topk - routing_weights, router_indices = softmax_topk( - router_logits, self.top_k, self.norm_topk_prob, + routing_weights, router_indices = router_topk( + router_logits, None, self.top_k, 1.0, self.norm_topk_prob, softmax=True, ) return routing_weights, router_indices, None router_logits = F.softmax(router_logits, dtype=torch.float, dim=-1) diff --git a/mstar/utils/fused_moe/decode.py b/mstar/utils/fused_moe/decode.py index 16159ef3ac..eeafbc9ab4 100644 --- a/mstar/utils/fused_moe/decode.py +++ b/mstar/utils/fused_moe/decode.py @@ -4,6 +4,7 @@ Up to a per-GPU token count (``_Tuning.pair_max_tokens``) gate/up and down are GEMV kernels over (token, expert) pairs; above it experts repeat across the batch, and grouped kernels read each expert's weights once. ``swiglu_limit`` picks the clamped SwiGLU or the plain one. +The router's top-k program, as ``router_topk``, also serves softmax routers (``TopKRouter``). """ from __future__ import annotations @@ -202,9 +203,11 @@ def _router_logits_kernel(x, w, part, T, E: tl.constexpr, H: tl.constexpr, TB: t @triton.jit def _router_topk_kernel(part, bias, w_out, id_out, T, scale, E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, - S: tl.constexpr, NORM: tl.constexpr, PDL: tl.constexpr): - """One program per token: sum the S partials in a fixed order, then sigmoid -> + bias - -> top-K (ties to the lower expert) -> gather the unbiased scores -> normalize. The + S: tl.constexpr, NORM: tl.constexpr, SOFTMAX: tl.constexpr, + PDL: tl.constexpr): + """One program per token: sum the S partials in a fixed order and score them, then + top-K (ties to the lower expert) -> gather the scores -> normalize -> * scale. Sigmoid + scores are picked by score + bias, fp32 softmax scores (SOFTMAX) by themselves. The ``1e-20`` guards an all-zero top-K and rounds away for any sum above 2^-42.""" if PDL: gdc_wait() @@ -214,15 +217,21 @@ def _router_topk_kernel(part, bias, w_out, id_out, T, scale, valid = e < E logit = tl.zeros([EP], tl.float32) for s in tl.static_range(S): - logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0) - score = tl.sigmoid(logit) - biased = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) + logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0).to(tl.float32) + if SOFTMAX: + x = tl.where(valid, logit, float("-inf")) + num = tl.where(valid, tl.exp(x - tl.max(x, 0)), 0.0) + score = num / tl.sum(num, 0) + rank = tl.where(valid, score, float("-inf")) + else: + score = tl.sigmoid(logit) + rank = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) kk = tl.arange(0, KP) ids = tl.zeros([KP], tl.int32) for j in tl.static_range(K): - idx = tl.argmax(biased, 0) + idx = tl.argmax(rank, 0) ids = tl.where(kk == j, idx, ids) - biased = tl.where(e == idx, float("-inf"), biased) + rank = tl.where(e == idx, float("-inf"), rank) ws = tl.sum(tl.where(e[None, :] == ids[:, None], score[None, :], 0.0), 1) ws = tl.where(kk < K, ws, 0.0) if NORM: @@ -489,8 +498,8 @@ def _launch_router(x, w, bias, part, topk_w, topk_ids, scale, normalize, pdl, tu num_warps=tuning.router["warps"]) _router_topk_kernel[(T,)]( part, bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), - K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, PDL=pdl, num_warps=1, - launch_pdl=pdl) + K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, SOFTMAX=False, PDL=pdl, + num_warps=1, launch_pdl=pdl) def _launch_gate_up(x, ids, w13, s13, sw13, act, block_size, limit, pdl, bj, bk, bjs, bks, warps): @@ -646,26 +655,28 @@ def route(x, w, bias, *, top_k, scale, normalize): @torch.library.custom_op("mstar::moe_router_topk", mutates_args=()) -def router_topk(logits: torch.Tensor, bias: torch.Tensor, top_k: int, scale: float, - normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: - """``route``'s top-k program on precomputed fp32 ``logits (T, E)``. An op opaque to - dynamo, so a compiled gate keeps it whole.""" +def router_topk(logits: torch.Tensor, bias: torch.Tensor | None, top_k: int, scale: float, + normalize: bool, softmax: bool = False) -> tuple[torch.Tensor, torch.Tensor]: + """``route``'s top-k program on precomputed ``logits (T, E)`` of any float dtype, read in + fp32. ``softmax``: rank and weight by the fp32 softmax, with no ``bias``; otherwise by + the sigmoid router. An op opaque to dynamo, so a compiled gate keeps it whole.""" T, E = logits.shape - assert logits.dtype == torch.float32 and logits.is_contiguous() + assert logits.is_contiguous() and (bias is None) == softmax topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) if T: _router_topk_kernel[(T,)]( - logits, bias, topk_w, topk_ids, T, float(scale), E=E, + logits, logits if bias is None else bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), K=top_k, KP=triton.next_power_of_2(top_k), S=1, - NORM=normalize, PDL=False, num_warps=1) + NORM=normalize, SOFTMAX=softmax, PDL=False, num_warps=1) return topk_w, topk_ids @router_topk.register_fake -def _(logits, bias, top_k, scale, normalize): +def _(logits, bias, top_k, scale, normalize, softmax=False): T = logits.shape[0] - return logits.new_empty(T, top_k), logits.new_empty(T, top_k, dtype=torch.int64) + return (logits.new_empty(T, top_k, dtype=torch.float32), + logits.new_empty(T, top_k, dtype=torch.int64)) def experts(x, w13, s13, w2, s2, shared_w13, shared_w2, topk_w, topk_ids, *, diff --git a/mstar/utils/fused_moe/router.py b/mstar/utils/fused_moe/router.py deleted file mode 100644 index b866614800..0000000000 --- a/mstar/utils/fused_moe/router.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Softmax top-k routing in one kernel. - -``softmax -> topk -> renormalize`` as separate ops is five launches per MoE layer -(``torch.topk`` alone is two, a radix select and a sort), which at decode batch sizes -costs more than the router GEMM. One program per token does it in registers. -""" -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _softmax_topk_kernel(logits, w_out, id_out, stride_t, - E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, - NORM: tl.constexpr): - """fp32 softmax over the E logits, the K largest probabilities in descending order, - optionally divided by their sum. Equal probabilities (common: the logits are bf16) go - to the lower expert first; ``torch.topk`` orders such ties arbitrarily, so the two can - pick different experts at a tie, both equally valid.""" - t = tl.program_id(0) - e = tl.arange(0, EP) - valid = e < E - x = tl.load(logits + t * stride_t + e, mask=valid, other=float("-inf")).to(tl.float32) - x = x - tl.max(x, 0) - num = tl.where(valid, tl.exp(x), 0.0) - prob = num / tl.sum(num, 0) - cand = tl.where(valid, prob, -1.0) - kk = tl.arange(0, KP) - ids = tl.zeros([KP], tl.int32) - ws = tl.zeros([KP], tl.float32) - for j in tl.static_range(K): - idx = tl.argmax(cand, 0) - ws = tl.where(kk == j, tl.max(cand, 0), ws) - ids = tl.where(kk == j, idx, ids) - cand = tl.where(e == idx, -1.0, cand) - if NORM: - ws = ws / tl.sum(tl.where(kk < K, ws, 0.0), 0) - tl.store(w_out + t * K + kk, ws, mask=kk < K) - tl.store(id_out + t * K + kk, ids.to(tl.int64), mask=kk < K) - - -@torch.library.custom_op("mstar::moe_softmax_topk", mutates_args=()) -def softmax_topk(logits: torch.Tensor, top_k: int, normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: - """Top-k of ``softmax(logits (T, E))`` as ``(fp32 weights, int64 expert ids)``, the - weights renormalized to sum to one when ``normalize``. Opaque to dynamo, so a compiled - router keeps it whole.""" - T, E = logits.shape - assert logits.stride(1) == 1 - topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) - topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) - if T: - _softmax_topk_kernel[(T,)]( - logits, topk_w, topk_ids, logits.stride(0), E=E, EP=triton.next_power_of_2(E), - K=top_k, KP=triton.next_power_of_2(top_k), NORM=normalize, num_warps=1) - return topk_w, topk_ids - - -@softmax_topk.register_fake -def _(logits, top_k, normalize): - T = logits.shape[0] - return (logits.new_empty(T, top_k, dtype=torch.float32), - logits.new_empty(T, top_k, dtype=torch.int64)) diff --git a/test/modular/test_fused_moe_decode.py b/test/modular/test_fused_moe_decode.py index 26903521d0..67e1ace1da 100644 --- a/test/modular/test_fused_moe_decode.py +++ b/test/modular/test_fused_moe_decode.py @@ -144,16 +144,21 @@ def test_router_launch_follows_the_table(table, launches, monkeypatch): assert logits["BK"] == bk and logits["num_warps"] == tuning.router["warps"] -def test_router_topk_op(launches): - w, ids = decode.router_topk(torch.randn(5, 64), torch.zeros(64), 8, 2.5, True) +@pytest.mark.parametrize("softmax", [False, True]) +def test_router_topk_op(launches, softmax): + # the softmax router passes bf16 logits and no bias; the weights are fp32 either way + dtype, bias = (torch.bfloat16, None) if softmax else (torch.float32, torch.zeros(64)) + w, ids = decode.router_topk(torch.randn(5, 64, dtype=dtype), bias, 8, 2.5, True, softmax) assert w.shape == ids.shape == (5, 8) and ids.dtype == torch.int64 (_, grid, _, kw), = launches - assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True + assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True and kw["SOFTMAX"] is softmax from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): - w, ids = decode.router_topk(torch.empty(5, 64), torch.empty(64), 8, 2.5, True) - assert w.shape == (5, 8) and ids.dtype == torch.int64 + fake_bias = None if softmax else torch.empty(64) + w, ids = decode.router_topk(torch.empty(5, 64, dtype=dtype), fake_bias, 8, 2.5, True, + softmax) + assert w.shape == (5, 8) and w.dtype == torch.float32 and ids.dtype == torch.int64 # ---------------------------------------------------------------- CUDA diff --git a/test/modular/test_moe_router.py b/test/modular/test_moe_router.py index 292064be55..36c7d028cf 100644 --- a/test/modular/test_moe_router.py +++ b/test/modular/test_moe_router.py @@ -1,4 +1,4 @@ -"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_softmax_topk``) +"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_router_topk``) against the op sequence it replaces, and where the swept MoE tiles apply.""" from __future__ import annotations @@ -12,7 +12,7 @@ @pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) @pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): - from mstar.utils.fused_moe.router import softmax_topk + from mstar.utils.fused_moe.decode import router_topk torch.manual_seed(1) logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 @@ -21,7 +21,7 @@ def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): if norm: ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) - w, ids = softmax_topk(logits, top_k, norm) + w, ids = router_topk(logits, None, top_k, 1.0, norm, softmax=True) assert ids.dtype == torch.int64 and w.dtype == torch.float32 # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the # selected probabilities, and that every pick is a distinct expert holding them diff --git a/test/modular/test_qwen3_omni_fused_kernels.py b/test/modular/test_qwen3_omni_fused_kernels.py index 522d25a9c4..ce37695554 100644 --- a/test/modular/test_qwen3_omni_fused_kernels.py +++ b/test/modular/test_qwen3_omni_fused_kernels.py @@ -3,7 +3,6 @@ * ``cp_qk_norm_rope_kv_write``: the code predictor's q/k norm + RoPE and its dense-KV write, bit-for-bit what two ``fused_qk_norm_rope`` calls and two strided copies produce. -* ``softmax_topk``: ``TopKRouter``'s softmax -> topk -> renormalize. """ from __future__ import annotations @@ -44,46 +43,6 @@ def test_cp_prologue_matches_the_unfused_ops(bs, seq_len, cache_pos): torch.testing.assert_close(cache, ref_cache, rtol=0, atol=0) -@pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) -@pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) -def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): - from mstar.utils.fused_moe.router import softmax_topk - - torch.manual_seed(1) - logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 - probs = torch.softmax(logits, dim=-1, dtype=torch.float) - ref_w, ref_ids = torch.topk(probs, top_k, dim=-1) - if norm: - ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) - - w, ids = softmax_topk(logits, top_k, norm) - assert ids.dtype == torch.int64 and w.dtype == torch.float32 - # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the - # selected probabilities, and that every pick is a distinct expert holding them - torch.testing.assert_close(w, ref_w, rtol=1e-6, atol=1e-7) - assert (ids.sort(dim=-1).values.diff(dim=-1) > 0).all() - picked = probs.gather(1, ids) - if norm: - picked = picked / picked.sum(dim=-1, keepdim=True) - torch.testing.assert_close(picked, w, rtol=1e-6, atol=1e-7) - # with no ties the choice and its order are torch's - distinct = (ref_w[:, :-1] != ref_w[:, 1:]).all(dim=-1) - assert torch.equal(ids[distinct], ref_ids[distinct]) - - -def test_router_compiles_with_the_fused_topk(): - from mstar.model.components.moe import TopKRouter - - router = TopKRouter(64, 128, 8, fused_softmax_topk=True).cuda().to(torch.bfloat16).requires_grad_(False) - torch.nn.init.normal_(router.weight) - x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) - with torch.inference_mode(): - eager = router(x) - compiled = torch.compile(router, fullgraph=True)(x) - assert torch.equal(eager[1], compiled[1]) - torch.testing.assert_close(eager[0], compiled[0], rtol=0, atol=0) - - @pytest.mark.parametrize("bs", [1, 8, 32]) def test_cp_embed_accumulate_matches_the_unfused_ops(bs): from mstar.model.qwen3_omni.components.cp_kernels import cp_embed_accumulate From 48eea8baa5a4daf62373d6d0eb0e7841d9fb5d43 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:03:54 -0700 Subject: [PATCH 25/59] test: bagel stubs carry the config fields the merged code reads Main's KV layout work (#808) sizes the decode loop by config.max_position_embeddings, which the cfg_prune walk-name tests' config stub predates. The branch's LLMSubmodule reads config.cfg_prune at construction, which main's guidance-sink stub (#831) predates; the real BagelModelConfig always has both. Neither test checks less. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_bagel_cfg_prune.py | 10 ++-------- test/modular/test_bagel_guidance_sink.py | 2 +- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py index 6597393bc5..27142e313b 100644 --- a/test/modular/test_bagel_cfg_prune.py +++ b/test/modular/test_bagel_cfg_prune.py @@ -272,10 +272,7 @@ def test_image_gen_walk_names_the_loop_the_counter_is_read_from(): from mstar.model.bagel.bagel_model import BagelModel model = BagelModel.__new__(BagelModel) - model.config = types.SimpleNamespace( - num_timesteps=N, temperature=1.0, top_k=0, top_p=1.0, - repetition_penalty=1.0, ignore_eos=False, vocab_size=32, - ) + model.config = types.SimpleNamespace(num_timesteps=N, max_position_embeddings=32768) model._has_cfg_parallel = False model._has_llm_disaggregation = False model._image_gen_remote_handoff = False @@ -367,10 +364,7 @@ def test_image_gen_cfg_walk_names_the_loop_too(): from mstar.model.bagel.bagel_model import BagelModel model = BagelModel.__new__(BagelModel) - model.config = types.SimpleNamespace( - num_timesteps=N, temperature=1.0, top_k=0, top_p=1.0, - repetition_penalty=1.0, ignore_eos=False, vocab_size=32, - ) + model.config = types.SimpleNamespace(num_timesteps=N, max_position_embeddings=32768) model._has_cfg_parallel = True model._has_llm_disaggregation = False model._image_gen_remote_handoff = False diff --git a/test/modular/test_bagel_guidance_sink.py b/test/modular/test_bagel_guidance_sink.py index b579462e0f..6606cca103 100644 --- a/test/modular/test_bagel_guidance_sink.py +++ b/test/modular/test_bagel_guidance_sink.py @@ -27,7 +27,7 @@ def _no_transfer(monkeypatch): def _llm() -> LLMSubmodule: return LLMSubmodule( language_model=SimpleNamespace(model=SimpleNamespace(embed_tokens=None), lm_head=None), llm2vae=None, - vae2llm=None, time_embedder=None, latent_pos_embed=None, config=SimpleNamespace(), + vae2llm=None, time_embedder=None, latent_pos_embed=None, config=SimpleNamespace(cfg_prune=True), boi_token_id=1, eoi_token_id=2, bos_token_id=3, eos_token_id=4, ) From 99a0518c5c3defbd83513fb7d739e938fdbb6d4b Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:03:54 -0700 Subject: [PATCH 26/59] test: orpheus and qwen3_tts loop-cap stubs match what the engine hands them Orpheus's LLM reads lm_head.tp_size at construction (the audio-vocab head slice), so its stub head carries one. Qwen3-TTS's talker postprocess reads request_info.graph_walk (the clone prefill's codec items), and the engine always passes a request info, so the test's seeding call passes one for the prefill instead of None. The caps and stop checks they compare are unchanged. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_model_kv_layouts.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index 11fd31c45f..ed7f3cc397 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -49,7 +49,8 @@ def test_orpheus_caps_decode_where_its_stop_check_ends_it(max_tokens): model = _orpheus() llm = OrpheusLLMSubmodule( - language_model=SimpleNamespace(model=SimpleNamespace(embed_tokens=None), lm_head=None), config=model.config, + language_model=SimpleNamespace(model=SimpleNamespace(embed_tokens=None), lm_head=SimpleNamespace(tp_size=1)), + config=model.config, ) cap = model.get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] @@ -523,7 +524,7 @@ def test_qwen3_tts_caps_decode_where_its_stop_check_ends_it(max_tokens): talker = _tiny_qwen3_tts_talker() outputs = {"new_token": [torch.tensor(NOT_EOS)]} talker.request_state("r").add_all(generated_frames=0) - talker.postprocess("r", None, outputs) + talker.postprocess("r", SimpleNamespace(graph_walk="talker_prefill"), outputs) def check_stop(request_id, request_info, outputs): talker.postprocess(request_id, request_info, outputs) From af45008c4259b256abbd31539a9c3759b89e7d21 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:25:51 -0700 Subject: [PATCH 27/59] BAGEL combine: take the CFG branch count from the branches that ran when its own counter isn't tracking When combine_cfg's loop counter disagrees with its first-iteration test it has no count of its own, and it read all three branches. In t2i the image branch sits out every step, so the request failed where the device-side combine had produced the right image. Take the count from the leading branches that arrived instead; fail only when the combine's own count reads a branch that sat out. Co-Authored-By: Claude Opus 5.5 --- mstar/model/bagel/submodules.py | 31 +++++++++++++++++----------- test/modular/test_bagel_cfg_prune.py | 16 +++++++++++--- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index 0a610705b5..58eb504c8b 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -2022,18 +2022,25 @@ def prepare_inputs( first_iter = "latents" not in inputs or not inputs["latents"] n_branches = self._cfg_branches(fwd_info, first_iter) - # The branches this step's combine reads; None reads all three (device-side interval test) - reads = CFG_BRANCH_LABELS[:n_branches] - result = {"v_main": inputs["v_main"][0]} - for label in CFG_BRANCH_LABELS[1:]: - v = inputs[f"v_{label}"][0] - if v.shape[0] > 0: - result[f"v_{label}"] = v - elif label in reads: - raise ValueError( - f"BAGEL parallel CFG: LLM_{label} sat out a step whose combine reads " - f"{list(reads)}; the nodes' loop counters disagree" - ) + # Branches run as a prefix of CFG_BRANCH_LABELS, so the leading non-empty ones ran + n_arrived = next( + (i for i, label in enumerate(CFG_BRANCH_LABELS) if inputs[f"v_{label}"][0].shape[0] == 0), + len(CFG_BRANCH_LABELS), + ) + if n_branches is None: + if n_arrived < len(CFG_BRANCH_LABELS): + # this node's counter isn't tracking the loop; the branch nodes counted from theirs + n_branches = n_arrived + elif n_branches > n_arrived: + raise ValueError( + f"BAGEL parallel CFG: LLM_{CFG_BRANCH_LABELS[n_arrived]} sat out a step whose " + f"combine reads {list(CFG_BRANCH_LABELS[:n_branches])}; the nodes' loop counters disagree" + ) + result = { + f"v_{label}": inputs[f"v_{label}"][0] + for label in CFG_BRANCH_LABELS + if inputs[f"v_{label}"][0].shape[0] > 0 + } if first_iter: H = fwd_info.step_metadata.get("height", 1024) W = fwd_info.step_metadata.get("width", 1024) diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py index 27142e313b..b503fc40c2 100644 --- a/test/modular/test_bagel_cfg_prune.py +++ b/test/modular/test_bagel_cfg_prune.py @@ -318,9 +318,10 @@ def test_parallel_combine_reduces_to_the_old_combine(): torch.testing.assert_close(new, old, rtol=0, atol=1e-6) -def _combine_inputs(k: int, sat_out: tuple[str, ...]): +def _combine_inputs(k: int, sat_out: tuple[str, ...], counter: int | None = None): """combine_cfg's prepare_inputs at step ``k`` of a t2i request, with the named - branch nodes having sat the step out.""" + branch nodes having sat the step out and the combine's loop counter reading + ``counter`` (default ``k``).""" from mstar.model.bagel.submodules import CombineCFGSubmodule, sat_out_velocity cfg = types.SimpleNamespace(cfg_prune=True, **{k_: _DEFAULTS[k_] for k_ in ( @@ -335,7 +336,7 @@ def _combine_inputs(k: int, sat_out: tuple[str, ...]): inputs.update(latents=[latents], time_index=[torch.full((16,), float(k))]) info = types.SimpleNamespace( step_metadata={"requires_cfg": True}, - dynamic_loop_iter_counts={IMAGE_GEN_LOOP: k}, + dynamic_loop_iter_counts={IMAGE_GEN_LOOP: k if counter is None else counter}, ) node = CombineCFGSubmodule(cfg) return node, node.prepare_inputs("image_gen_cfg", info, inputs) @@ -359,6 +360,15 @@ def test_combine_never_reads_a_branch_that_sat_the_step_out(): with pytest.raises(ValueError, match="cfg_text sat out"): _combine_inputs(k_in, sat_out=("cfg_text", "cfg_img")) + # The combine's own counter isn't tracking the loop (reads 0 on a later step): it takes + # the count from the branches that ran instead of failing the request. + node, prepared = _combine_inputs(k_in, sat_out=("cfg_img",), counter=0) + assert prepared.kwargs["cfg_branches"] == 2 + out = node.forward("image_gen_cfg", None, **prepared.tensor_inputs, **prepared.kwargs) + _, reference = _combine_inputs(k_in, sat_out=("cfg_img",)) + expected = node.forward("image_gen_cfg", None, **reference.tensor_inputs, **reference.kwargs) + torch.testing.assert_close(out["latents"][0], expected["latents"][0]) + def test_image_gen_cfg_walk_names_the_loop_too(): from mstar.model.bagel.bagel_model import BagelModel From 108b1250daa3d1c6e844f9ffe563e68ededaa3cc Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:27:33 -0700 Subject: [PATCH 28/59] NodeSubmodule.get_device: set the cache on a move, never on read A compiled forward that asked for the device traced its first call with the cache empty, filled it, and failed its guard on the second, recompiling once (the Whisper encoder). Set the cache from the moved parameters in _apply instead of clearing it, so it changes only when the engine moves the submodule at load. Co-Authored-By: Claude Opus 5.5 --- mstar/model/submodule_base.py | 15 ++++++++++----- test/modular/test_submodule_alloc_dtype.py | 16 +++++++++++++--- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 3ea2d7f1d1..8382a49784 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -608,19 +608,24 @@ def request_state(self, request_id: str) -> PerRequestState: return state # Cached: prepare_inputs and preprocess ask per row per step, and finding - # a parameter walks the module tree. ``_apply`` (``.to()``, ``.cuda()``, - # ``to_empty()``) clears it. + # a parameter walks the module tree. Set by ``_apply`` (``.to()``, + # ``.cuda()``, ``to_empty()``), which the engine runs on every submodule at + # load, and never on read, so a compiled forward that asks doesn't see it + # change between calls. ``load_state_dict(assign=True)`` and reassigning a + # parameter don't update it; the engine's ``.to()`` after them does. _device_cache: torch.device | None = None def get_device(self): device = self._device_cache if device is None: - device = self._device_cache = next(self.parameters()).device + device = next(self.parameters()).device return device def _apply(self, fn, *args, **kwargs): - self._device_cache = None - return super()._apply(fn, *args, **kwargs) + module = super()._apply(fn, *args, **kwargs) + param = next(self.parameters(), None) + self._device_cache = None if param is None else param.device + return module @abstractmethod def prepare_inputs( diff --git a/test/modular/test_submodule_alloc_dtype.py b/test/modular/test_submodule_alloc_dtype.py index 9fad26d508..be8a07600a 100644 --- a/test/modular/test_submodule_alloc_dtype.py +++ b/test/modular/test_submodule_alloc_dtype.py @@ -56,7 +56,10 @@ def test_to_on_meta_is_in_place_and_zero_alloc(): def test_submodule_device_follows_materialisation(): - """``NodeSubmodule.get_device`` is cached; moving the module clears it.""" + """``NodeSubmodule.get_device`` is cached by a move, which updates it, and a + compiled forward that asks for it doesn't recompile on its second call.""" + from torch._dynamo.utils import counters + from mstar.model.submodule_base import NodeSubmodule class _Sub(NodeSubmodule): @@ -68,10 +71,17 @@ def __init__(self): def prepare_inputs(self, *args, **kwargs): raise NotImplementedError - def forward(self, *args, **kwargs): - raise NotImplementedError + def forward(self, x): + return self.layer(x) + torch.ones(2, device=self.get_device()) sub = _Sub() assert sub.get_device().type == "meta" sub.to_empty(device="cpu") + + torch._dynamo.reset() + counters.clear() + forward = torch.compile(sub.forward, backend="eager") + for _ in range(2): + forward(torch.randn(3, 2)) + assert counters["stats"]["unique_graphs"] == 1 assert sub.get_device().type == "cpu" From 09f20e021165951f6bf0a409bd3a2ac19df4ab24 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:30:12 -0700 Subject: [PATCH 29/59] Whisper, Qwen3-Omni Thinker, Chatterbox T3: declare their stop as a TokenStop Each check_stop was the TokenStop check spelled out (stop on the end token unless ignore_eos, or at max_tokens), so declare it and delete the method. The Thinker, which hands its tokens over as batch rows, now checks the batch in one call. Co-Authored-By: Claude Opus 5.5 --- mstar/model/chatterbox/submodules.py | 20 +++++--------------- mstar/model/qwen3_omni/submodules.py | 21 ++++++--------------- mstar/model/whisper/submodules.py | 20 ++++---------------- test/modular/test_model_kv_layouts.py | 2 +- 4 files changed, 16 insertions(+), 47 deletions(-) diff --git a/mstar/model/chatterbox/submodules.py b/mstar/model/chatterbox/submodules.py index 96491b40cb..f874df5bf5 100644 --- a/mstar/model/chatterbox/submodules.py +++ b/mstar/model/chatterbox/submodules.py @@ -60,6 +60,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) logger = logging.getLogger(__name__) @@ -220,6 +221,10 @@ def __init__( self.t3 = config.t3 self.builtin_voice = builtin_voice self.supports_cfg = not config.is_turbo + self.token_stop = TokenStop( + loop=DECODE_LOOP, buffer=SPEECH_TOKENS, stop_ids=frozenset({config.t3.stop_speech_token}), + sampler=T3_SAMPLER, generated_key="generated", + ) # -- request knobs ------------------------------------------------------- @@ -441,21 +446,6 @@ def postprocess( state = self.request_state(request_id) state.add("generated", int(state.get("generated", 0)) + 1) - def check_stop( - self, request_id: str, request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if SPEECH_TOKENS not in outputs: - return set() - token = int(outputs[SPEECH_TOKENS][0].reshape(-1)[0].item()) - generated = int(self.request_state(request_id).get("generated", 0)) - max_new = request_info.max_tokens - sampling = request_info.resource_configs.get(T3_SAMPLER) - ignore_eos = bool(getattr(sampling, "ignore_eos", False)) - if (not ignore_eos and token == self.t3.stop_speech_token) or generated >= max_new: - return {DECODE_LOOP} - return set() - # -- batching and capture --------------------------------------------------- def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 68334778c7..112abb7110 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -64,6 +64,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) logger = logging.getLogger(__name__) @@ -351,6 +352,11 @@ def __init__( super().__init__() self.model = thinker_model # Qwen3OmniThinkerModel self.config = config + # only a decode row's token ends the loop; a prefill is not in it yet + self.token_stop = TokenStop( + loop="thinker_decode_loop", buffer="new_token", stop_ids=frozenset({config.im_end_token_id}), + sampler=THINKER_SAMPLER, walk="thinker_decode", + ) # Pre-compute inverse frequencies for 3D MRoPE self._inv_freq: torch.Tensor | None = None @@ -1143,21 +1149,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs or request_info.graph_walk != "thinker_decode": - return set() - token = outputs["new_token"][0].item() - ignore_eos= request_info.resource_configs[THINKER_SAMPLER].ignore_eos - eos_token_id = self.config.im_end_token_id - if (not ignore_eos and eos_token_id == token) or \ - (request_info.dynamic_loop_iter_counts.get("thinker_decode_loop", 0) + 1 >= request_info.max_tokens): - return {"thinker_decode_loop"} - return set() - # =================================================================== # 4. TalkerSubmodule (ar engine) -- SECOND MOST COMPLEX diff --git a/mstar/model/whisper/submodules.py b/mstar/model/whisper/submodules.py index e8969d625c..4de63cad6e 100644 --- a/mstar/model/whisper/submodules.py +++ b/mstar/model/whisper/submodules.py @@ -30,6 +30,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.model.whisper.components.decoder import WhisperDecoderModel from mstar.model.whisper.config import ( @@ -124,6 +125,9 @@ def __init__(self, decoder: WhisperDecoderModel, config: WhisperModelConfig): super().__init__() self.decoder = decoder self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=frozenset({config.eos_token_id}), sampler=SAMPLER, + ) self._suppress_ids: torch.Tensor | None = None self._begin_suppress_ids: torch.Tensor | None = None @@ -339,19 +343,3 @@ def postprocess( if "new_token" not in outputs: return outputs["text_inputs"] = outputs["new_token"] - - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - if (not ignore_eos and token == self.config.eos_token_id) or \ - decoded_tokens >= request_info.max_tokens: - return {"decode_loop"} - return set() diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index ed7f3cc397..d66858861a 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -486,7 +486,7 @@ def test_qwen3_omni_caps_the_thinkers_decode_where_its_stop_check_ends_it(max_to cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), max_output_tokens=max_tokens) assert cap["thinker_decode_loop"] == _stops_after( - lambda *args: ThinkerSubmodule.check_stop(_omni_submodule(), *args), "thinker_decode_loop", "thinker_decode", + ThinkerSubmodule(None, _omni_submodule().config).check_stop, "thinker_decode_loop", "thinker_decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, resource_configs={THINKER_SAMPLER: SamplingReqConfig(ignore_eos=True)}, ), "the cap and the stop check end decode on different iterations" From 6faef2bd3be1a4512c17487bd42570f9297ab4fa Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:30:50 -0700 Subject: [PATCH 30/59] BAGEL image_gen shape inputs: hold 64 sizes and build a miss without blocking Eight entries sat below the shipped configs' 16 concurrent requests, so concurrent sizes could evict each other, and a miss built its tensors with a pageable copy and torch.tensor(device=), each waiting on the stream. Build the position ids on the device and stage the text indexes through to_device_nonblocking. Co-Authored-By: Claude Opus 5.5 --- mstar/model/bagel/submodules.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index 58eb504c8b..0977d873fc 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -138,8 +138,10 @@ def vit_patches( # Absent means the old behaviour, all three whenever guidance is on. CFG_BRANCHES = "cfg_branches" -# Output sizes whose image_gen shape inputs an LLM node keeps on the device -_IMAGE_GEN_SHAPES_CACHED = 8 +# Output sizes whose image_gen shape inputs an LLM node keeps on the device: above +# the shipped configs' concurrency (16), so concurrent sizes don't evict each other. +# An entry is the size's position ids and indexes, tens of KB at 1024x1024. +_IMAGE_GEN_SHAPES_CACHED = 64 def _f32(x: float) -> float: @@ -947,8 +949,8 @@ def _get_text_vae_idxs( self, seq_len: int, device: str ): tensor_inputs = {} - tensor_inputs["text_indexes"] = torch.tensor( - [0, seq_len-1], dtype=torch.long, device=device + tensor_inputs["text_indexes"] = to_device_nonblocking( + torch.tensor([0, seq_len-1], dtype=torch.long), device ) tensor_inputs["vae_token_indexes"] = torch.arange( 1, seq_len-1, @@ -1131,7 +1133,8 @@ def _image_gen_shape_inputs(self, H: int, W: int, device) -> dict[str, torch.Ten if cached is None: pos = get_flattened_position_ids_extrapolate( H, W, self.config.latent_downsample, max_num_patches_per_side=self.config.max_latent_size, - ).to(device) + device=device, + ) cached = { "vae_position_ids": pos, **self._get_text_vae_idxs(pos.shape[0] + IMAGE_SENTINELS, device), From c10020b44a2bce66661b0b501b925ab6c71c1ba6 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:32:18 -0700 Subject: [PATCH 31/59] Validate an image request's output_format once, at API-server submit Only BAGEL checked it at intake; FLUX.2-klein and Z-Image failed a bad format after the denoise. Every image-output request passes APIServer.submit_request on either frontend, so check it there with the shared writer's accepted set, and drop BAGEL's own call. Co-Authored-By: Claude Opus 5.5 --- mstar/api_server/entrypoint.py | 3 +++ mstar/model/bagel/bagel_model.py | 4 +--- mstar/model/components/diffusion/image_io.py | 4 ++-- test/modular/test_bagel_postprocess.py | 11 +++++++++-- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index cd6dc9063d..8b798666b2 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -38,6 +38,7 @@ from mstar.api_server.data_worker import PreprocessWorker from mstar.api_server.request_types import APIServerMessage, PreprocessInput, ResultChunk from mstar.communication.communicator import CommProtocol, make_communicator +from mstar.model.components.diffusion.image_io import image_output_format from mstar.model.multimodal import PromptPart from mstar.model.registry import model_init_kwargs from mstar.observability.catalog import STATUS_INTERVAL_S @@ -572,6 +573,8 @@ def submit_request( f"Output modality {names} requires streaming=True; raw frame " "chunks cannot be returned as an aggregated response." ) + if "image" in output_modalities: + image_output_format(model_kwargs) # fail a format the writer can't produce now, not after generation # Register pending request with self.request_lock: diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index a90d8d9b4b..f945cfe0d8 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -95,7 +95,7 @@ Span, prompt_spans, ) -from mstar.model.components.diffusion.image_io import encode_image, image_output_format +from mstar.model.components.diffusion.image_io import encode_image from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -633,8 +633,6 @@ def process_prompt( Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ - if "image" in output_modalities: - image_output_format(kwargs) # fail a bad output_format now, not after the denoise result: NameToTensorList = {} if prompt is not None: diff --git a/mstar/model/components/diffusion/image_io.py b/mstar/model/components/diffusion/image_io.py index a9088373c0..9fe9a44060 100644 --- a/mstar/model/components/diffusion/image_io.py +++ b/mstar/model/components/diffusion/image_io.py @@ -148,8 +148,8 @@ def deflate(index: int, part: bytes) -> bytes: def image_output_format(request_kwargs: dict | None) -> str: """The request's ``output_format`` (default ``png``; ``jpg`` means ``jpeg``), or ValueError - for one :func:`encode_image` cannot write. A model calls it at intake too, so a bad value - fails the request before generation rather than after.""" + for one :func:`encode_image` cannot write. ``APIServer.submit_request`` calls it for every + image-output request, so a bad value fails before generation rather than after.""" fmt = str((request_kwargs or {}).get("output_format") or "png").lower() if fmt == "jpg": fmt = "jpeg" diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py index e731ad4cad..ef0a17a9f8 100644 --- a/test/modular/test_bagel_postprocess.py +++ b/test/modular/test_bagel_postprocess.py @@ -23,6 +23,13 @@ def test_image_png_decodes_to_the_truncated_pixels(): def test_unsupported_output_format_fails_at_intake(): - model = BagelModel.__new__(BagelModel) + """Every image-output request passes the API server's submit, so a format the shared + writer can't produce fails there, for BAGEL and the DiT image models alike.""" + from mstar.api_server.entrypoint import APIServer + + server = APIServer.__new__(APIServer) with pytest.raises(ValueError, match="output_format"): - model.process_prompt(None, [], ["image"], output_format="gif") + server.submit_request( + text="a cat", input_modalities=["text"], output_modalities=["image"], + model_kwargs={"output_format": "gif"}, streaming=False, + ) From 935ea414999ba0c95dbc9987b275697693a84753 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:32:49 -0700 Subject: [PATCH 32/59] Wan2.2 DiT: drop the missing-_addmm_activation fallback Every torch the project supports (>=2.9.1,<2.13) has it, so the fallback and its test were dead. Co-Authored-By: Claude Opus 5.5 --- mstar/model/wan22/components/dit.py | 9 +-------- test/modular/test_wan22_model.py | 13 ------------- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/mstar/model/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index d79d09dd7f..61b317a4a6 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -29,7 +29,6 @@ two-linear MLPs ARE the shared ``components.mlp.MLP``. """ -import logging import math from typing import NamedTuple @@ -41,8 +40,6 @@ from mstar.model.components.mlp import MLP from mstar.model.wan22.config import Wan22Config -logger = logging.getLogger(__name__) - # 3D RoPE base frequency (reference WanRotaryPosEmbed default `theta`). ROPE_THETA = 10000.0 @@ -355,11 +352,7 @@ def __init__(self, config: Wan22Config): def set_fuse_ffn_gelu(self, enabled: bool) -> None: """Run every block's FFN GELU in the GEMM epilogue (see ``_ffn_gelu_epilogue``); off by default, so the module alone stays the - reference port. ``torch._addmm_activation`` is private: without it the - blocks keep the unfused GELU.""" - if enabled and not hasattr(torch, "_addmm_activation"): - logger.warning("Wan2.2 DiT: torch._addmm_activation is missing; the FFN GELU stays unfused") - enabled = False + reference port.""" for block in self.blocks: block.fuse_ffn_gelu = enabled diff --git a/test/modular/test_wan22_model.py b/test/modular/test_wan22_model.py index 8ea75170a3..11d4b59ed9 100644 --- a/test/modular/test_wan22_model.py +++ b/test/modular/test_wan22_model.py @@ -177,19 +177,6 @@ def test_wan22_forwards_run_without_autograd(): assert not torch.is_autocast_enabled() - -def test_wan22_dit_keeps_the_unfused_gelu_without_the_private_op(monkeypatch): - """The fused FFN GELU calls ``torch._addmm_activation``, which torch does not promise.""" - from types import SimpleNamespace - - from mstar.model.wan22.components.dit import Wan22DiT - - monkeypatch.delattr(torch, "_addmm_activation", raising=False) - dit = SimpleNamespace(blocks=[SimpleNamespace(fuse_ffn_gelu=False)]) - Wan22DiT.set_fuse_ffn_gelu(dit, True) - assert dit.blocks[0].fuse_ffn_gelu is False - - def test_wan22_encoder_walks_persist_embeddings(): model = _make_model() walks = model.get_graph_walk_graphs() From 24e071d786effc1d649cffaec1cbe002ac0958b8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:33:24 -0700 Subject: [PATCH 33/59] Move to_device_nonblocking into mstar.utils.h2d beside PinnedStager It is a host-to-device staging helper, so it lives with the others. Co-Authored-By: Claude Opus 5.5 --- mstar/model/bagel/submodules.py | 2 +- mstar/model/components/diffusion/denoise_loop.py | 2 +- mstar/model/higgs_audio/submodules.py | 2 +- mstar/model/omnivoice/components/backbone.py | 2 +- mstar/model/qwen3_5/submodules.py | 2 +- mstar/model/submodule_base.py | 2 +- mstar/utils/device.py | 10 ---------- mstar/utils/h2d.py | 10 ++++++++++ 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index 0977d873fc..52a867e129 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -47,7 +47,7 @@ NodeSubmodule, StackingMethod, ) -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index 9ddcc26e48..e120a70007 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -58,7 +58,7 @@ from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 7c2e5caebe..2edec64990 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -43,7 +43,7 @@ TokenStop, device_loopback_enabled, ) -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index ed1907cf06..bb227d1faa 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -34,7 +34,7 @@ import torch.nn.functional as F from torch import nn -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index fc02311cd1..eff201d40d 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -76,7 +76,7 @@ device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 8382a49784..32bdb1f56a 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -13,7 +13,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep -from mstar.utils.device import to_device_nonblocking +from mstar.utils.h2d import to_device_nonblocking from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) diff --git a/mstar/utils/device.py b/mstar/utils/device.py index 2b8e339368..cbf9541448 100644 --- a/mstar/utils/device.py +++ b/mstar/utils/device.py @@ -54,16 +54,6 @@ def is_capturing(self) -> bool: ) -def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: - """A host tensor onto ``device`` without blocking the host. A copy from - pageable memory waits for everything already queued on the stream; a - pinned non-blocking one does not.""" - device = torch.device(device) - if t.device.type != "cpu" or not DeviceCapability.of(device).pin_memory: - return t.to(device) - return t.pin_memory().to(device, non_blocking=True) - - @lru_cache(maxsize=None) def _device_capability(device: torch.device) -> DeviceCapability: return DeviceCapability(device) diff --git a/mstar/utils/h2d.py b/mstar/utils/h2d.py index bc80f15a92..c413baf46f 100644 --- a/mstar/utils/h2d.py +++ b/mstar/utils/h2d.py @@ -18,6 +18,16 @@ from mstar.utils.device import DeviceCapability +def to_device_nonblocking(t: torch.Tensor, device) -> torch.Tensor: + """A host tensor onto ``device`` without blocking the host. A copy from + pageable memory waits for everything already queued on the stream; a + pinned non-blocking one does not.""" + device = torch.device(device) + if t.device.type != "cpu" or not DeviceCapability.of(device).pin_memory: + return t.to(device) + return t.pin_memory().to(device, non_blocking=True) + + class H2DMirror: """What the stager last copied into one device destination. From a6f1a52a0b4c57bcb6a68cd7115cb47f76b54abf Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:36:11 -0700 Subject: [PATCH 34/59] Delete RaggedAttentionCallable: callers bind the label on the resource's run It had become a pass-through to attn.run(q, k, v, label=label). Callers now hold functools.partial(resource.run, label=...), still built once per resource binding. The identity-guard rationale the old docstrings gave doesn't hold for a partial (a fresh one per call compiles one graph), so the comments now say what the memoization buys: one build per binding. Co-Authored-By: Claude Opus 5.5 --- .claude/skills/diffusion-scaffold/SKILL.md | 2 +- docs/adding_models.rst | 2 +- mstar/engine/resources/attn/ragged/config.py | 2 +- mstar/engine/resources/convenience.py | 46 +------------------ .../components/diffusion/denoise_loop.py | 16 +++---- mstar/model/ltx2_5/submodules.py | 7 ++- mstar/model/omnilingual_asr/submodules.py | 10 ++-- mstar/model/wan22/submodules.py | 7 ++- test/modular/test_diffusion_denoise_loop.py | 16 +++---- test/modular/test_ragged_attention.py | 11 ++--- 10 files changed, 34 insertions(+), 85 deletions(-) diff --git a/.claude/skills/diffusion-scaffold/SKILL.md b/.claude/skills/diffusion-scaffold/SKILL.md index b921d45181..28416973f1 100644 --- a/.claude/skills/diffusion-scaffold/SKILL.md +++ b/.claude/skills/diffusion-scaffold/SKILL.md @@ -41,7 +41,7 @@ A recipe such as LTX-2.5's two-stage distilled one (half-resolution denoise → Every attention in the DiT should go through a ragged resource, so the model gets the engine's kernels and its backend can change in config. - `attention_segments(bucket_key)` declares every span a request's layers attend over, as `(label, length)`; the default is one `"main"` span of all its tokens. A layer that attends over a sub-span (a refiner over the image tokens alone, a joint model's per-modality self-attention) gets its own label, and `ragged_for(label)` gives that layer its callable. -- `ragged_for` memoizes one callable per label for the life of a binding. The compiled transformer guards on the identity of the callables it is handed; a fresh one per step recompiles every step. +- `ragged_for` memoizes one callable per label (the resource's `run` with the label bound) for the life of a binding, and a rebinding drops them so none attends through the old resource. - A model with more than one head geometry (LTX-2.5's video 32×128 and audio 32×64 heads) declares one resource per geometry. - Cross-attention between spans of the same request (video queries over text keys) is a separate cross-attention resource kind (`RaggedCrossAttentionSpec`, added in #372) whose step names `(q_label, kv_label)` pairs; see "Cross-attention between spans" in `docs/adding_models.rst`. Don't fall back to SDPA for it because the shapes are fixed inside a bucket. - `joint_attention` in `attention.py` packs per-row q/k/v into the ragged layout (or runs SDPA when no resource is bound, which is what a parity backend uses). diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 3780a6d44d..80fe3221d8 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -772,7 +772,7 @@ step names the ``(q_label, kv_label)`` pairs: Each pair is planned as its own layout, pairing each request's ``q_label`` span with that request's ``kv_label`` span; cross-attention is never causal. A layer attends through -``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``, passing the key +``resource.run(q, k, v, label=cross_label(q_label, kv_label))``, passing the key span's tokens as ``k`` and ``v``. One resource has one head geometry, so a model whose attentions differ in head count or head dim declares one spec per (kind, geometry). LTX-2.5 is the reference (``mstar/model/ltx2_5/submodules.py``). diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 1fa1fdeb10..e258a0756a 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -159,7 +159,7 @@ def resource_class(self) -> "type[Resource]": class RaggedCrossAttentionSpec(_RaggedSpecBase): """Cacheless cross-attention between two spans of each request: a forward's ``(q_label, kv_label)`` pairs (``RaggedCrossAttentionStep``), each run through - ``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``. + ``resource.run(q, k, v, label=cross_label(q_label, kv_label))``. Separate from ``RaggedAttentionSpec`` so a resource is one kind of attention: a model with both declares one of each per head geometry. diff --git a/mstar/engine/resources/convenience.py b/mstar/engine/resources/convenience.py index ed46c83b15..7a26f06c2a 100644 --- a/mstar/engine/resources/convenience.py +++ b/mstar/engine/resources/convenience.py @@ -1,11 +1,6 @@ import torch from mstar.engine.resources.attn.base import AttentionManager -from mstar.engine.resources.attn.ragged.base import ( - RaggedAttnManager, - RaggedBlockCausalAttnManager, - RaggedCrossAttnManager, -) from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.linear_attn.base import LinearAttnManager from mstar.engine.resources.recurrent.pool import RecurrentStatePool @@ -26,7 +21,8 @@ class AttentionCallable: varies its label per layer today; one that needs to should thread the label explicitly rather than use this. - For a cacheless attention, see :class:`RaggedAttentionCallable`. + A cacheless (ragged) attention needs no callable: call the resource's + ``run(q, k, v, label=...)`` with the span's label bound. """ def __init__(self, kv: KVManager, attn: AttentionManager | None=None): @@ -128,44 +124,6 @@ def __call__( ) -class RaggedAttentionCallable: - """`AttentionCallable` for a cacheless attention: ``(q, k, v) -> out`` over - one declared span of a `RaggedAttnManager`. - - No ``kv`` and no layer cursor, because there is neither to carry: a ragged - plan is keyed by the segment label, and the wrapper behind it is per - (bucket, slot, label) rather than per layer. The label is therefore bound - here at construction instead of being a cursor on the resource -- which is - what lets one stack attend over several spans at once (a refiner over the - image tokens next to a pass over all of them), the case `AttentionCallable` - says to thread explicitly. A cross-attention pair of a - ``RaggedCrossAttentionSpec`` resource is bound the same way, under - ``cross_label(q_label, kv_label)``; its ``k`` / ``v`` are then the key span's - tokens, packed in the same request order as ``q``. A - ``RaggedBlockCausalAttentionSpec`` resource binds like a self-attention one. - - One instance per label, held for the life of the resource binding: a - compiled transformer region guards on the identity of the callables it is - handed, so handing it a fresh one per step retraces that frame every step - until Dynamo gives up. - """ - - def __init__( - self, - attn: RaggedAttnManager | RaggedCrossAttnManager | RaggedBlockCausalAttnManager, - label: str, - ): - self.attn = attn - self.label = label - - # Not compiler-disabled: a compiled region traces through to the resource's - # ``run``, which keeps the attention in its graph as one custom op. - def __call__( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - ) -> torch.Tensor: - return self.attn.run(q, k, v, label=self.label) - - class Mamba2Callable: """``LinearAttnCallable`` for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor. diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index e120a70007..2b7dea73fc 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -44,6 +44,7 @@ from __future__ import annotations +import functools import logging from collections import OrderedDict from collections.abc import Hashable, Mapping, Sequence @@ -54,7 +55,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep -from mstar.engine.resources.convenience import RaggedAttentionCallable +from mstar.model.components.diffusion.attention import RaggedAttentionFn from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule @@ -142,7 +143,7 @@ def __init__( self._noise_stagers: dict[torch.dtype, NoiseStager] = {} # One callable per attention label for the life of a binding; see # :meth:`ragged_for`. - self._ragged_fns: dict[str, RaggedAttentionCallable] = {} + self._ragged_fns: dict[str, RaggedAttentionFn] = {} @property def loop_back_names(self) -> tuple[str, ...]: @@ -279,7 +280,7 @@ def layout(self, bucket_key: Hashable, device: torch.device) -> Any: def step_index(self, fwd_info: CurrentForwardPassInfo) -> int: return int(fwd_info.dynamic_loop_iter_counts.get(self.loop_name, 0)) - def _ragged(self) -> RaggedAttentionCallable | None: + def _ragged(self) -> RaggedAttentionFn | None: """The attention callable for the ``"main"`` span, or None (SDPA).""" return self.ragged_for(self.MAIN_SPAN) @@ -287,20 +288,19 @@ def bind_node_resources(self, resources) -> None: super().bind_node_resources(resources) self._ragged_fns.clear() # the memoized callables hold the old resource - def ragged_for(self, label: str) -> RaggedAttentionCallable | None: + def ragged_for(self, label: str) -> RaggedAttentionFn | None: """``(q, k, v) -> out`` over the segments declared under ``label``, or None when no ragged resource is bound (layers then fall back to SDPA). - One callable per label for the life of the binding: the compiled transformer - region guards on the identity of the callables it is handed, and a fresh one - per step would recompile it every step until dynamo gave up.""" + The resource's ``run`` with the label bound, built once per label for the life + of the binding rather than per step.""" if self.attn_resource_key is None: return None if label not in self._ragged_fns: resource = self.node_resources.get(self.attn_resource_key) if resource is None: return None - self._ragged_fns[label] = RaggedAttentionCallable(resource, label) + self._ragged_fns[label] = functools.partial(resource.run, label=label) return self._ragged_fns[label] # --------------------------------------------------------- engine contract diff --git a/mstar/model/ltx2_5/submodules.py b/mstar/model/ltx2_5/submodules.py index 9205c32a31..626223a24f 100644 --- a/mstar/model/ltx2_5/submodules.py +++ b/mstar/model/ltx2_5/submodules.py @@ -31,7 +31,6 @@ SubmoduleStep, cross_label, ) -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.components.batched_rows import BatchedRows from mstar.model.components.diffusion.attention import joint_attention, sdpa_attention from mstar.model.components.diffusion.compile_utils import compile_transformer_forward @@ -303,15 +302,15 @@ def bind_node_resources(self, resources) -> None: self._attends = None # the memoized callables hold the old resources def attends(self) -> LTX2Attends: - """The six attentions' callables, built once per resource binding: a compiled - region guards on their identity (see ``DenoiseLoopSubmodule.ragged_for``).""" + """The six attentions' callables, built once per resource binding rather than + per step.""" if self._attends is None: def on(key, label): resource = self.node_resources.get(key) if self.use_ragged_attention else None if resource is None: return sdpa_attention # every row has the bucket's spans, so [B, L, H, D] packs per request - return functools.partial(joint_attention, ragged=RaggedAttentionCallable(resource, label)) + return functools.partial(joint_attention, ragged=functools.partial(resource.run, label=label)) self._attends = LTX2Attends( video_self=on(VIDEO_ATTN, VIDEO_SPAN), audio_self=on(AUDIO_ATTN, AUDIO_SPAN), diff --git a/mstar/model/omnilingual_asr/submodules.py b/mstar/model/omnilingual_asr/submodules.py index fae5b48cad..221003f49b 100644 --- a/mstar/model/omnilingual_asr/submodules.py +++ b/mstar/model/omnilingual_asr/submodules.py @@ -19,7 +19,8 @@ from __future__ import annotations -from collections.abc import Mapping +import functools +from collections.abc import Callable, Mapping from contextlib import contextmanager from typing import Any @@ -35,7 +36,6 @@ ) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources.attn.config import AttentionStep -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep from mstar.model.omnilingual_asr.components.wav2vec2 import Wav2Vec2CTC from mstar.model.omnilingual_asr.config import ENCODER_ATTN, SAMPLE_RATE, TOKEN_IDS, WAVEFORM @@ -92,7 +92,7 @@ def __init__(self, model: Wav2Vec2CTC, max_batch_size: int, max_batch_seconds: f self._compiled_features = with_inductor_overrides( torch.compile(model.encoder_frontend.feature_extractor, dynamic=True) ) - self._attn_fns: dict[int, RaggedAttentionCallable] = {} + self._attn_fns: dict[int, Callable] = {} def warm_compile(self, device: str) -> None: """Compile the feature extractor at load rather than on a live request: @@ -179,12 +179,12 @@ def _capture_encoder(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: attn = self._attn_fn(inp.resources[ENCODER_ATTN]) return {"ids": self.model.encode_packed(inp.static_inputs["hidden"], attn).argmax(dim=-1)} - def _attn_fn(self, resource) -> RaggedAttentionCallable: + def _attn_fn(self, resource) -> Callable: """The resource's attention over the step's segments, one callable per bound resource.""" fn = self._attn_fns.get(id(resource)) if fn is None: - fn = self._attn_fns[id(resource)] = RaggedAttentionCallable(resource, "main") + fn = self._attn_fns[id(resource)] = functools.partial(resource.run, label="main") return fn def max_batch_size(self, graph_walk: str) -> int: diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 1a09d2da18..f171a2c8cd 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -17,6 +17,7 @@ the v1 engine supplies neither the autocast nor the no-grad wrapper. """ +import functools import logging import os from contextlib import contextmanager @@ -28,7 +29,6 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule from mstar.model.wan22.components.dit import Wan22AttentionFns from mstar.model.wan22.components.unipc import ( @@ -497,14 +497,13 @@ def bind_node_resources(self, resources) -> None: self._attn_fns = None # the memoized callables hold the old resources def _attention_fns(self) -> Wan22AttentionFns | None: - """The ragged callables, built once per resource binding (the compiled - transformer guards on their identity), or None to run SDPA.""" + """The ragged callables, built once per resource binding, or None to run SDPA.""" if self._attn_fns is None: self_attn = self.node_resources.get(self.SELF_ATTN) if self_attn is None: return None self._attn_fns = Wan22AttentionFns( - self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN), cross_attn=None, + self_attn=functools.partial(self_attn.run, label=self.VIDEO_SPAN), cross_attn=None, ) logger.info("Wan2.2 DiT: self-attention through the engine's ragged attention resource") return self._attn_fns diff --git a/test/modular/test_diffusion_denoise_loop.py b/test/modular/test_diffusion_denoise_loop.py index 38897e3a1b..4222bd7a00 100644 --- a/test/modular/test_diffusion_denoise_loop.py +++ b/test/modular/test_diffusion_denoise_loop.py @@ -17,7 +17,6 @@ sys.path.insert(0, ".") from mstar.conductor.request_info import CurrentForwardPassInfo # noqa: E402 -from mstar.engine.resources.convenience import RaggedAttentionCallable # noqa: E402 from mstar.model.components.diffusion.denoise_loop import ( # noqa: E402 LATENTS, DenoiseLoopSubmodule, @@ -397,25 +396,22 @@ class OffsetPastTheEnd(ToyDenoise): def test_ragged_for_binds_the_label(): class Resource: def run(self, q, k, v, label=None): - return label + return self, label sub = TwoSpanDenoise(attn_resource_key="dit_attn") resource = Resource() sub.bind_node_resources({"dit_attn": resource}) image = sub.ragged_for("image") - assert isinstance(image, RaggedAttentionCallable) - assert image.attn is resource and image.label == "image" - assert image(None, None, None) == "image" - assert sub._ragged()(None, None, None) == "main" - # one callable object per label, so a compiled region's identity guards hold across steps + assert image(None, None, None) == (resource, "image") + assert sub._ragged()(None, None, None) == (resource, "main") + # built once per label for the binding, not per step assert sub.ragged_for("image") is image and sub._ragged() is sub._ragged() # a rebinding starts over: a callable holding the old resource would attend # through a plan nothing writes to any more new_resource = Resource() sub.bind_node_resources({"dit_attn": new_resource}) - assert sub.ragged_for("image") is not image - assert sub.ragged_for("image").attn is new_resource - assert sub.ragged_for("main")(None, None, None) == "main" + assert sub.ragged_for("image")(None, None, None) == (new_resource, "image") + assert sub.ragged_for("main")(None, None, None) == (new_resource, "main") assert ToyDenoise().ragged_for("image") is None # no resource declared -> SDPA unbound = TwoSpanDenoise(attn_resource_key="dit_attn") assert unbound.ragged_for("image") is None # declared but not bound yet -> SDPA diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index bcf588b930..3eac741837 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -12,6 +12,7 @@ from __future__ import annotations +import functools import sys sys.path.insert(0, ".") @@ -325,13 +326,11 @@ class _Group: def test_compiled_caller_keeps_attention_in_its_graph(): """Traced by ``torch.compile``, the resource's ``run`` is one custom op: the region keeps one graph around it (no break) and computes what eager does.""" - from mstar.engine.resources.convenience import RaggedAttentionCallable - mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fn = RaggedAttentionCallable(mgr, "main") + fn = functools.partial(mgr.run, label="main") def region(q, k, v): return fn(q * 1.0, k, v) + 1.0 @@ -346,18 +345,16 @@ def region(q, k, v): def test_callable_built_inside_a_compiled_region(): """A node that builds its callables lazily on first forward does so inside the compiled region; later calls still reach the resource.""" - from mstar.engine.resources.convenience import RaggedAttentionCallable - mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fns: dict[int, RaggedAttentionCallable] = {} + fns: dict[int, functools.partial] = {} def region(q, k, v): fn = fns.get(id(mgr)) if fn is None: - fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main") + fn = fns[id(mgr)] = functools.partial(mgr.run, label="main") return fn(q, k, v) + 1.0 torch._dynamo.reset() From de211c3d0163427a85405b921241ccc6e2a7f0f5 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:36:42 -0700 Subject: [PATCH 35/59] one_row_per_step: accept capture_batch_sizes given as a tuple The check compared the sequence to [1], so (1,) was refused. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/accelerator_graph_config.py | 2 +- test/modular/test_piecewise_can_lease.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/mstar/engine/accelerator_graph_config.py b/mstar/engine/accelerator_graph_config.py index 1134da1e89..44e1c9a3b0 100644 --- a/mstar/engine/accelerator_graph_config.py +++ b/mstar/engine/accelerator_graph_config.py @@ -320,7 +320,7 @@ def __post_init__(self): if self.one_row_per_step and ( self.get_config_type() is not PiecewiseConfigType.PACKED or not self.lease_before_step - or self.capture_batch_sizes != [1] + or list(self.capture_batch_sizes) != [1] ): raise ValueError( "one_row_per_step needs a packed lease_before_step region " diff --git a/test/modular/test_piecewise_can_lease.py b/test/modular/test_piecewise_can_lease.py index d45207c80d..7bfa063753 100644 --- a/test/modular/test_piecewise_can_lease.py +++ b/test/modular/test_piecewise_can_lease.py @@ -94,6 +94,7 @@ def test_one_row_per_step_is_refused_where_it_means_nothing(batched, overrides): "lease_before_step": True, "one_row_per_step": True, "capture_batch_sizes": [1], } PiecewisePackedConfig(**kwargs, total_tokens=[8]) # the shape it is for + PiecewisePackedConfig(**{**kwargs, "capture_batch_sizes": (1,)}, total_tokens=[8]) with pytest.raises(ValueError, match="one_row_per_step"): if batched: PiecewiseBatchedConfig(**kwargs, seq_len=8) From 57a5f6b43488bc88f1855eb74b66355068743348 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:37:12 -0700 Subject: [PATCH 36/59] test: pin that a compiled region reaches a rebuilt ragged resource, not the private registry The test asserted on _MANAGERS. Check the behaviour instead: a region compiled against one resource attends through the plan of the resource rebuilt under the same key. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_ragged_attention.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 3eac741837..4ffc5162ab 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -364,15 +364,24 @@ def region(q, k, v): close(compiled(q, k, v), expected) -def test_rebuilding_the_resource_replaces_its_op_entry(): - """The op names its manager by resource key, so a rebuild replaces the entry - a compiled graph reaches instead of adding one per build.""" - from mstar.engine.resources.attn.ragged.flashinfer import _MANAGERS +def test_a_compiled_region_attends_through_a_rebuilt_resource(): + """The op names its manager by resource key: after the resource is rebuilt under + the same key, a region compiled against the old one attends through the new one's + plan, not the old one's.""" + q, k, v = qkv(224) + + def region(fn, q, k, v): + return fn(q, k, v) + 1.0 + + torch._dynamo.reset() + compiled = torch.compile(region) + old = manager() + old.plan(step(128, 96), ctx(["r0", "r1"])) + close(compiled(functools.partial(old.run, label="main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) - manager() - count = len(_MANAGERS) rebuilt = manager() - assert len(_MANAGERS) == count and _MANAGERS[ATTN] is rebuilt + rebuilt.plan(step(64, 160), ctx(["r0", "r1"])) + close(compiled(functools.partial(rebuilt.run, label="main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0) def test_plan_then_run_matches_reference(): From d3f042e32ed9534dcbc2320ba60abccf71698b89 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:38:36 -0700 Subject: [PATCH 37/59] test: Higgs loop-back decode prep uploads its row index without a blocking copy Fails on the old Higgs path, which built the index with torch.tensor(..., device=) on every request's first decode step. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_step_prep_host_syncs.py | 51 ++++++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/test/modular/test_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index 221b675c50..4e6f71327c 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -1,4 +1,5 @@ -"""Denoise step prep reads no tensor on the host (Cosmos3, Wan 2.2, OmniVoice). +"""Denoise step prep reads no tensor on the host (Cosmos3, Wan 2.2, OmniVoice), +and AR decode prep copies no pageable host data to the device (Higgs Audio). The loop-carried step counter (``time_index`` / ``step_index``) is the previous step's device output, so an ``.item()`` on it in ``prepare_inputs`` @@ -336,3 +337,51 @@ def test_cosmos3_cfg_active_uses_host_timesteps(): assert reads.reads == [] want = k < len(sched.timesteps) and 400.0 <= float(sched.timesteps[k].item()) <= 1000.0 assert active == want + + +# ---------------------------------------------------------------- Higgs Audio + + +class _PageableH2D(TorchFunctionMode): + """Records every copy of host data to another device that isn't a non-blocking + one (from pinned memory): each waits for the stream's queued work.""" + + def __init__(self): + super().__init__() + self.copies: list[str] = [] + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + device = kwargs.get("device") + if func in (torch.tensor, torch.as_tensor) and device is not None and torch.device(device).type != "cpu": + self.copies.append(func.__name__) + if func is _T.to and args[0].device.type == "cpu" and not kwargs.get("non_blocking"): + target = next((a for a in (*args[1:], device) if isinstance(a, (str, torch.device))), None) + if target is not None and torch.device(target).type != "cpu": + self.copies.append("to") + return func(*args, **kwargs) + + +def test_higgs_loopback_decode_prep_copies_no_pageable_index(monkeypatch): + """A device loop-back decode step whose rows include one fresh off its prefill overlays + that row's token by index: the index goes up through pinned memory, not a blocking copy + (every request's first decode step takes this path). The meta device stands in for the + accelerator.""" + from mstar.model.higgs_audio.config import SAMPLER, HiggsAudioModelConfig + from mstar.model.higgs_audio.submodules import HiggsAudioLLMSubmodule + from mstar.model.submodule_base import ARNodeInputs + from mstar.utils import h2d + + monkeypatch.setattr(h2d.DeviceCapability, "of", staticmethod(lambda device: SimpleNamespace(pin_memory=True))) + monkeypatch.setattr(_T, "pin_memory", lambda self: self) + sub = HiggsAudioLLMSubmodule(None, HiggsAudioModelConfig()) + sampler = SimpleNamespace(loopback_tokens=lambda rids: torch.zeros(len(rids), dtype=torch.long, device="meta")) + engine_inputs = SimpleNamespace(resources={SAMPLER: sampler}, request_ids=["fresh", "looped"]) + inputs = [ + ARNodeInputs(input_ids=torch.zeros(1, dtype=torch.long, device="meta"), input_seq_len=1), + ARNodeInputs(input_seq_len=1), + ] + with _PageableH2D() as h2d_copies: + out = sub.preprocess("decode", engine_inputs, inputs) + assert h2d_copies.copies == [] + assert out["input_ids"].shape == (2,) From 3d87ff74b11c2d8a4d233f1a1b20d7ae7ec00cc4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:39:36 -0700 Subject: [PATCH 38/59] test: the Qwen3 text encoder builds one rotary table across prompt lengths Fails on the per-length table cache, which allocated a device table for every new prompt length; also checks a slice matches a table built at that length and that a prompt past the table is refused. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_diffusion_numerics.py | 27 +++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/test/modular/test_diffusion_numerics.py b/test/modular/test_diffusion_numerics.py index 62fa6bdda9..5b1bf8e5cd 100644 --- a/test/modular/test_diffusion_numerics.py +++ b/test/modular/test_diffusion_numerics.py @@ -76,6 +76,33 @@ def test_qwen3_rotary_tables_match_hf(seq, head_dim, theta): torch.testing.assert_close(sin, sin_ref[0], rtol=0, atol=0) +def test_qwen3_encoder_builds_one_rotary_table_across_prompt_lengths(monkeypatch): + """The encoder builds its rotary table once per device, at its longest length, and + slices it: a table per prompt length was a device allocation per new length, never + freed. A slice computes what a table built at that length does.""" + from mstar.model.components.diffusion.qwen3 import encoder as qwen3_encoder + + torch.manual_seed(0) + dims = dict( + vocab_size=32, hidden_size=16, intermediate_size=32, num_heads=2, num_kv_heads=1, + head_dim=8, rms_norm_eps=1e-6, rope_theta=10_000.0, tap_layers=(1,), + ) + encoder = qwen3_encoder.Qwen3HiddenStateEncoder(**dims, max_positions=16).eval() + builds = [] + build = qwen3_encoder.qwen3_rotary_tables + monkeypatch.setattr(qwen3_encoder, "qwen3_rotary_tables", lambda *a: builds.append(a) or build(*a)) + prompts = [torch.randint(0, 32, (1, seq)) for seq in (3, 7, 16)] + outputs = [encoder(ids, torch.ones_like(ids, dtype=torch.bool), padded=False) for ids in prompts] + assert len(builds) == 1 + for ids, out in zip(prompts, outputs, strict=True): + exact = qwen3_encoder.Qwen3HiddenStateEncoder(**dims, max_positions=ids.shape[1]).eval() + exact.load_state_dict(encoder.state_dict()) + want = exact(ids, torch.ones_like(ids, dtype=torch.bool), padded=False) + torch.testing.assert_close(out, want, rtol=0, atol=0) + with pytest.raises(ValueError, match="prompt tokens"): + encoder(torch.zeros(1, 17, dtype=torch.long), torch.ones(1, 17, dtype=torch.bool), padded=False) + + def _decode(data: bytes) -> torch.Tensor: import io From b565b3c88372a2211e07207b770a8db925e9ebbe Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 18:42:00 -0700 Subject: [PATCH 39/59] OmniVoice: run the backbone's RMSNorms through the shared norm, not flashinfer The backbone wraps Hugging Face modules whose norms aren't M*'s RMSNorm, so it called flashinfer's rmsnorm and fused_add_rmsnorm directly. RMSNorm's forward and forward_residual now have function forms, rms_norm and add_rms_norm, that take a weight and eps; the backbone uses add_rms_norm for its residual norms and the mstar::rms_norm op for the rest, the same FlashInfer kernels on CUDA in half precision and the portable path elsewhere. Co-Authored-By: Claude Opus 5.5 --- mstar/model/components/norm.py | 74 ++++++++++++-------- mstar/model/omnivoice/components/backbone.py | 37 +++++----- 2 files changed, 64 insertions(+), 47 deletions(-) diff --git a/mstar/model/components/norm.py b/mstar/model/components/norm.py index cfe9c33ad7..58b94ddc57 100644 --- a/mstar/model/components/norm.py +++ b/mstar/model/components/norm.py @@ -20,6 +20,50 @@ from mstar.engine.resources.rms_norm import run_rms_norm +def rms_norm( + hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, gemma_mode: bool = False, +) -> torch.Tensor: + """``RMSNorm``'s forward, for a norm whose weight lives elsewhere (e.g. a wrapped + Hugging Face module's): ``mstar::rms_norm``, or Gemma's fp32 ``(1 + weight)``.""" + if gemma_mode: + orig_dtype = hidden_states.dtype + x = hidden_states.to(torch.float32) + var = x.square().mean(dim=-1, keepdim=True) + normed = x * torch.rsqrt(var + eps) + normed = normed * (1.0 + weight.to(torch.float32)) + return normed.to(orig_dtype) + # Fused kernels expect 2D input. Reshape/restore so callers can pass + # arbitrary leading dims, such as [batch, sequence, hidden_size]. + orig_shape = hidden_states.shape + flat = hidden_states.reshape(-1, orig_shape[-1]) + return run_rms_norm(flat, weight, eps=eps).reshape(orig_shape) + + +def add_rms_norm( + hidden_states: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, + gemma_mode: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """``r = hidden_states + residual``; returns ``(rms_norm(r), r)``. + + On CUDA in half precision, one FlashInfer kernel, in place: ``residual += + hidden_states``, then ``hidden_states = norm(residual)``, so both inputs are + consumed. Inductor's own fusion of the pair is a reduction over a few rows that + takes ~3x as long at decode batch sizes. The kernel normalizes the sum before + rounding it to half precision, where the unfused pair rounds first. + """ + if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): + from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm + + orig_shape = hidden_states.shape + h = hidden_states.reshape(-1, orig_shape[-1]) + r = residual.reshape(-1, orig_shape[-1]) + fused = gemma_fused_add_rmsnorm if gemma_mode else fused_add_rmsnorm + fused(h, r, weight, eps) + return h.reshape(orig_shape), r.reshape(orig_shape) + r = hidden_states + residual + return rms_norm(r, weight, eps, gemma_mode), r + + class RMSNorm(nn.Module): """RMSNorm with optional Gemma-style ``(1 + weight)`` scaling. @@ -44,19 +88,7 @@ def __init__(self, hidden_size: int, eps: float = 1e-6, gemma_mode: bool = False self.weight = nn.Parameter(init) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.gemma_mode: - orig_dtype = hidden_states.dtype - x = hidden_states.to(torch.float32) - var = x.square().mean(dim=-1, keepdim=True) - normed = x * torch.rsqrt(var + self.variance_epsilon) - normed = normed * (1.0 + self.weight.to(torch.float32)) - return normed.to(orig_dtype) - # Fused kernels expect 2D input. Reshape/restore so callers can pass - # arbitrary leading dims, such as [batch, sequence, hidden_size]. - orig_shape = hidden_states.shape - flat = hidden_states.reshape(-1, orig_shape[-1]) - out = run_rms_norm(flat, self.weight, eps=self.variance_epsilon) - return out.reshape(orig_shape) + return rms_norm(hidden_states, self.weight, self.variance_epsilon, self.gemma_mode) def forward_residual( self, hidden_states: torch.Tensor, residual: torch.Tensor, @@ -78,21 +110,7 @@ def forward_residual( weight_bias=1.0 if self.gemma_mode else 0.0, ) return normed.reshape(orig_shape), r.reshape(orig_shape) - if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): - # One FlashInfer kernel, in place: `residual += hidden_states`, then - # `hidden_states = norm(residual)`. Inductor's own fusion of the - # pair is a reduction over a few rows that takes ~3x as long at - # decode batch sizes. Both inputs are consumed by the caller. - from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm - - orig_shape = hidden_states.shape - h = hidden_states.reshape(-1, orig_shape[-1]) - r = residual.reshape(-1, orig_shape[-1]) - fused = gemma_fused_add_rmsnorm if self.gemma_mode else fused_add_rmsnorm - fused(h, r, self.weight, self.variance_epsilon) - return h.reshape(orig_shape), r.reshape(orig_shape) - r = hidden_states + residual - return self(r), r + return add_rms_norm(hidden_states, residual, self.weight, self.variance_epsilon, self.gemma_mode) def extra_repr(self) -> str: return f"{self.hidden_size}, eps={self.variance_epsilon}, gemma_mode={self.gemma_mode}" diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index bb227d1faa..f77ba81c4a 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -7,11 +7,12 @@ nothing, and the head GEMM runs only at the positions whose logits are read. The per-step kernels are the reference's own fastest mode: ``apply_flashinfer`` -fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE, -RMSNorm and ``silu_and_mul``. ``OmniVoiceBackbone.encode`` drives those fused -weights with its own layer loop rather than through the patched ``llm``: it -norms q and k straight off the fused QKV output (the reference copies each out -first) and folds every residual add into the next RMSNorm. At 16 requests a +fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE +and ``silu_and_mul`` (its RMSNorms go through M*'s shared ones). +``OmniVoiceBackbone.encode`` drives those fused weights with its own layer loop +rather than through the patched ``llm``: it norms q and k straight off the fused +QKV output (the reference copies each out first) and folds every residual add +into the next RMSNorm. At 16 requests a step those copies and adds were a tenth of the step's GPU time. What M* adds on top is what the reference cannot do from inside one @@ -34,6 +35,8 @@ import torch.nn.functional as F from torch import nn +from mstar.engine.resources.rms_norm import run_rms_norm +from mstar.model.components.norm import add_rms_norm from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) @@ -293,30 +296,26 @@ def encode( norm = layer.input_layernorm if residual is None: residual = hidden - x = flashinfer.norm.rmsnorm(hidden, norm.weight, eps=norm.variance_epsilon) + x = run_rms_norm(hidden, norm.weight, eps=norm.variance_epsilon) else: - # residual += hidden; hidden = norm(residual), in place. The kernel - # normalizes the sum before rounding it to half precision, where - # the reference rounds first: logits match it to that rounding, - # not bit for bit. - flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) - x = hidden + # The fused add-norm normalizes the sum before rounding it to half + # precision, where the reference rounds first: logits match it to + # that rounding, not bit for bit. + x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) n, d = x.shape[0], sa.head_dim heads_q, heads_k, _ = (width // d for width in sa._fi_qkv_split) qkv = F.linear(x, sa._fi_w_qkv).view(n, -1, d) # the q/k norms read their heads straight off the fused output - q = flashinfer.norm.rmsnorm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) - k = flashinfer.norm.rmsnorm( - qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon, - ) + q = run_rms_norm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) + k = run_rms_norm(qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon) flashinfer.rope.apply_rope_pos_ids_inplace( q, k, pos_ids, rope_theta=sa._fi_rope_theta, interleave=False, ) hidden = sa.o_proj(attn.run(q, k, qkv[:, heads_q + heads_k:]).reshape(n, -1)) norm = layer.post_attention_layernorm - flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) - hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(hidden, mlp._fi_w_gate_up))) - flashinfer.norm.fused_add_rmsnorm(hidden, residual, llm.norm.weight, eps=llm.norm.variance_epsilon) + x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) + hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(x, mlp._fi_w_gate_up))) + hidden, _ = add_rms_norm(hidden, residual, llm.norm.weight, llm.norm.variance_epsilon) return hidden def head(self, hidden: torch.Tensor, tgt_index: torch.Tensor) -> torch.Tensor: From ab3db58f91610049802449d0b7fcc8e009f00f1e Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:02:29 -0700 Subject: [PATCH 40/59] Revert "Whisper, Qwen3-Omni Thinker, Chatterbox T3: declare their stop as a TokenStop" This reverts commit 09f20e021165951f6bf0a409bd3a2ac19df4ab24. --- mstar/model/chatterbox/submodules.py | 20 +++++++++++++++----- mstar/model/qwen3_omni/submodules.py | 21 +++++++++++++++------ mstar/model/whisper/submodules.py | 20 ++++++++++++++++---- test/modular/test_model_kv_layouts.py | 2 +- 4 files changed, 47 insertions(+), 16 deletions(-) diff --git a/mstar/model/chatterbox/submodules.py b/mstar/model/chatterbox/submodules.py index f874df5bf5..96491b40cb 100644 --- a/mstar/model/chatterbox/submodules.py +++ b/mstar/model/chatterbox/submodules.py @@ -60,7 +60,6 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, ) logger = logging.getLogger(__name__) @@ -221,10 +220,6 @@ def __init__( self.t3 = config.t3 self.builtin_voice = builtin_voice self.supports_cfg = not config.is_turbo - self.token_stop = TokenStop( - loop=DECODE_LOOP, buffer=SPEECH_TOKENS, stop_ids=frozenset({config.t3.stop_speech_token}), - sampler=T3_SAMPLER, generated_key="generated", - ) # -- request knobs ------------------------------------------------------- @@ -446,6 +441,21 @@ def postprocess( state = self.request_state(request_id) state.add("generated", int(state.get("generated", 0)) + 1) + def check_stop( + self, request_id: str, request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if SPEECH_TOKENS not in outputs: + return set() + token = int(outputs[SPEECH_TOKENS][0].reshape(-1)[0].item()) + generated = int(self.request_state(request_id).get("generated", 0)) + max_new = request_info.max_tokens + sampling = request_info.resource_configs.get(T3_SAMPLER) + ignore_eos = bool(getattr(sampling, "ignore_eos", False)) + if (not ignore_eos and token == self.t3.stop_speech_token) or generated >= max_new: + return {DECODE_LOOP} + return set() + # -- batching and capture --------------------------------------------------- def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 112abb7110..68334778c7 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -64,7 +64,6 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, ) logger = logging.getLogger(__name__) @@ -352,11 +351,6 @@ def __init__( super().__init__() self.model = thinker_model # Qwen3OmniThinkerModel self.config = config - # only a decode row's token ends the loop; a prefill is not in it yet - self.token_stop = TokenStop( - loop="thinker_decode_loop", buffer="new_token", stop_ids=frozenset({config.im_end_token_id}), - sampler=THINKER_SAMPLER, walk="thinker_decode", - ) # Pre-compute inverse frequencies for 3D MRoPE self._inv_freq: torch.Tensor | None = None @@ -1149,6 +1143,21 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] + def check_stop( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if "new_token" not in outputs or request_info.graph_walk != "thinker_decode": + return set() + token = outputs["new_token"][0].item() + ignore_eos= request_info.resource_configs[THINKER_SAMPLER].ignore_eos + eos_token_id = self.config.im_end_token_id + if (not ignore_eos and eos_token_id == token) or \ + (request_info.dynamic_loop_iter_counts.get("thinker_decode_loop", 0) + 1 >= request_info.max_tokens): + return {"thinker_decode_loop"} + return set() + # =================================================================== # 4. TalkerSubmodule (ar engine) -- SECOND MOST COMPLEX diff --git a/mstar/model/whisper/submodules.py b/mstar/model/whisper/submodules.py index 4de63cad6e..e8969d625c 100644 --- a/mstar/model/whisper/submodules.py +++ b/mstar/model/whisper/submodules.py @@ -30,7 +30,6 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, ) from mstar.model.whisper.components.decoder import WhisperDecoderModel from mstar.model.whisper.config import ( @@ -125,9 +124,6 @@ def __init__(self, decoder: WhisperDecoderModel, config: WhisperModelConfig): super().__init__() self.decoder = decoder self.config = config - self.token_stop = TokenStop( - loop="decode_loop", buffer="new_token", stop_ids=frozenset({config.eos_token_id}), sampler=SAMPLER, - ) self._suppress_ids: torch.Tensor | None = None self._begin_suppress_ids: torch.Tensor | None = None @@ -343,3 +339,19 @@ def postprocess( if "new_token" not in outputs: return outputs["text_inputs"] = outputs["new_token"] + + def check_stop( + self, + request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if "new_token" not in outputs: + return set() + token = outputs["new_token"][0].item() + ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos + decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 + if (not ignore_eos and token == self.config.eos_token_id) or \ + decoded_tokens >= request_info.max_tokens: + return {"decode_loop"} + return set() diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index d66858861a..ed7f3cc397 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -486,7 +486,7 @@ def test_qwen3_omni_caps_the_thinkers_decode_where_its_stop_check_ends_it(max_to cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), max_output_tokens=max_tokens) assert cap["thinker_decode_loop"] == _stops_after( - ThinkerSubmodule(None, _omni_submodule().config).check_stop, "thinker_decode_loop", "thinker_decode", + lambda *args: ThinkerSubmodule.check_stop(_omni_submodule(), *args), "thinker_decode_loop", "thinker_decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, resource_configs={THINKER_SAMPLER: SamplingReqConfig(ignore_eos=True)}, ), "the cap and the stop check end decode on different iterations" From 4bf4ea0fcead564f820721b53d6fccc9809486a8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:02:29 -0700 Subject: [PATCH 41/59] Revert "Delete RaggedAttentionCallable: callers bind the label on the resource's run" This reverts commit a6f1a52a0b4c57bcb6a68cd7115cb47f76b54abf. --- .claude/skills/diffusion-scaffold/SKILL.md | 2 +- docs/adding_models.rst | 2 +- mstar/engine/resources/attn/ragged/config.py | 2 +- mstar/engine/resources/convenience.py | 46 ++++++++++++++++++- .../components/diffusion/denoise_loop.py | 16 +++---- mstar/model/ltx2_5/submodules.py | 7 +-- mstar/model/omnilingual_asr/submodules.py | 10 ++-- mstar/model/wan22/submodules.py | 7 +-- test/modular/test_diffusion_denoise_loop.py | 16 ++++--- test/modular/test_ragged_attention.py | 11 +++-- 10 files changed, 85 insertions(+), 34 deletions(-) diff --git a/.claude/skills/diffusion-scaffold/SKILL.md b/.claude/skills/diffusion-scaffold/SKILL.md index 28416973f1..b921d45181 100644 --- a/.claude/skills/diffusion-scaffold/SKILL.md +++ b/.claude/skills/diffusion-scaffold/SKILL.md @@ -41,7 +41,7 @@ A recipe such as LTX-2.5's two-stage distilled one (half-resolution denoise → Every attention in the DiT should go through a ragged resource, so the model gets the engine's kernels and its backend can change in config. - `attention_segments(bucket_key)` declares every span a request's layers attend over, as `(label, length)`; the default is one `"main"` span of all its tokens. A layer that attends over a sub-span (a refiner over the image tokens alone, a joint model's per-modality self-attention) gets its own label, and `ragged_for(label)` gives that layer its callable. -- `ragged_for` memoizes one callable per label (the resource's `run` with the label bound) for the life of a binding, and a rebinding drops them so none attends through the old resource. +- `ragged_for` memoizes one callable per label for the life of a binding. The compiled transformer guards on the identity of the callables it is handed; a fresh one per step recompiles every step. - A model with more than one head geometry (LTX-2.5's video 32×128 and audio 32×64 heads) declares one resource per geometry. - Cross-attention between spans of the same request (video queries over text keys) is a separate cross-attention resource kind (`RaggedCrossAttentionSpec`, added in #372) whose step names `(q_label, kv_label)` pairs; see "Cross-attention between spans" in `docs/adding_models.rst`. Don't fall back to SDPA for it because the shapes are fixed inside a bucket. - `joint_attention` in `attention.py` packs per-row q/k/v into the ragged layout (or runs SDPA when no resource is bound, which is what a parity backend uses). diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 80fe3221d8..3780a6d44d 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -772,7 +772,7 @@ step names the ``(q_label, kv_label)`` pairs: Each pair is planned as its own layout, pairing each request's ``q_label`` span with that request's ``kv_label`` span; cross-attention is never causal. A layer attends through -``resource.run(q, k, v, label=cross_label(q_label, kv_label))``, passing the key +``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``, passing the key span's tokens as ``k`` and ``v``. One resource has one head geometry, so a model whose attentions differ in head count or head dim declares one spec per (kind, geometry). LTX-2.5 is the reference (``mstar/model/ltx2_5/submodules.py``). diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index e258a0756a..1fa1fdeb10 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -159,7 +159,7 @@ def resource_class(self) -> "type[Resource]": class RaggedCrossAttentionSpec(_RaggedSpecBase): """Cacheless cross-attention between two spans of each request: a forward's ``(q_label, kv_label)`` pairs (``RaggedCrossAttentionStep``), each run through - ``resource.run(q, k, v, label=cross_label(q_label, kv_label))``. + ``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``. Separate from ``RaggedAttentionSpec`` so a resource is one kind of attention: a model with both declares one of each per head geometry. diff --git a/mstar/engine/resources/convenience.py b/mstar/engine/resources/convenience.py index 7a26f06c2a..ed46c83b15 100644 --- a/mstar/engine/resources/convenience.py +++ b/mstar/engine/resources/convenience.py @@ -1,6 +1,11 @@ import torch from mstar.engine.resources.attn.base import AttentionManager +from mstar.engine.resources.attn.ragged.base import ( + RaggedAttnManager, + RaggedBlockCausalAttnManager, + RaggedCrossAttnManager, +) from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.linear_attn.base import LinearAttnManager from mstar.engine.resources.recurrent.pool import RecurrentStatePool @@ -21,8 +26,7 @@ class AttentionCallable: varies its label per layer today; one that needs to should thread the label explicitly rather than use this. - A cacheless (ragged) attention needs no callable: call the resource's - ``run(q, k, v, label=...)`` with the span's label bound. + For a cacheless attention, see :class:`RaggedAttentionCallable`. """ def __init__(self, kv: KVManager, attn: AttentionManager | None=None): @@ -124,6 +128,44 @@ def __call__( ) +class RaggedAttentionCallable: + """`AttentionCallable` for a cacheless attention: ``(q, k, v) -> out`` over + one declared span of a `RaggedAttnManager`. + + No ``kv`` and no layer cursor, because there is neither to carry: a ragged + plan is keyed by the segment label, and the wrapper behind it is per + (bucket, slot, label) rather than per layer. The label is therefore bound + here at construction instead of being a cursor on the resource -- which is + what lets one stack attend over several spans at once (a refiner over the + image tokens next to a pass over all of them), the case `AttentionCallable` + says to thread explicitly. A cross-attention pair of a + ``RaggedCrossAttentionSpec`` resource is bound the same way, under + ``cross_label(q_label, kv_label)``; its ``k`` / ``v`` are then the key span's + tokens, packed in the same request order as ``q``. A + ``RaggedBlockCausalAttentionSpec`` resource binds like a self-attention one. + + One instance per label, held for the life of the resource binding: a + compiled transformer region guards on the identity of the callables it is + handed, so handing it a fresh one per step retraces that frame every step + until Dynamo gives up. + """ + + def __init__( + self, + attn: RaggedAttnManager | RaggedCrossAttnManager | RaggedBlockCausalAttnManager, + label: str, + ): + self.attn = attn + self.label = label + + # Not compiler-disabled: a compiled region traces through to the resource's + # ``run``, which keeps the attention in its graph as one custom op. + def __call__( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + ) -> torch.Tensor: + return self.attn.run(q, k, v, label=self.label) + + class Mamba2Callable: """``LinearAttnCallable`` for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor. diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index 2b7dea73fc..e120a70007 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -44,7 +44,6 @@ from __future__ import annotations -import functools import logging from collections import OrderedDict from collections.abc import Hashable, Mapping, Sequence @@ -55,7 +54,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep -from mstar.model.components.diffusion.attention import RaggedAttentionFn +from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule @@ -143,7 +142,7 @@ def __init__( self._noise_stagers: dict[torch.dtype, NoiseStager] = {} # One callable per attention label for the life of a binding; see # :meth:`ragged_for`. - self._ragged_fns: dict[str, RaggedAttentionFn] = {} + self._ragged_fns: dict[str, RaggedAttentionCallable] = {} @property def loop_back_names(self) -> tuple[str, ...]: @@ -280,7 +279,7 @@ def layout(self, bucket_key: Hashable, device: torch.device) -> Any: def step_index(self, fwd_info: CurrentForwardPassInfo) -> int: return int(fwd_info.dynamic_loop_iter_counts.get(self.loop_name, 0)) - def _ragged(self) -> RaggedAttentionFn | None: + def _ragged(self) -> RaggedAttentionCallable | None: """The attention callable for the ``"main"`` span, or None (SDPA).""" return self.ragged_for(self.MAIN_SPAN) @@ -288,19 +287,20 @@ def bind_node_resources(self, resources) -> None: super().bind_node_resources(resources) self._ragged_fns.clear() # the memoized callables hold the old resource - def ragged_for(self, label: str) -> RaggedAttentionFn | None: + def ragged_for(self, label: str) -> RaggedAttentionCallable | None: """``(q, k, v) -> out`` over the segments declared under ``label``, or None when no ragged resource is bound (layers then fall back to SDPA). - The resource's ``run`` with the label bound, built once per label for the life - of the binding rather than per step.""" + One callable per label for the life of the binding: the compiled transformer + region guards on the identity of the callables it is handed, and a fresh one + per step would recompile it every step until dynamo gave up.""" if self.attn_resource_key is None: return None if label not in self._ragged_fns: resource = self.node_resources.get(self.attn_resource_key) if resource is None: return None - self._ragged_fns[label] = functools.partial(resource.run, label=label) + self._ragged_fns[label] = RaggedAttentionCallable(resource, label) return self._ragged_fns[label] # --------------------------------------------------------- engine contract diff --git a/mstar/model/ltx2_5/submodules.py b/mstar/model/ltx2_5/submodules.py index 626223a24f..9205c32a31 100644 --- a/mstar/model/ltx2_5/submodules.py +++ b/mstar/model/ltx2_5/submodules.py @@ -31,6 +31,7 @@ SubmoduleStep, cross_label, ) +from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.components.batched_rows import BatchedRows from mstar.model.components.diffusion.attention import joint_attention, sdpa_attention from mstar.model.components.diffusion.compile_utils import compile_transformer_forward @@ -302,15 +303,15 @@ def bind_node_resources(self, resources) -> None: self._attends = None # the memoized callables hold the old resources def attends(self) -> LTX2Attends: - """The six attentions' callables, built once per resource binding rather than - per step.""" + """The six attentions' callables, built once per resource binding: a compiled + region guards on their identity (see ``DenoiseLoopSubmodule.ragged_for``).""" if self._attends is None: def on(key, label): resource = self.node_resources.get(key) if self.use_ragged_attention else None if resource is None: return sdpa_attention # every row has the bucket's spans, so [B, L, H, D] packs per request - return functools.partial(joint_attention, ragged=functools.partial(resource.run, label=label)) + return functools.partial(joint_attention, ragged=RaggedAttentionCallable(resource, label)) self._attends = LTX2Attends( video_self=on(VIDEO_ATTN, VIDEO_SPAN), audio_self=on(AUDIO_ATTN, AUDIO_SPAN), diff --git a/mstar/model/omnilingual_asr/submodules.py b/mstar/model/omnilingual_asr/submodules.py index 221003f49b..fae5b48cad 100644 --- a/mstar/model/omnilingual_asr/submodules.py +++ b/mstar/model/omnilingual_asr/submodules.py @@ -19,8 +19,7 @@ from __future__ import annotations -import functools -from collections.abc import Callable, Mapping +from collections.abc import Mapping from contextlib import contextmanager from typing import Any @@ -36,6 +35,7 @@ ) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources.attn.config import AttentionStep +from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep from mstar.model.omnilingual_asr.components.wav2vec2 import Wav2Vec2CTC from mstar.model.omnilingual_asr.config import ENCODER_ATTN, SAMPLE_RATE, TOKEN_IDS, WAVEFORM @@ -92,7 +92,7 @@ def __init__(self, model: Wav2Vec2CTC, max_batch_size: int, max_batch_seconds: f self._compiled_features = with_inductor_overrides( torch.compile(model.encoder_frontend.feature_extractor, dynamic=True) ) - self._attn_fns: dict[int, Callable] = {} + self._attn_fns: dict[int, RaggedAttentionCallable] = {} def warm_compile(self, device: str) -> None: """Compile the feature extractor at load rather than on a live request: @@ -179,12 +179,12 @@ def _capture_encoder(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: attn = self._attn_fn(inp.resources[ENCODER_ATTN]) return {"ids": self.model.encode_packed(inp.static_inputs["hidden"], attn).argmax(dim=-1)} - def _attn_fn(self, resource) -> Callable: + def _attn_fn(self, resource) -> RaggedAttentionCallable: """The resource's attention over the step's segments, one callable per bound resource.""" fn = self._attn_fns.get(id(resource)) if fn is None: - fn = self._attn_fns[id(resource)] = functools.partial(resource.run, label="main") + fn = self._attn_fns[id(resource)] = RaggedAttentionCallable(resource, "main") return fn def max_batch_size(self, graph_walk: str) -> int: diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index f171a2c8cd..1a09d2da18 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -17,7 +17,6 @@ the v1 engine supplies neither the autocast nor the no-grad wrapper. """ -import functools import logging import os from contextlib import contextmanager @@ -29,6 +28,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep +from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule from mstar.model.wan22.components.dit import Wan22AttentionFns from mstar.model.wan22.components.unipc import ( @@ -497,13 +497,14 @@ def bind_node_resources(self, resources) -> None: self._attn_fns = None # the memoized callables hold the old resources def _attention_fns(self) -> Wan22AttentionFns | None: - """The ragged callables, built once per resource binding, or None to run SDPA.""" + """The ragged callables, built once per resource binding (the compiled + transformer guards on their identity), or None to run SDPA.""" if self._attn_fns is None: self_attn = self.node_resources.get(self.SELF_ATTN) if self_attn is None: return None self._attn_fns = Wan22AttentionFns( - self_attn=functools.partial(self_attn.run, label=self.VIDEO_SPAN), cross_attn=None, + self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN), cross_attn=None, ) logger.info("Wan2.2 DiT: self-attention through the engine's ragged attention resource") return self._attn_fns diff --git a/test/modular/test_diffusion_denoise_loop.py b/test/modular/test_diffusion_denoise_loop.py index 4222bd7a00..38897e3a1b 100644 --- a/test/modular/test_diffusion_denoise_loop.py +++ b/test/modular/test_diffusion_denoise_loop.py @@ -17,6 +17,7 @@ sys.path.insert(0, ".") from mstar.conductor.request_info import CurrentForwardPassInfo # noqa: E402 +from mstar.engine.resources.convenience import RaggedAttentionCallable # noqa: E402 from mstar.model.components.diffusion.denoise_loop import ( # noqa: E402 LATENTS, DenoiseLoopSubmodule, @@ -396,22 +397,25 @@ class OffsetPastTheEnd(ToyDenoise): def test_ragged_for_binds_the_label(): class Resource: def run(self, q, k, v, label=None): - return self, label + return label sub = TwoSpanDenoise(attn_resource_key="dit_attn") resource = Resource() sub.bind_node_resources({"dit_attn": resource}) image = sub.ragged_for("image") - assert image(None, None, None) == (resource, "image") - assert sub._ragged()(None, None, None) == (resource, "main") - # built once per label for the binding, not per step + assert isinstance(image, RaggedAttentionCallable) + assert image.attn is resource and image.label == "image" + assert image(None, None, None) == "image" + assert sub._ragged()(None, None, None) == "main" + # one callable object per label, so a compiled region's identity guards hold across steps assert sub.ragged_for("image") is image and sub._ragged() is sub._ragged() # a rebinding starts over: a callable holding the old resource would attend # through a plan nothing writes to any more new_resource = Resource() sub.bind_node_resources({"dit_attn": new_resource}) - assert sub.ragged_for("image")(None, None, None) == (new_resource, "image") - assert sub.ragged_for("main")(None, None, None) == (new_resource, "main") + assert sub.ragged_for("image") is not image + assert sub.ragged_for("image").attn is new_resource + assert sub.ragged_for("main")(None, None, None) == "main" assert ToyDenoise().ragged_for("image") is None # no resource declared -> SDPA unbound = TwoSpanDenoise(attn_resource_key="dit_attn") assert unbound.ragged_for("image") is None # declared but not bound yet -> SDPA diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 4ffc5162ab..5ca122195a 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -12,7 +12,6 @@ from __future__ import annotations -import functools import sys sys.path.insert(0, ".") @@ -326,11 +325,13 @@ class _Group: def test_compiled_caller_keeps_attention_in_its_graph(): """Traced by ``torch.compile``, the resource's ``run`` is one custom op: the region keeps one graph around it (no break) and computes what eager does.""" + from mstar.engine.resources.convenience import RaggedAttentionCallable + mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fn = functools.partial(mgr.run, label="main") + fn = RaggedAttentionCallable(mgr, "main") def region(q, k, v): return fn(q * 1.0, k, v) + 1.0 @@ -345,16 +346,18 @@ def region(q, k, v): def test_callable_built_inside_a_compiled_region(): """A node that builds its callables lazily on first forward does so inside the compiled region; later calls still reach the resource.""" + from mstar.engine.resources.convenience import RaggedAttentionCallable + mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fns: dict[int, functools.partial] = {} + fns: dict[int, RaggedAttentionCallable] = {} def region(q, k, v): fn = fns.get(id(mgr)) if fn is None: - fn = fns[id(mgr)] = functools.partial(mgr.run, label="main") + fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main") return fn(q, k, v) + 1.0 torch._dynamo.reset() From c30a4c57dc992d66a9d44a23f687ce31f2f8fcfa Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:02:29 -0700 Subject: [PATCH 42/59] Revert "OmniVoice: run the backbone's RMSNorms through the shared norm, not flashinfer" This reverts commit b565b3c88372a2211e07207b770a8db925e9ebbe. --- mstar/model/components/norm.py | 74 ++++++++------------ mstar/model/omnivoice/components/backbone.py | 37 +++++----- 2 files changed, 47 insertions(+), 64 deletions(-) diff --git a/mstar/model/components/norm.py b/mstar/model/components/norm.py index 58b94ddc57..cfe9c33ad7 100644 --- a/mstar/model/components/norm.py +++ b/mstar/model/components/norm.py @@ -20,50 +20,6 @@ from mstar.engine.resources.rms_norm import run_rms_norm -def rms_norm( - hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, gemma_mode: bool = False, -) -> torch.Tensor: - """``RMSNorm``'s forward, for a norm whose weight lives elsewhere (e.g. a wrapped - Hugging Face module's): ``mstar::rms_norm``, or Gemma's fp32 ``(1 + weight)``.""" - if gemma_mode: - orig_dtype = hidden_states.dtype - x = hidden_states.to(torch.float32) - var = x.square().mean(dim=-1, keepdim=True) - normed = x * torch.rsqrt(var + eps) - normed = normed * (1.0 + weight.to(torch.float32)) - return normed.to(orig_dtype) - # Fused kernels expect 2D input. Reshape/restore so callers can pass - # arbitrary leading dims, such as [batch, sequence, hidden_size]. - orig_shape = hidden_states.shape - flat = hidden_states.reshape(-1, orig_shape[-1]) - return run_rms_norm(flat, weight, eps=eps).reshape(orig_shape) - - -def add_rms_norm( - hidden_states: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, - gemma_mode: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """``r = hidden_states + residual``; returns ``(rms_norm(r), r)``. - - On CUDA in half precision, one FlashInfer kernel, in place: ``residual += - hidden_states``, then ``hidden_states = norm(residual)``, so both inputs are - consumed. Inductor's own fusion of the pair is a reduction over a few rows that - takes ~3x as long at decode batch sizes. The kernel normalizes the sum before - rounding it to half precision, where the unfused pair rounds first. - """ - if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): - from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm - - orig_shape = hidden_states.shape - h = hidden_states.reshape(-1, orig_shape[-1]) - r = residual.reshape(-1, orig_shape[-1]) - fused = gemma_fused_add_rmsnorm if gemma_mode else fused_add_rmsnorm - fused(h, r, weight, eps) - return h.reshape(orig_shape), r.reshape(orig_shape) - r = hidden_states + residual - return rms_norm(r, weight, eps, gemma_mode), r - - class RMSNorm(nn.Module): """RMSNorm with optional Gemma-style ``(1 + weight)`` scaling. @@ -88,7 +44,19 @@ def __init__(self, hidden_size: int, eps: float = 1e-6, gemma_mode: bool = False self.weight = nn.Parameter(init) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return rms_norm(hidden_states, self.weight, self.variance_epsilon, self.gemma_mode) + if self.gemma_mode: + orig_dtype = hidden_states.dtype + x = hidden_states.to(torch.float32) + var = x.square().mean(dim=-1, keepdim=True) + normed = x * torch.rsqrt(var + self.variance_epsilon) + normed = normed * (1.0 + self.weight.to(torch.float32)) + return normed.to(orig_dtype) + # Fused kernels expect 2D input. Reshape/restore so callers can pass + # arbitrary leading dims, such as [batch, sequence, hidden_size]. + orig_shape = hidden_states.shape + flat = hidden_states.reshape(-1, orig_shape[-1]) + out = run_rms_norm(flat, self.weight, eps=self.variance_epsilon) + return out.reshape(orig_shape) def forward_residual( self, hidden_states: torch.Tensor, residual: torch.Tensor, @@ -110,7 +78,21 @@ def forward_residual( weight_bias=1.0 if self.gemma_mode else 0.0, ) return normed.reshape(orig_shape), r.reshape(orig_shape) - return add_rms_norm(hidden_states, residual, self.weight, self.variance_epsilon, self.gemma_mode) + if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): + # One FlashInfer kernel, in place: `residual += hidden_states`, then + # `hidden_states = norm(residual)`. Inductor's own fusion of the + # pair is a reduction over a few rows that takes ~3x as long at + # decode batch sizes. Both inputs are consumed by the caller. + from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm + + orig_shape = hidden_states.shape + h = hidden_states.reshape(-1, orig_shape[-1]) + r = residual.reshape(-1, orig_shape[-1]) + fused = gemma_fused_add_rmsnorm if self.gemma_mode else fused_add_rmsnorm + fused(h, r, self.weight, self.variance_epsilon) + return h.reshape(orig_shape), r.reshape(orig_shape) + r = hidden_states + residual + return self(r), r def extra_repr(self) -> str: return f"{self.hidden_size}, eps={self.variance_epsilon}, gemma_mode={self.gemma_mode}" diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index f77ba81c4a..bb227d1faa 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -7,12 +7,11 @@ nothing, and the head GEMM runs only at the positions whose logits are read. The per-step kernels are the reference's own fastest mode: ``apply_flashinfer`` -fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE -and ``silu_and_mul`` (its RMSNorms go through M*'s shared ones). -``OmniVoiceBackbone.encode`` drives those fused weights with its own layer loop -rather than through the patched ``llm``: it norms q and k straight off the fused -QKV output (the reference copies each out first) and folds every residual add -into the next RMSNorm. At 16 requests a +fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE, +RMSNorm and ``silu_and_mul``. ``OmniVoiceBackbone.encode`` drives those fused +weights with its own layer loop rather than through the patched ``llm``: it +norms q and k straight off the fused QKV output (the reference copies each out +first) and folds every residual add into the next RMSNorm. At 16 requests a step those copies and adds were a tenth of the step's GPU time. What M* adds on top is what the reference cannot do from inside one @@ -35,8 +34,6 @@ import torch.nn.functional as F from torch import nn -from mstar.engine.resources.rms_norm import run_rms_norm -from mstar.model.components.norm import add_rms_norm from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) @@ -296,26 +293,30 @@ def encode( norm = layer.input_layernorm if residual is None: residual = hidden - x = run_rms_norm(hidden, norm.weight, eps=norm.variance_epsilon) + x = flashinfer.norm.rmsnorm(hidden, norm.weight, eps=norm.variance_epsilon) else: - # The fused add-norm normalizes the sum before rounding it to half - # precision, where the reference rounds first: logits match it to - # that rounding, not bit for bit. - x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) + # residual += hidden; hidden = norm(residual), in place. The kernel + # normalizes the sum before rounding it to half precision, where + # the reference rounds first: logits match it to that rounding, + # not bit for bit. + flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) + x = hidden n, d = x.shape[0], sa.head_dim heads_q, heads_k, _ = (width // d for width in sa._fi_qkv_split) qkv = F.linear(x, sa._fi_w_qkv).view(n, -1, d) # the q/k norms read their heads straight off the fused output - q = run_rms_norm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) - k = run_rms_norm(qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon) + q = flashinfer.norm.rmsnorm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) + k = flashinfer.norm.rmsnorm( + qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon, + ) flashinfer.rope.apply_rope_pos_ids_inplace( q, k, pos_ids, rope_theta=sa._fi_rope_theta, interleave=False, ) hidden = sa.o_proj(attn.run(q, k, qkv[:, heads_q + heads_k:]).reshape(n, -1)) norm = layer.post_attention_layernorm - x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) - hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(x, mlp._fi_w_gate_up))) - hidden, _ = add_rms_norm(hidden, residual, llm.norm.weight, llm.norm.variance_epsilon) + flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) + hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(hidden, mlp._fi_w_gate_up))) + flashinfer.norm.fused_add_rmsnorm(hidden, residual, llm.norm.weight, eps=llm.norm.variance_epsilon) return hidden def head(self, hidden: torch.Tensor, tgt_index: torch.Tensor) -> torch.Tensor: From c24d9339ea20ec899387d3b1f56e724b3c6cab3f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:02:30 -0700 Subject: [PATCH 43/59] Revert "Validate an image request's output_format once, at API-server submit" This reverts commit c10020b44a2bce66661b0b501b925ab6c71c1ba6. --- mstar/api_server/entrypoint.py | 3 --- mstar/model/bagel/bagel_model.py | 4 +++- mstar/model/components/diffusion/image_io.py | 4 ++-- test/modular/test_bagel_postprocess.py | 11 ++--------- 4 files changed, 7 insertions(+), 15 deletions(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 8b798666b2..cd6dc9063d 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -38,7 +38,6 @@ from mstar.api_server.data_worker import PreprocessWorker from mstar.api_server.request_types import APIServerMessage, PreprocessInput, ResultChunk from mstar.communication.communicator import CommProtocol, make_communicator -from mstar.model.components.diffusion.image_io import image_output_format from mstar.model.multimodal import PromptPart from mstar.model.registry import model_init_kwargs from mstar.observability.catalog import STATUS_INTERVAL_S @@ -573,8 +572,6 @@ def submit_request( f"Output modality {names} requires streaming=True; raw frame " "chunks cannot be returned as an aggregated response." ) - if "image" in output_modalities: - image_output_format(model_kwargs) # fail a format the writer can't produce now, not after generation # Register pending request with self.request_lock: diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index f945cfe0d8..a90d8d9b4b 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -95,7 +95,7 @@ Span, prompt_spans, ) -from mstar.model.components.diffusion.image_io import encode_image +from mstar.model.components.diffusion.image_io import encode_image, image_output_format from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -633,6 +633,8 @@ def process_prompt( Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ + if "image" in output_modalities: + image_output_format(kwargs) # fail a bad output_format now, not after the denoise result: NameToTensorList = {} if prompt is not None: diff --git a/mstar/model/components/diffusion/image_io.py b/mstar/model/components/diffusion/image_io.py index 9fe9a44060..a9088373c0 100644 --- a/mstar/model/components/diffusion/image_io.py +++ b/mstar/model/components/diffusion/image_io.py @@ -148,8 +148,8 @@ def deflate(index: int, part: bytes) -> bytes: def image_output_format(request_kwargs: dict | None) -> str: """The request's ``output_format`` (default ``png``; ``jpg`` means ``jpeg``), or ValueError - for one :func:`encode_image` cannot write. ``APIServer.submit_request`` calls it for every - image-output request, so a bad value fails before generation rather than after.""" + for one :func:`encode_image` cannot write. A model calls it at intake too, so a bad value + fails the request before generation rather than after.""" fmt = str((request_kwargs or {}).get("output_format") or "png").lower() if fmt == "jpg": fmt = "jpeg" diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py index ef0a17a9f8..e731ad4cad 100644 --- a/test/modular/test_bagel_postprocess.py +++ b/test/modular/test_bagel_postprocess.py @@ -23,13 +23,6 @@ def test_image_png_decodes_to_the_truncated_pixels(): def test_unsupported_output_format_fails_at_intake(): - """Every image-output request passes the API server's submit, so a format the shared - writer can't produce fails there, for BAGEL and the DiT image models alike.""" - from mstar.api_server.entrypoint import APIServer - - server = APIServer.__new__(APIServer) + model = BagelModel.__new__(BagelModel) with pytest.raises(ValueError, match="output_format"): - server.submit_request( - text="a cat", input_modalities=["text"], output_modalities=["image"], - model_kwargs={"output_format": "gif"}, streaming=False, - ) + model.process_prompt(None, [], ["image"], output_format="gif") From 289056f52aefaa41feac5f46fe755bba5c41bcb5 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:02:57 -0700 Subject: [PATCH 44/59] Fix the BAGEL combine test's unseeded draws, the reverted callable in the ragged rebuild test, and one_row_per_step with no capture sizes Co-Authored-By: Claude Opus 5.5 --- mstar/engine/accelerator_graph_config.py | 2 +- test/modular/test_bagel_cfg_prune.py | 5 +++-- test/modular/test_ragged_attention.py | 6 ++++-- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/mstar/engine/accelerator_graph_config.py b/mstar/engine/accelerator_graph_config.py index 44e1c9a3b0..6312fd3bd7 100644 --- a/mstar/engine/accelerator_graph_config.py +++ b/mstar/engine/accelerator_graph_config.py @@ -320,7 +320,7 @@ def __post_init__(self): if self.one_row_per_step and ( self.get_config_type() is not PiecewiseConfigType.PACKED or not self.lease_before_step - or list(self.capture_batch_sizes) != [1] + or list(self.capture_batch_sizes or ()) != [1] ): raise ValueError( "one_row_per_step needs a packed lease_before_step region " diff --git a/test/modular/test_bagel_cfg_prune.py b/test/modular/test_bagel_cfg_prune.py index b503fc40c2..2898b11a59 100644 --- a/test/modular/test_bagel_cfg_prune.py +++ b/test/modular/test_bagel_cfg_prune.py @@ -328,9 +328,10 @@ def _combine_inputs(k: int, sat_out: tuple[str, ...], counter: int | None = None "num_timesteps", "timestep_shift", "cfg_text_scale", "cfg_img_scale", "cfg_interval", "cfg_renorm_type", "cfg_renorm_min", )}) - latents = torch.randn(16, 8, generator=torch.Generator().manual_seed(0)) + g = torch.Generator().manual_seed(0) + latents = torch.randn(16, 8, generator=g) inputs = { - name: [sat_out_velocity(latents) if name[2:] in sat_out else torch.randn(16, 8)] + name: [sat_out_velocity(latents) if name[2:] in sat_out else torch.randn(16, 8, generator=g)] for name in ("v_main", "v_cfg_text", "v_cfg_img") } inputs.update(latents=[latents], time_index=[torch.full((16,), float(k))]) diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 5ca122195a..5d1bb9e052 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -371,6 +371,8 @@ def test_a_compiled_region_attends_through_a_rebuilt_resource(): """The op names its manager by resource key: after the resource is rebuilt under the same key, a region compiled against the old one attends through the new one's plan, not the old one's.""" + from mstar.engine.resources.convenience import RaggedAttentionCallable + q, k, v = qkv(224) def region(fn, q, k, v): @@ -380,11 +382,11 @@ def region(fn, q, k, v): compiled = torch.compile(region) old = manager() old.plan(step(128, 96), ctx(["r0", "r1"])) - close(compiled(functools.partial(old.run, label="main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) + close(compiled(RaggedAttentionCallable(old, "main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) rebuilt = manager() rebuilt.plan(step(64, 160), ctx(["r0", "r1"])) - close(compiled(functools.partial(rebuilt.run, label="main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0) + close(compiled(RaggedAttentionCallable(rebuilt, "main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0) def test_plan_then_run_matches_reference(): From 38a7a01741393396f9ae558b1e25527d480aa43c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:10:24 -0700 Subject: [PATCH 45/59] Revert "One token stop check for AR submodules, declared as TokenStop" This reverts commit 781c5a7f83bfc252e630935b6b880bce9497458c. --- docs/adding_models.rst | 5 +- mstar/model/higgs_audio/submodules.py | 53 ++++++++++- mstar/model/orpheus/submodules.py | 48 +++++++++- mstar/model/qwen3_5/submodules.py | 70 ++++++++++++-- mstar/model/qwen3_tts/submodules.py | 62 +++++++++++-- mstar/model/submodule_base.py | 91 ------------------- test/modular/test_model_kv_layouts.py | 7 +- .../test_qwen3_5_check_stop_batched.py | 5 +- 8 files changed, 219 insertions(+), 122 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 3780a6d44d..3264c7e9cb 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -954,10 +954,7 @@ methods are: ``check_stop(...) -> set[str]`` (optional) Runs off the GPU thread, and may read tensor values. Return the names of the ``Loop`` objects to stop, for example after seeing the EOS token. This is how a decode loop - terminates. An AR submodule whose loop ends on a stop token or the token budget sets - ``self.token_stop = TokenStop(...)`` in ``__init__`` instead of writing one; that also - gives it ``check_stop_batched``, which checks a whole step off one host copy of its - tokens. Qwen3.5's ``LLMSubmodule`` is an example. + terminates. ``cleanup_request(self, request_id)`` (optional) Free per-request state held inside the submodule when a request finishes: buffers, diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 2edec64990..4a84991fd1 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -37,10 +37,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, device_loopback_enabled, ) from mstar.utils.h2d import to_device_nonblocking @@ -269,10 +269,6 @@ def __init__(self, llm: nn.Module, config: HiggsAudioModelConfig): super().__init__() self.model = llm self.config = config - self.token_stop = TokenStop( - loop="decode_loop", buffer="new_token", stop_ids=frozenset(config.stop_token_ids), - sampler=SAMPLER, walk="decode", - ) def get_accelerator_graph_configs( self, device: torch.device, tp_world_size: int = 1, @@ -487,6 +483,22 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] + def check_stop( + self, + request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if "new_token" not in outputs or request_info.graph_walk != "decode": + return set() + token = outputs["new_token"][0].item() + ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos + decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 + if (not ignore_eos and token in self.config.stop_token_ids) or \ + decoded_tokens >= request_info.max_tokens: + return {"decode_loop"} + return set() + def step_is_input_free(self, graph_walk: str) -> bool: # a decode row is one token, like the capture template's return graph_walk == "decode" @@ -495,3 +507,34 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # the decode loop's client edge carries the token the stop check # already has on the host; prefill's client edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} + + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for the whole batch off one ``tolist``.""" + if step_walk is not None and step_walk != "decode": + return {} + tokens = host_rows.buffers.get("new_token") + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + stop_ids = self.config.stop_token_ids + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + if step_walk is None and info.graph_walk != "decode": + continue + if ( + values[i] in stop_ids + and not info.resource_configs[SAMPLER].ignore_eos + ) or info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens: + stops[rid] = {"decode_loop"} + return stops diff --git a/mstar/model/orpheus/submodules.py b/mstar/model/orpheus/submodules.py index 75a10dfd3d..3194142765 100644 --- a/mstar/model/orpheus/submodules.py +++ b/mstar/model/orpheus/submodules.py @@ -23,11 +23,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo @@ -53,10 +53,6 @@ def __init__( self.embed_tokens = language_model.model.embed_tokens self.lm_head = language_model.lm_head self.config = config - self.token_stop = TokenStop( - loop="decode_loop", buffer="new_token", - stop_ids=frozenset({config.stop_token_id}), sampler=SAMPLER, - ) # Orpheus only ever emits ```` ids (SNAC codes and the # speech/turn markers): the LM head computes the logits of rows # ``[custom_token_base_id, vocab)`` alone, about 18% of the head, and @@ -226,6 +222,48 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] + def check_stop( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if "new_token" not in outputs: + return set() + token = outputs["new_token"][0].item() + eos_token_id = self.config.stop_token_id + ignore_eos= request_info.resource_configs[SAMPLER].ignore_eos + if (not ignore_eos and eos_token_id == token) or \ + (request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= request_info.max_tokens): + return {"decode_loop"} + return set() + + + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for every row off one host copy of the step's tokens.""" + tokens = host_rows.buffers.get("new_token") + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(-1).tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + eos = self.config.stop_token_id + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + if (values[i] == eos and not info.resource_configs[SAMPLER].ignore_eos) or ( + info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens + ): + stops[rid] = {"decode_loop"} + return stops + class SNACDecoderSubmodule(NodeSubmodule): """SNAC 24kHz streaming decoder submodule. diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index eff201d40d..c527e2a575 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -69,10 +69,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides @@ -154,11 +154,6 @@ def __init__( self.model = model self.config = config self.vision_config = vision_config - # only a decode row's token ends decode_loop; a prefill is not in it yet - self.token_stop = TokenStop( - loop="decode_loop", buffer="new_token", stop_ids=config.stop_token_ids, - sampler=SAMPLER, walk="decode", - ) self._vision_sentinels: tuple[torch.Tensor, torch.Tensor] | None = None # ------------------------------------------------------------------ @@ -657,6 +652,24 @@ def _token_is_real(request_info: CurrentForwardPassInfo) -> bool: "last_prefill", False, ) + def check_stop( + self, + request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + # only a decode row's stop can end decode_loop; a prefill is not in it yet + if request_info.graph_walk != "decode" or "new_token" not in outputs: + return set() + token = outputs["new_token"][0].item() + ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos + hit_eos = not ignore_eos and token in self.config.stop_token_ids + out_of_budget = ( + request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 + >= request_info.max_tokens + ) + return {"decode_loop"} if hit_eos or out_of_budget else set() + def step_is_input_free(self, graph_walk: str) -> bool: # decode rows are one token each, like the template; the prefill walks # feed `prefill_tokens` from the inputs @@ -667,6 +680,51 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # check already has on the host; prefill's new_token edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for every request off one ``tolist``.""" + if step_walk is not None and step_walk != "decode": + return {} # as check_stop: only a decode row's stop ends decode_loop + tokens = host_rows.buffers.get("new_token") + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() + rows = host_rows.request_ids + if len(rows) >= len(request_ids) and all( + a == b for a, b in zip(request_ids, rows, strict=False) + ): + # the common case: the batch in forward order, no row map needed + pairs = zip(request_ids, values, strict=False) + else: + row_of = {rid: i for i, rid in enumerate(rows)} + pairs = ( + (rid, values[i]) for rid in request_ids + if (i := row_of.get(rid)) is not None and i < len(values) + ) # no row, no output: the per-request path skips it too + stop_ids = self.config.stop_token_ids + stops: dict[str, set[str]] = {} + mixed = step_walk is None + for rid, value in pairs: + info = request_infos[rid] + # a mixed step reads each row's walk; a decode step need not + if mixed and info.graph_walk != "decode": + continue + # the set test first: a stop id is rare, the config read behind it + # goes through the wrapper + if ( + value in stop_ids + and not info.resource_configs[SAMPLER].ignore_eos + ) or ( + info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 + >= info.max_tokens + ): + stops[rid] = {"decode_loop"} + return stops # The vision tower's block loop, as a piecewise capture region diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index c8f66ceb8e..35478346f1 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -85,11 +85,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, + HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, - TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo from mstar.utils.h2d import PinnedStager @@ -162,12 +162,6 @@ def __init__( self.talker_config = config.talker self.cp_config = config.talker.code_predictor self.num_codes = config.talker.num_code_groups - # ``postprocess`` counts the frames, the prefill's included - self.token_stop = TokenStop( - loop="talker_decode_loop", buffer="layer0_codes", - stop_ids=frozenset({config.talker.codec_eos_token_id}), - sampler=TALKER_SAMPLER, generated_key="generated_frames", - ) self._suppress_mask: torch.Tensor | None = None self._decode_last_token_indices: torch.Tensor | None = None self._flag_stager = PinnedStager(torch.bool, numel=self.MAX_BATCH_SIZE) @@ -802,6 +796,60 @@ def postprocess( "generated_frames", int(state.get("generated_frames", 0)) + 1 ) + def check_stop( + self, + request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + """Stop the graph loop on codec EOS or the request frame limit. + + This callback runs off the GPU execution thread, so reading the sampled + token with ``item`` is allowed here rather than in ``postprocess``. + """ + if "layer0_codes" not in outputs: + return set() + token = int(outputs["layer0_codes"][0].item()) + generated = int( + self.request_state(request_id).get("generated_frames", 0) + ) + max_tokens = request_info.max_tokens + sampling_config = request_info.resource_configs.get(TALKER_SAMPLER) + ignore_eos = bool(getattr(sampling_config, "ignore_eos", False)) + reached_eos = ( + not ignore_eos and token == self.talker_config.codec_eos_token_id + ) + if reached_eos or generated >= max_tokens: + logger.debug("talker %s: stop after %d frames (eos=%s)", request_id, generated, reached_eos) + return {"talker_decode_loop"} + return set() + + def check_stop_batched( + self, + request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + """``check_stop`` for every row off one host copy of the group-0 codes.""" + codes = host_rows.buffers.get("layer0_codes") + if not torch.is_tensor(codes) or codes.dim() == 0: + return None + values = codes.reshape(-1).tolist() + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + eos = self.talker_config.codec_eos_token_id + stops: dict[str, set[str]] = {} + for rid in request_ids: + i = row_of.get(rid) + if i is None or i >= len(values): + continue # no row, no output: the per-request path skips it too + info = request_infos[rid] + generated = int(self.request_state(rid).get("generated_frames", 0)) + ignore_eos = bool(getattr(info.resource_configs.get(TALKER_SAMPLER), "ignore_eos", False)) + if (values[i] == eos and not ignore_eos) or generated >= info.max_tokens: + stops[rid] = {"talker_decode_loop"} + return stops + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Admit compatible prefill/decode requests to continuous batching. diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 32bdb1f56a..19a30efaea 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -28,28 +28,6 @@ class HostRows(NamedTuple): buffers: dict[str, torch.Tensor] -class TokenStop(NamedTuple): - """When an AR submodule's sampled token ends its decode loop: on a stop - id (unless the request's sampler config sets ``ignore_eos``), or once - the request has generated ``max_tokens`` tokens. - - Declared as ``ARNodeSubmodule.token_stop``; ``check_stop`` and - ``check_stop_batched`` follow from it. - """ - - loop: str - # the output (and ``check_stop_buffers`` entry) holding the step's token - buffer: str - stop_ids: frozenset[int] - # resource key of the sampler whose request config carries ``ignore_eos`` - sampler: str - # the only walk whose rows can stop the loop; None for every walk - walk: str | None = None - # request-state key counting the tokens generated so far; None counts - # the loop's iterations, plus the token this step sampled - generated_key: str | None = None - - if TYPE_CHECKING: from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner @@ -1019,75 +997,6 @@ def cleanup_request(self, request_id: str): class ARNodeSubmodule(NodeSubmodule): - # Set (in ``__init__``, where the stop ids are known) by a submodule whose - # sampled token ends a dynamic loop; the stop checks below follow from it. - token_stop: TokenStop | None = None - - def check_stop( - self, request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - stop = self.token_stop - if stop is None: - return super().check_stop(request_id, request_info, outputs) - if (stop.walk is not None and request_info.graph_walk != stop.walk) or stop.buffer not in outputs: - return set() - token = outputs[stop.buffer][0].item() - return self._token_stops(((request_id, token),), {request_id: request_info}, False).get(request_id, set()) - - def check_stop_batched( - self, request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - stop = self.token_stop - if stop is None: - return super().check_stop_batched(request_ids, request_infos, host_rows, step_walk) - if stop.walk is not None and step_walk is not None and step_walk != stop.walk: - return {} - tokens = host_rows.buffers.get(stop.buffer) - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - rows = host_rows.request_ids - if len(rows) >= len(request_ids) and all( - a == b for a, b in zip(request_ids, rows, strict=False) - ): - # the common case: the batch in forward order, no row map needed - pairs = zip(request_ids, values, strict=False) - else: - row_of = {rid: i for i, rid in enumerate(rows)} - pairs = ( - (rid, values[i]) for rid in request_ids - if (i := row_of.get(rid)) is not None and i < len(values) - ) # no row, no output: the per-request path skips it too - # a mixed step reads each row's walk; a step of one walk need not - return self._token_stops(pairs, request_infos, stop.walk is not None and step_walk is None) - - def _token_stops(self, pairs, request_infos, check_walk: bool) -> dict[str, set[str]]: - """Both stop checks' verdict on ``(rid, token)`` pairs.""" - stop = self.token_stop - stop_ids, loop, key = stop.stop_ids, stop.loop, stop.generated_key - stops: dict[str, set[str]] = {} - for rid, token in pairs: - info = request_infos[rid] - if check_walk and info.graph_walk != stop.walk: - continue - # the set test first: a stop id is rare, and the config read - # behind it goes through the info wrapper - if token in stop_ids and not info.resource_configs[stop.sampler].ignore_eos: - stops[rid] = {loop} - continue - generated = ( - info.dynamic_loop_iter_counts.get(loop, 0) + 1 if key is None - else int(self.request_state(rid).get(key, 0)) - ) - if generated >= info.max_tokens: - stops[rid] = {loop} - return stops - def split_inputs( self, graph_walk: str, diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index ed7f3cc397..cfce625f9d 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -462,12 +462,15 @@ def test_qwen3_5_caps_decode_where_its_stop_check_ends_it(max_tokens): from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel from mstar.model.qwen3_5.submodules import LLMSubmodule - llm = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset())) + llm = SimpleNamespace(config=SimpleNamespace(stop_token_ids=())) + + def check_stop(request_id, request_info, outputs): + return LLMSubmodule.check_stop(llm, request_id, request_info, outputs) cap = Qwen3_5DenseModel.get_loop_caps(object.__new__(Qwen3_5DenseModel), max_output_tokens=max_tokens) assert cap["decode_loop"] == _stops_after( - llm.check_stop, "decode_loop", "decode", 2048, + check_stop, "decode_loop", "decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, ), "the cap and the stop check end decode on different iterations" diff --git a/test/modular/test_qwen3_5_check_stop_batched.py b/test/modular/test_qwen3_5_check_stop_batched.py index f50b2d65bc..0441960842 100644 --- a/test/modular/test_qwen3_5_check_stop_batched.py +++ b/test/modular/test_qwen3_5_check_stop_batched.py @@ -23,7 +23,7 @@ def _info(ignore_eos=False, done=0, max_tokens=100, walk="decode", last_prefill= def _sub(): - return LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) + return SimpleNamespace(config=SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) @pytest.mark.parametrize("shape", ["flat", "column"]) @@ -118,7 +118,8 @@ def test_rows_out_of_batch_order_still_map_by_request(): from mstar.model.qwen3_5.submodules import LLMSubmodule from mstar.model.submodule_base import HostRows - sub = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({7}))) + sub = LLMSubmodule.__new__(LLMSubmodule) + sub.config = SimpleNamespace(stop_token_ids={7}) rows = HostRows(("a", "b", "c"), {"new_token": torch.tensor([1, 7, 1])}) def info(max_tokens, iters=0): From e75ec283edf7dc84425aaa022a9e6f5bf5d3f0dc Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:10:44 -0700 Subject: [PATCH 46/59] Revert "Share the device loop-back decode input ids on ARNodeSubmodule" Keeps the review fix it carried for Higgs: its loop-back row index goes up through to_device_nonblocking instead of a blocking pageable copy. Co-Authored-By: Claude Opus 5.5 --- mstar/model/higgs_audio/submodules.py | 12 +++++++++--- mstar/model/qwen3_5/submodules.py | 23 ++++++++++++++++++++--- mstar/model/submodule_base.py | 23 ----------------------- 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 4a84991fd1..0ea8b5e33b 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -390,9 +390,15 @@ def preprocess( return {"input_ids": torch.cat([inp.input_ids for inp in inputs])} # device loop-back: each row's last sampled token, read by slot; # rows whose previous step was the prefill brought theirs - return {"input_ids": self.loopback_input_ids( - engine_inputs.resources[SAMPLER], engine_inputs.request_ids, inputs, - )} + sampler: SamplerResource = engine_inputs.resources[SAMPLER] + ids = sampler.loopback_tokens(engine_inputs.request_ids) + routed = [i for i, inp in enumerate(inputs) if inp.input_ids is not None] + if routed: + ids.index_copy_( + 0, to_device_nonblocking(torch.tensor(routed), ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + return {"input_ids": ids} return { "input_embeds": torch.cat( [inp.input_embeds for inp in inputs], dim=0 diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index c527e2a575..10448bab06 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -76,6 +76,7 @@ device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides +from mstar.utils.device import DeviceCapability from mstar.utils.h2d import to_device_nonblocking from mstar.utils.profiler import PHASE_PERIOD, phase_record @@ -437,17 +438,33 @@ def preprocess( # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] t0 = perf_counter() if PHASE_PERIOD else 0.0 + ids = sampler.loopback_tokens(engine_inputs.request_ids) + # rows whose last step ran elsewhere brought their token + routed = [ + i for i, inp in enumerate(inputs) + if inp.input_ids is not None and inp.input_seq_len + ] if graph_walk != "decode": # a mixed step packs its decode rows' single tokens beside the # prompts: each row takes its own ids, or its loop-back token - ids = sampler.loopback_tokens(engine_inputs.request_ids) ids = torch.cat([ inp.input_ids.to(ids.dtype) if inp.input_ids is not None else ids[i:i + 1] for i, inp in enumerate(inputs) ]) - else: - ids = self.loopback_input_ids(sampler, engine_inputs.request_ids, inputs) + elif routed: + # Through pinned memory, non-blocking: a pageable + # ``torch.tensor(..., device=cuda)`` is a synchronous copy + # that waits for the queued step, and every request's first + # decode step after its prefill takes this path. + index = torch.tensor( + routed, dtype=torch.long, + pin_memory=DeviceCapability.of(ids.device).pin_memory, + ).to(ids.device, non_blocking=True) + ids.index_copy_( + 0, index, + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) out["input_ids"] = ids if PHASE_PERIOD: phase_record("engine.preprocess.loopback", perf_counter() - t0) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 19a30efaea..2cbc12330d 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -13,7 +13,6 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep -from mstar.utils.h2d import to_device_nonblocking from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) @@ -32,7 +31,6 @@ class HostRows(NamedTuple): from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner from mstar.engine.engine import ExecutingBatch - from mstar.engine.resources.sampler.resource import SamplerResource from mstar.streaming.stream_buffer import StreamChunkInfo @@ -1060,24 +1058,3 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: # input name to tensor pass - - @staticmethod - def loopback_input_ids( - sampler: SamplerResource, request_ids: list[str], inputs: Sequence[ARNodeInputs], - ) -> torch.Tensor: - """A device loop-back decode step's input ids: each row's last sampled - token, read off the sampler on the device, except where a row brought - its own (its previous step ran elsewhere, e.g. the prefill).""" - ids = sampler.loopback_tokens(request_ids) - routed = [ - i for i, inp in enumerate(inputs) - if inp.input_ids is not None and inp.input_seq_len - ] - if routed: - # every request's first decode step after its prefill comes here, - # so the row index goes up without a blocking copy - ids.index_copy_( - 0, to_device_nonblocking(torch.tensor(routed, dtype=torch.long), ids.device), - torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), - ) - return ids From 4327e2238f2475a588030292abffa68989fecb87 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:11:27 -0700 Subject: [PATCH 47/59] Revert "fused_moe: one router top-k op for the softmax and sigmoid routers" This reverts commit b45e8f7169a43327cce65daf853844d865e4e76a. --- mstar/model/components/moe.py | 8 +-- mstar/utils/fused_moe/decode.py | 49 ++++++-------- mstar/utils/fused_moe/router.py | 64 +++++++++++++++++++ test/modular/test_fused_moe_decode.py | 15 ++--- test/modular/test_moe_router.py | 6 +- test/modular/test_qwen3_omni_fused_kernels.py | 41 ++++++++++++ 6 files changed, 136 insertions(+), 47 deletions(-) create mode 100644 mstar/utils/fused_moe/router.py diff --git a/mstar/model/components/moe.py b/mstar/model/components/moe.py index 9a9fc24f4c..17d864c63b 100644 --- a/mstar/model/components/moe.py +++ b/mstar/model/components/moe.py @@ -97,7 +97,7 @@ class TopKRouter(nn.Module): norm_topk_prob: if True, renormalize the top-k probabilities so they sum to 1. fused_softmax_topk: on CUDA, route through one Triton kernel - (``mstar::moe_router_topk``) instead of softmax, ``torch.topk`` + (``mstar::moe_softmax_topk``) instead of softmax, ``torch.topk`` and the renormalize. Equal probabilities go to the lower expert id, where ``torch.topk`` may pick another, so a model opts in once it has checked its outputs. @@ -138,10 +138,10 @@ def forward( hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) router_logits = F.linear(hidden_states, self.weight) if self.fused_softmax_topk and router_logits.is_cuda: - from mstar.utils.fused_moe.decode import router_topk + from mstar.utils.fused_moe.router import softmax_topk - routing_weights, router_indices = router_topk( - router_logits, None, self.top_k, 1.0, self.norm_topk_prob, softmax=True, + routing_weights, router_indices = softmax_topk( + router_logits, self.top_k, self.norm_topk_prob, ) return routing_weights, router_indices, None router_logits = F.softmax(router_logits, dtype=torch.float, dim=-1) diff --git a/mstar/utils/fused_moe/decode.py b/mstar/utils/fused_moe/decode.py index eeafbc9ab4..16159ef3ac 100644 --- a/mstar/utils/fused_moe/decode.py +++ b/mstar/utils/fused_moe/decode.py @@ -4,7 +4,6 @@ Up to a per-GPU token count (``_Tuning.pair_max_tokens``) gate/up and down are GEMV kernels over (token, expert) pairs; above it experts repeat across the batch, and grouped kernels read each expert's weights once. ``swiglu_limit`` picks the clamped SwiGLU or the plain one. -The router's top-k program, as ``router_topk``, also serves softmax routers (``TopKRouter``). """ from __future__ import annotations @@ -203,11 +202,9 @@ def _router_logits_kernel(x, w, part, T, E: tl.constexpr, H: tl.constexpr, TB: t @triton.jit def _router_topk_kernel(part, bias, w_out, id_out, T, scale, E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, - S: tl.constexpr, NORM: tl.constexpr, SOFTMAX: tl.constexpr, - PDL: tl.constexpr): - """One program per token: sum the S partials in a fixed order and score them, then - top-K (ties to the lower expert) -> gather the scores -> normalize -> * scale. Sigmoid - scores are picked by score + bias, fp32 softmax scores (SOFTMAX) by themselves. The + S: tl.constexpr, NORM: tl.constexpr, PDL: tl.constexpr): + """One program per token: sum the S partials in a fixed order, then sigmoid -> + bias + -> top-K (ties to the lower expert) -> gather the unbiased scores -> normalize. The ``1e-20`` guards an all-zero top-K and rounds away for any sum above 2^-42.""" if PDL: gdc_wait() @@ -217,21 +214,15 @@ def _router_topk_kernel(part, bias, w_out, id_out, T, scale, valid = e < E logit = tl.zeros([EP], tl.float32) for s in tl.static_range(S): - logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0).to(tl.float32) - if SOFTMAX: - x = tl.where(valid, logit, float("-inf")) - num = tl.where(valid, tl.exp(x - tl.max(x, 0)), 0.0) - score = num / tl.sum(num, 0) - rank = tl.where(valid, score, float("-inf")) - else: - score = tl.sigmoid(logit) - rank = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) + logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0) + score = tl.sigmoid(logit) + biased = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) kk = tl.arange(0, KP) ids = tl.zeros([KP], tl.int32) for j in tl.static_range(K): - idx = tl.argmax(rank, 0) + idx = tl.argmax(biased, 0) ids = tl.where(kk == j, idx, ids) - rank = tl.where(e == idx, float("-inf"), rank) + biased = tl.where(e == idx, float("-inf"), biased) ws = tl.sum(tl.where(e[None, :] == ids[:, None], score[None, :], 0.0), 1) ws = tl.where(kk < K, ws, 0.0) if NORM: @@ -498,8 +489,8 @@ def _launch_router(x, w, bias, part, topk_w, topk_ids, scale, normalize, pdl, tu num_warps=tuning.router["warps"]) _router_topk_kernel[(T,)]( part, bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), - K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, SOFTMAX=False, PDL=pdl, - num_warps=1, launch_pdl=pdl) + K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, PDL=pdl, num_warps=1, + launch_pdl=pdl) def _launch_gate_up(x, ids, w13, s13, sw13, act, block_size, limit, pdl, bj, bk, bjs, bks, warps): @@ -655,28 +646,26 @@ def route(x, w, bias, *, top_k, scale, normalize): @torch.library.custom_op("mstar::moe_router_topk", mutates_args=()) -def router_topk(logits: torch.Tensor, bias: torch.Tensor | None, top_k: int, scale: float, - normalize: bool, softmax: bool = False) -> tuple[torch.Tensor, torch.Tensor]: - """``route``'s top-k program on precomputed ``logits (T, E)`` of any float dtype, read in - fp32. ``softmax``: rank and weight by the fp32 softmax, with no ``bias``; otherwise by - the sigmoid router. An op opaque to dynamo, so a compiled gate keeps it whole.""" +def router_topk(logits: torch.Tensor, bias: torch.Tensor, top_k: int, scale: float, + normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: + """``route``'s top-k program on precomputed fp32 ``logits (T, E)``. An op opaque to + dynamo, so a compiled gate keeps it whole.""" T, E = logits.shape - assert logits.is_contiguous() and (bias is None) == softmax + assert logits.dtype == torch.float32 and logits.is_contiguous() topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) if T: _router_topk_kernel[(T,)]( - logits, logits if bias is None else bias, topk_w, topk_ids, T, float(scale), E=E, + logits, bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), K=top_k, KP=triton.next_power_of_2(top_k), S=1, - NORM=normalize, SOFTMAX=softmax, PDL=False, num_warps=1) + NORM=normalize, PDL=False, num_warps=1) return topk_w, topk_ids @router_topk.register_fake -def _(logits, bias, top_k, scale, normalize, softmax=False): +def _(logits, bias, top_k, scale, normalize): T = logits.shape[0] - return (logits.new_empty(T, top_k, dtype=torch.float32), - logits.new_empty(T, top_k, dtype=torch.int64)) + return logits.new_empty(T, top_k), logits.new_empty(T, top_k, dtype=torch.int64) def experts(x, w13, s13, w2, s2, shared_w13, shared_w2, topk_w, topk_ids, *, diff --git a/mstar/utils/fused_moe/router.py b/mstar/utils/fused_moe/router.py new file mode 100644 index 0000000000..b866614800 --- /dev/null +++ b/mstar/utils/fused_moe/router.py @@ -0,0 +1,64 @@ +"""Softmax top-k routing in one kernel. + +``softmax -> topk -> renormalize`` as separate ops is five launches per MoE layer +(``torch.topk`` alone is two, a radix select and a sort), which at decode batch sizes +costs more than the router GEMM. One program per token does it in registers. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _softmax_topk_kernel(logits, w_out, id_out, stride_t, + E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, + NORM: tl.constexpr): + """fp32 softmax over the E logits, the K largest probabilities in descending order, + optionally divided by their sum. Equal probabilities (common: the logits are bf16) go + to the lower expert first; ``torch.topk`` orders such ties arbitrarily, so the two can + pick different experts at a tie, both equally valid.""" + t = tl.program_id(0) + e = tl.arange(0, EP) + valid = e < E + x = tl.load(logits + t * stride_t + e, mask=valid, other=float("-inf")).to(tl.float32) + x = x - tl.max(x, 0) + num = tl.where(valid, tl.exp(x), 0.0) + prob = num / tl.sum(num, 0) + cand = tl.where(valid, prob, -1.0) + kk = tl.arange(0, KP) + ids = tl.zeros([KP], tl.int32) + ws = tl.zeros([KP], tl.float32) + for j in tl.static_range(K): + idx = tl.argmax(cand, 0) + ws = tl.where(kk == j, tl.max(cand, 0), ws) + ids = tl.where(kk == j, idx, ids) + cand = tl.where(e == idx, -1.0, cand) + if NORM: + ws = ws / tl.sum(tl.where(kk < K, ws, 0.0), 0) + tl.store(w_out + t * K + kk, ws, mask=kk < K) + tl.store(id_out + t * K + kk, ids.to(tl.int64), mask=kk < K) + + +@torch.library.custom_op("mstar::moe_softmax_topk", mutates_args=()) +def softmax_topk(logits: torch.Tensor, top_k: int, normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: + """Top-k of ``softmax(logits (T, E))`` as ``(fp32 weights, int64 expert ids)``, the + weights renormalized to sum to one when ``normalize``. Opaque to dynamo, so a compiled + router keeps it whole.""" + T, E = logits.shape + assert logits.stride(1) == 1 + topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) + topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) + if T: + _softmax_topk_kernel[(T,)]( + logits, topk_w, topk_ids, logits.stride(0), E=E, EP=triton.next_power_of_2(E), + K=top_k, KP=triton.next_power_of_2(top_k), NORM=normalize, num_warps=1) + return topk_w, topk_ids + + +@softmax_topk.register_fake +def _(logits, top_k, normalize): + T = logits.shape[0] + return (logits.new_empty(T, top_k, dtype=torch.float32), + logits.new_empty(T, top_k, dtype=torch.int64)) diff --git a/test/modular/test_fused_moe_decode.py b/test/modular/test_fused_moe_decode.py index 67e1ace1da..26903521d0 100644 --- a/test/modular/test_fused_moe_decode.py +++ b/test/modular/test_fused_moe_decode.py @@ -144,21 +144,16 @@ def test_router_launch_follows_the_table(table, launches, monkeypatch): assert logits["BK"] == bk and logits["num_warps"] == tuning.router["warps"] -@pytest.mark.parametrize("softmax", [False, True]) -def test_router_topk_op(launches, softmax): - # the softmax router passes bf16 logits and no bias; the weights are fp32 either way - dtype, bias = (torch.bfloat16, None) if softmax else (torch.float32, torch.zeros(64)) - w, ids = decode.router_topk(torch.randn(5, 64, dtype=dtype), bias, 8, 2.5, True, softmax) +def test_router_topk_op(launches): + w, ids = decode.router_topk(torch.randn(5, 64), torch.zeros(64), 8, 2.5, True) assert w.shape == ids.shape == (5, 8) and ids.dtype == torch.int64 (_, grid, _, kw), = launches - assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True and kw["SOFTMAX"] is softmax + assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): - fake_bias = None if softmax else torch.empty(64) - w, ids = decode.router_topk(torch.empty(5, 64, dtype=dtype), fake_bias, 8, 2.5, True, - softmax) - assert w.shape == (5, 8) and w.dtype == torch.float32 and ids.dtype == torch.int64 + w, ids = decode.router_topk(torch.empty(5, 64), torch.empty(64), 8, 2.5, True) + assert w.shape == (5, 8) and ids.dtype == torch.int64 # ---------------------------------------------------------------- CUDA diff --git a/test/modular/test_moe_router.py b/test/modular/test_moe_router.py index 88eb59cc54..c914e7c8e8 100644 --- a/test/modular/test_moe_router.py +++ b/test/modular/test_moe_router.py @@ -1,4 +1,4 @@ -"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_router_topk``) +"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_softmax_topk``) against the op sequence it replaces, and where the swept MoE tiles apply.""" from __future__ import annotations @@ -12,7 +12,7 @@ @pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) @pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): - from mstar.utils.fused_moe.decode import router_topk + from mstar.utils.fused_moe.router import softmax_topk torch.manual_seed(1) logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 @@ -21,7 +21,7 @@ def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): if norm: ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) - w, ids = router_topk(logits, None, top_k, 1.0, norm, softmax=True) + w, ids = softmax_topk(logits, top_k, norm) assert ids.dtype == torch.int64 and w.dtype == torch.float32 # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the # selected probabilities, and that every pick is a distinct expert holding them diff --git a/test/modular/test_qwen3_omni_fused_kernels.py b/test/modular/test_qwen3_omni_fused_kernels.py index ce37695554..522d25a9c4 100644 --- a/test/modular/test_qwen3_omni_fused_kernels.py +++ b/test/modular/test_qwen3_omni_fused_kernels.py @@ -3,6 +3,7 @@ * ``cp_qk_norm_rope_kv_write``: the code predictor's q/k norm + RoPE and its dense-KV write, bit-for-bit what two ``fused_qk_norm_rope`` calls and two strided copies produce. +* ``softmax_topk``: ``TopKRouter``'s softmax -> topk -> renormalize. """ from __future__ import annotations @@ -43,6 +44,46 @@ def test_cp_prologue_matches_the_unfused_ops(bs, seq_len, cache_pos): torch.testing.assert_close(cache, ref_cache, rtol=0, atol=0) +@pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) +@pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) +def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): + from mstar.utils.fused_moe.router import softmax_topk + + torch.manual_seed(1) + logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 + probs = torch.softmax(logits, dim=-1, dtype=torch.float) + ref_w, ref_ids = torch.topk(probs, top_k, dim=-1) + if norm: + ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) + + w, ids = softmax_topk(logits, top_k, norm) + assert ids.dtype == torch.int64 and w.dtype == torch.float32 + # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the + # selected probabilities, and that every pick is a distinct expert holding them + torch.testing.assert_close(w, ref_w, rtol=1e-6, atol=1e-7) + assert (ids.sort(dim=-1).values.diff(dim=-1) > 0).all() + picked = probs.gather(1, ids) + if norm: + picked = picked / picked.sum(dim=-1, keepdim=True) + torch.testing.assert_close(picked, w, rtol=1e-6, atol=1e-7) + # with no ties the choice and its order are torch's + distinct = (ref_w[:, :-1] != ref_w[:, 1:]).all(dim=-1) + assert torch.equal(ids[distinct], ref_ids[distinct]) + + +def test_router_compiles_with_the_fused_topk(): + from mstar.model.components.moe import TopKRouter + + router = TopKRouter(64, 128, 8, fused_softmax_topk=True).cuda().to(torch.bfloat16).requires_grad_(False) + torch.nn.init.normal_(router.weight) + x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + eager = router(x) + compiled = torch.compile(router, fullgraph=True)(x) + assert torch.equal(eager[1], compiled[1]) + torch.testing.assert_close(eager[0], compiled[0], rtol=0, atol=0) + + @pytest.mark.parametrize("bs", [1, 8, 32]) def test_cp_embed_accumulate_matches_the_unfused_ops(bs): from mstar.model.qwen3_omni.components.cp_kernels import cp_embed_accumulate From 772749aff0660ac5168fb21478ab438c6b98e3ec Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:11:27 -0700 Subject: [PATCH 48/59] Revert "warmup: let a spec's concurrency be a list of batch sizes" This reverts commit eb70ae905a77c9d08a875df5bfaaf312a9b733d2. --- configs/cosmos3_droid.yaml | 36 +++++++++++++++++- configs/cosmos3_edge_droid.yaml | 32 +++++++++++++++- docs/serving.rst | 3 +- mstar/api_server/entrypoint.py | 57 ++++++++++++++--------------- test/modular/test_startup_warmup.py | 9 +---- 5 files changed, 95 insertions(+), 42 deletions(-) diff --git a/configs/cosmos3_droid.yaml b/configs/cosmos3_droid.yaml index 0df754f65c..35a5e9b5a6 100644 --- a/configs/cosmos3_droid.yaml +++ b/configs/cosmos3_droid.yaml @@ -38,4 +38,38 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: [1, 2, 3, 4, 5, 6, 7, 8] + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 2 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 3 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 4 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 5 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 6 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 7 + - text: warmup + files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 8 diff --git a/configs/cosmos3_edge_droid.yaml b/configs/cosmos3_edge_droid.yaml index 16fcf44df2..2632aa69e2 100644 --- a/configs/cosmos3_edge_droid.yaml +++ b/configs/cosmos3_edge_droid.yaml @@ -45,4 +45,34 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: [2, 3, 4, 5, 6, 7, 8] + concurrency: 2 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 3 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 4 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 5 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 6 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 7 + - text: warmup + files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] + output_modalities: [action] + model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} + concurrency: 8 diff --git a/docs/serving.rst b/docs/serving.rst index 1cc84e5c5b..38309d2cb7 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -262,8 +262,7 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key comes from the extension, or give ``{path: ..., modality: ...}``. ``repeat: N`` sends ``text`` N times over, a long prompt for a path only long prompts take (a long-context prefill). ``concurrency: N`` submits N - copies at once, so they batch, for a step that compiles per batch size; a - list (``concurrency: [1, 2, 4]``) runs the entry once per size, in order. A + copies at once, so they batch, for a step that compiles per batch size. A malformed or failing warmup is logged and skipped. * - ``dist_timeout_s`` - *(optional)* Timeout in seconds for the NCCL world group and its parallel diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index cd6dc9063d..8ed9cfda0b 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -1348,14 +1348,13 @@ class WarmupRequest: """One entry of the deployment's ``warmup_requests``: a ``/generate`` request in yaml form. ``files`` are the media inputs in prompt order (any modality, any number), followed by ``text`` as on ``/generate``, - ``repeat`` times over (a long prompt for a long-context path). - ``concurrency`` holds the batch sizes it runs at, one warmup each, in order.""" + ``repeat`` times over (a long prompt for a long-context path).""" text: str | None = None files: list[WarmupMedia] = field(default_factory=list) output_modalities: list[str] = field(default_factory=lambda: ["text"]) model_kwargs: dict | None = None - concurrency: tuple[int, ...] = (1,) + concurrency: int = 1 @classmethod def from_config(cls, spec: Any) -> "WarmupRequest": @@ -1373,9 +1372,8 @@ def from_config(cls, spec: Any) -> "WarmupRequest": if repeat > 1 and not text: raise ValueError("repeat needs text") concurrency = spec.get("concurrency", 1) - levels = concurrency if isinstance(concurrency, list) else [concurrency] - if not levels or any(isinstance(n, bool) or not isinstance(n, int) or n < 1 for n in levels): - raise ValueError(f"concurrency must be a positive integer or a list of them, got {concurrency!r}") + if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: + raise ValueError(f"concurrency must be a positive integer, got {concurrency!r}") out_mods = spec.get("output_modalities") or ["text"] if isinstance(out_mods, str): out_mods = [m.strip() for m in out_mods.split(",") if m.strip()] @@ -1390,7 +1388,7 @@ def from_config(cls, spec: Any) -> "WarmupRequest": files=[WarmupMedia.from_config(f) for f in files], output_modalities=list(out_mods), model_kwargs=model_kwargs, - concurrency=tuple(levels), + concurrency=concurrency, ) def submit(self, server: APIServer, request_id: str) -> str: @@ -1434,11 +1432,11 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non warmup that is malformed or fails is logged and skipped; it must not keep the server from coming up. - Each spec runs as ``copies`` x its ``concurrency`` concurrent requests, - once per value when ``concurrency`` is a list: ``copies`` warms every - replica of a data-parallel deployment, since the conductor sends - concurrent requests to different replicas (the least-loaded one each), and - ``concurrency`` lets them batch, for a step that compiles per batch size. + Each spec runs as ``copies`` x its ``concurrency`` concurrent requests: + ``copies`` warms every replica of a data-parallel deployment, since the + conductor sends concurrent requests to different replicas (the + least-loaded one each), and ``concurrency`` lets them batch, for a step + that compiles per batch size. """ if not isinstance(specs, list): logger.warning("warmup_requests must be a list; ignoring %r", type(specs).__name__) @@ -1447,8 +1445,9 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non async def _one(request: WarmupRequest, request_id: str) -> None: await server.collect_results(request.submit(server, request_id)) - async def _all(request: WarmupRequest, base_id: str, n: int) -> None: - ids = [base_id] + [f"{base_id}-{c}" for c in range(1, n)] + async def _all(request: WarmupRequest, index: int) -> None: + n = copies * request.concurrency + ids = [f"warmup-{index}"] + [f"warmup-{index}-{c}" for c in range(1, n)] # every copy runs to its end before a failure is reported, so none is # cancelled mid-request and still in the server when the next spec starts results = await asyncio.gather(*(_one(request, rid) for rid in ids), return_exceptions=True) @@ -1462,23 +1461,21 @@ async def _all(request: WarmupRequest, base_id: str, n: int) -> None: except Exception as exc: # noqa: BLE001 — a repeat too large to build, say logger.warning("warmup_requests[%d] skipped: %s", index, exc) continue - for level, concurrency in enumerate(request.concurrency): - t0 = time.perf_counter() - try: - base_id = f"warmup-{index}" + (f".{level}" if level else "") - asyncio.run(_all(request, base_id, copies * concurrency)) - except Exception as exc: # noqa: BLE001 — a cold server is better than none - logger.warning( - "warmup request %d/%d failed after %.1f s at concurrency %d: %s", - index + 1, len(specs), time.perf_counter() - t0, concurrency, - getattr(exc, "detail", exc), - ) - continue - logger.info( - "warmup request %d/%d (%s) done in %.1f s at concurrency %d", - index + 1, len(specs), ",".join(request.output_modalities), - time.perf_counter() - t0, concurrency, + t0 = time.perf_counter() + try: + asyncio.run(_all(request, index)) + except Exception as exc: # noqa: BLE001 — a cold server is better than none + logger.warning( + "warmup request %d/%d failed after %.1f s: %s", + index + 1, len(specs), time.perf_counter() - t0, + getattr(exc, "detail", exc), ) + continue + logger.info( + "warmup request %d/%d (%s) done in %.1f s", + index + 1, len(specs), ",".join(request.output_modalities), + time.perf_counter() - t0, + ) # ------------------------------------------------------------------ diff --git a/test/modular/test_startup_warmup.py b/test/modular/test_startup_warmup.py index b2e6988b23..b948b84db0 100644 --- a/test/modular/test_startup_warmup.py +++ b/test/modular/test_startup_warmup.py @@ -171,15 +171,8 @@ def test_concurrency_runs_copies_together_for_batched_shapes() -> None: fake = _FakeServer() entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 3}], copies=2) assert [s["request_id"] for s in fake.submitted] == [f"warmup-0{s}" for s in ("", "-1", "-2", "-3", "-4", "-5")] - # a list runs one warmup per batch size, in order, each under its own ids fake = _FakeServer() - entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": [1, 2]}]) - assert [s["request_id"] for s in fake.submitted] == ["warmup-0", "warmup-0.1", "warmup-0.1-1"] - fake = _FakeServer() - entrypoint._run_warmup_requests(fake, [ - {"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}, - {"text": "c", "concurrency": []}, {"text": "d", "concurrency": [2, 0]}, - ]) + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}]) assert fake.submitted == [] From 0d1370f72c0021dafcc06837174420cb1d0ade23 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:11:27 -0700 Subject: [PATCH 49/59] Revert "bagel: reject an unsupported output_format at intake" This reverts commit b143d72ccdd720378e8dac318bac1dd297b71409. --- docs/clients.rst | 3 +-- mstar/model/bagel/bagel_model.py | 4 +--- mstar/model/components/diffusion/image_io.py | 18 +++++------------- test/modular/test_bagel_postprocess.py | 9 +-------- 4 files changed, 8 insertions(+), 26 deletions(-) diff --git a/docs/clients.rst b/docs/clients.rst index 88ebe28240..cb81cc44bd 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -282,8 +282,7 @@ Waypoint emits live RGB frame chunks and is not routed through the encoded-video Per-model notes: - **BAGEL** — chat returns text only; use ``/v1/images/generations`` and - ``/v1/images/edits`` for image output, with ``output_format`` and ``output_compression`` - as for FLUX.2 klein below. + ``/v1/images/edits`` for image output. - **FLUX.2 klein / Z-Image** — ``/v1/images/generations`` (and ``/v1/images/edits`` for klein) only: ``size`` as ``WxH`` (multiples of 16); ``seed``; ``num_inference_steps``, ``output_format`` (``png`` default, ``jpeg``, ``webp``), ``output_compression`` (JPEG / diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index a90d8d9b4b..f945cfe0d8 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -95,7 +95,7 @@ Span, prompt_spans, ) -from mstar.model.components.diffusion.image_io import encode_image, image_output_format +from mstar.model.components.diffusion.image_io import encode_image from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -633,8 +633,6 @@ def process_prompt( Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ - if "image" in output_modalities: - image_output_format(kwargs) # fail a bad output_format now, not after the denoise result: NameToTensorList = {} if prompt is not None: diff --git a/mstar/model/components/diffusion/image_io.py b/mstar/model/components/diffusion/image_io.py index a9088373c0..b4f541d88b 100644 --- a/mstar/model/components/diffusion/image_io.py +++ b/mstar/model/components/diffusion/image_io.py @@ -146,18 +146,6 @@ def deflate(index: int, part: bytes) -> bytes: OUTPUT_FORMATS = ("png", "jpeg", "webp") -def image_output_format(request_kwargs: dict | None) -> str: - """The request's ``output_format`` (default ``png``; ``jpg`` means ``jpeg``), or ValueError - for one :func:`encode_image` cannot write. A model calls it at intake too, so a bad value - fails the request before generation rather than after.""" - fmt = str((request_kwargs or {}).get("output_format") or "png").lower() - if fmt == "jpg": - fmt = "jpeg" - if fmt not in OUTPUT_FORMATS: - raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") - return fmt - - def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> bytes: """Encode a uint8 ``[3, H, W]`` image for the client per the OpenAI images knobs. @@ -167,7 +155,11 @@ def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> byt Everything else decodes to the pixels the VAE produced (PNG losslessly). """ kwargs = request_kwargs or {} - fmt = image_output_format(kwargs) + fmt = str(kwargs.get("output_format") or "png").lower() + if fmt == "jpg": + fmt = "jpeg" + if fmt not in OUTPUT_FORMATS: + raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") if fmt == "png": return uint8_to_png(image, compress_level=int(kwargs.get("png_compress_level", 1))) from PIL import Image diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py index e731ad4cad..9bbea4ff4f 100644 --- a/test/modular/test_bagel_postprocess.py +++ b/test/modular/test_bagel_postprocess.py @@ -1,5 +1,4 @@ -"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path -sent, and a format it cannot write fails at intake.""" +"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path sent.""" import io import numpy as np @@ -20,9 +19,3 @@ def test_image_png_decodes_to_the_truncated_pixels(): # the old path: permute to HWC, scale, truncate to uint8 expected = (output[0].permute(1, 2, 0) * 255).to(torch.uint8).numpy() assert np.array_equal(decoded, expected) - - -def test_unsupported_output_format_fails_at_intake(): - model = BagelModel.__new__(BagelModel) - with pytest.raises(ValueError, match="output_format"): - model.process_prompt(None, [], ["image"], output_format="gif") From 899e14e3ac9763a291e31ed596092b8e8a7ef35a Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 50/59] Revert "Revert "bagel: reject an unsupported output_format at intake"" This reverts commit 0d1370f72c0021dafcc06837174420cb1d0ade23. --- docs/clients.rst | 3 ++- mstar/model/bagel/bagel_model.py | 4 +++- mstar/model/components/diffusion/image_io.py | 18 +++++++++++++----- test/modular/test_bagel_postprocess.py | 9 ++++++++- 4 files changed, 26 insertions(+), 8 deletions(-) diff --git a/docs/clients.rst b/docs/clients.rst index cb81cc44bd..88ebe28240 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -282,7 +282,8 @@ Waypoint emits live RGB frame chunks and is not routed through the encoded-video Per-model notes: - **BAGEL** — chat returns text only; use ``/v1/images/generations`` and - ``/v1/images/edits`` for image output. + ``/v1/images/edits`` for image output, with ``output_format`` and ``output_compression`` + as for FLUX.2 klein below. - **FLUX.2 klein / Z-Image** — ``/v1/images/generations`` (and ``/v1/images/edits`` for klein) only: ``size`` as ``WxH`` (multiples of 16); ``seed``; ``num_inference_steps``, ``output_format`` (``png`` default, ``jpeg``, ``webp``), ``output_compression`` (JPEG / diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index f945cfe0d8..a90d8d9b4b 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -95,7 +95,7 @@ Span, prompt_spans, ) -from mstar.model.components.diffusion.image_io import encode_image +from mstar.model.components.diffusion.image_io import encode_image, image_output_format from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -633,6 +633,8 @@ def process_prompt( Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ + if "image" in output_modalities: + image_output_format(kwargs) # fail a bad output_format now, not after the denoise result: NameToTensorList = {} if prompt is not None: diff --git a/mstar/model/components/diffusion/image_io.py b/mstar/model/components/diffusion/image_io.py index b4f541d88b..a9088373c0 100644 --- a/mstar/model/components/diffusion/image_io.py +++ b/mstar/model/components/diffusion/image_io.py @@ -146,6 +146,18 @@ def deflate(index: int, part: bytes) -> bytes: OUTPUT_FORMATS = ("png", "jpeg", "webp") +def image_output_format(request_kwargs: dict | None) -> str: + """The request's ``output_format`` (default ``png``; ``jpg`` means ``jpeg``), or ValueError + for one :func:`encode_image` cannot write. A model calls it at intake too, so a bad value + fails the request before generation rather than after.""" + fmt = str((request_kwargs or {}).get("output_format") or "png").lower() + if fmt == "jpg": + fmt = "jpeg" + if fmt not in OUTPUT_FORMATS: + raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") + return fmt + + def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> bytes: """Encode a uint8 ``[3, H, W]`` image for the client per the OpenAI images knobs. @@ -155,11 +167,7 @@ def encode_image(image: torch.Tensor, request_kwargs: dict | None = None) -> byt Everything else decodes to the pixels the VAE produced (PNG losslessly). """ kwargs = request_kwargs or {} - fmt = str(kwargs.get("output_format") or "png").lower() - if fmt == "jpg": - fmt = "jpeg" - if fmt not in OUTPUT_FORMATS: - raise ValueError(f"output_format must be one of {OUTPUT_FORMATS}, got {fmt!r}") + fmt = image_output_format(kwargs) if fmt == "png": return uint8_to_png(image, compress_level=int(kwargs.get("png_compress_level", 1))) from PIL import Image diff --git a/test/modular/test_bagel_postprocess.py b/test/modular/test_bagel_postprocess.py index 9bbea4ff4f..e731ad4cad 100644 --- a/test/modular/test_bagel_postprocess.py +++ b/test/modular/test_bagel_postprocess.py @@ -1,4 +1,5 @@ -"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path sent.""" +"""BAGEL's image response: the shared PNG writer must decode to the pixels the old PIL path +sent, and a format it cannot write fails at intake.""" import io import numpy as np @@ -19,3 +20,9 @@ def test_image_png_decodes_to_the_truncated_pixels(): # the old path: permute to HWC, scale, truncate to uint8 expected = (output[0].permute(1, 2, 0) * 255).to(torch.uint8).numpy() assert np.array_equal(decoded, expected) + + +def test_unsupported_output_format_fails_at_intake(): + model = BagelModel.__new__(BagelModel) + with pytest.raises(ValueError, match="output_format"): + model.process_prompt(None, [], ["image"], output_format="gif") From f661b8d1f7c80f9027575b9d45c6e3c57c67364b Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 51/59] Revert "Revert "warmup: let a spec's concurrency be a list of batch sizes"" This reverts commit 772749aff0660ac5168fb21478ab438c6b98e3ec. --- configs/cosmos3_droid.yaml | 36 +----------------- configs/cosmos3_edge_droid.yaml | 32 +--------------- docs/serving.rst | 3 +- mstar/api_server/entrypoint.py | 57 +++++++++++++++-------------- test/modular/test_startup_warmup.py | 9 ++++- 5 files changed, 42 insertions(+), 95 deletions(-) diff --git a/configs/cosmos3_droid.yaml b/configs/cosmos3_droid.yaml index 35a5e9b5a6..0df754f65c 100644 --- a/configs/cosmos3_droid.yaml +++ b/configs/cosmos3_droid.yaml @@ -38,38 +38,4 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 2 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 3 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 4 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 5 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 6 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 7 - - text: warmup - files: [hf://nvidia/Cosmos3-Nano/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 8 + concurrency: [1, 2, 3, 4, 5, 6, 7, 8] diff --git a/configs/cosmos3_edge_droid.yaml b/configs/cosmos3_edge_droid.yaml index 2632aa69e2..16fcf44df2 100644 --- a/configs/cosmos3_edge_droid.yaml +++ b/configs/cosmos3_edge_droid.yaml @@ -45,34 +45,4 @@ warmup_requests: files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] output_modalities: [action] model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 2 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 3 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 4 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 5 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 6 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 7 - - text: warmup - files: [hf://nvidia/Cosmos3-Edge/assets/example_i2v_input.jpg] - output_modalities: [action] - model_kwargs: {action_mode: policy, domain_name: droid_lerobot, raw_action_dim: 10, action_chunk_size: 32} - concurrency: 8 + concurrency: [2, 3, 4, 5, 6, 7, 8] diff --git a/docs/serving.rst b/docs/serving.rst index 38309d2cb7..1cc84e5c5b 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -262,7 +262,8 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key comes from the extension, or give ``{path: ..., modality: ...}``. ``repeat: N`` sends ``text`` N times over, a long prompt for a path only long prompts take (a long-context prefill). ``concurrency: N`` submits N - copies at once, so they batch, for a step that compiles per batch size. A + copies at once, so they batch, for a step that compiles per batch size; a + list (``concurrency: [1, 2, 4]``) runs the entry once per size, in order. A malformed or failing warmup is logged and skipped. * - ``dist_timeout_s`` - *(optional)* Timeout in seconds for the NCCL world group and its parallel diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 8ed9cfda0b..cd6dc9063d 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -1348,13 +1348,14 @@ class WarmupRequest: """One entry of the deployment's ``warmup_requests``: a ``/generate`` request in yaml form. ``files`` are the media inputs in prompt order (any modality, any number), followed by ``text`` as on ``/generate``, - ``repeat`` times over (a long prompt for a long-context path).""" + ``repeat`` times over (a long prompt for a long-context path). + ``concurrency`` holds the batch sizes it runs at, one warmup each, in order.""" text: str | None = None files: list[WarmupMedia] = field(default_factory=list) output_modalities: list[str] = field(default_factory=lambda: ["text"]) model_kwargs: dict | None = None - concurrency: int = 1 + concurrency: tuple[int, ...] = (1,) @classmethod def from_config(cls, spec: Any) -> "WarmupRequest": @@ -1372,8 +1373,9 @@ def from_config(cls, spec: Any) -> "WarmupRequest": if repeat > 1 and not text: raise ValueError("repeat needs text") concurrency = spec.get("concurrency", 1) - if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: - raise ValueError(f"concurrency must be a positive integer, got {concurrency!r}") + levels = concurrency if isinstance(concurrency, list) else [concurrency] + if not levels or any(isinstance(n, bool) or not isinstance(n, int) or n < 1 for n in levels): + raise ValueError(f"concurrency must be a positive integer or a list of them, got {concurrency!r}") out_mods = spec.get("output_modalities") or ["text"] if isinstance(out_mods, str): out_mods = [m.strip() for m in out_mods.split(",") if m.strip()] @@ -1388,7 +1390,7 @@ def from_config(cls, spec: Any) -> "WarmupRequest": files=[WarmupMedia.from_config(f) for f in files], output_modalities=list(out_mods), model_kwargs=model_kwargs, - concurrency=concurrency, + concurrency=tuple(levels), ) def submit(self, server: APIServer, request_id: str) -> str: @@ -1432,11 +1434,11 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non warmup that is malformed or fails is logged and skipped; it must not keep the server from coming up. - Each spec runs as ``copies`` x its ``concurrency`` concurrent requests: - ``copies`` warms every replica of a data-parallel deployment, since the - conductor sends concurrent requests to different replicas (the - least-loaded one each), and ``concurrency`` lets them batch, for a step - that compiles per batch size. + Each spec runs as ``copies`` x its ``concurrency`` concurrent requests, + once per value when ``concurrency`` is a list: ``copies`` warms every + replica of a data-parallel deployment, since the conductor sends + concurrent requests to different replicas (the least-loaded one each), and + ``concurrency`` lets them batch, for a step that compiles per batch size. """ if not isinstance(specs, list): logger.warning("warmup_requests must be a list; ignoring %r", type(specs).__name__) @@ -1445,9 +1447,8 @@ def _run_warmup_requests(server: APIServer, specs: list, copies: int = 1) -> Non async def _one(request: WarmupRequest, request_id: str) -> None: await server.collect_results(request.submit(server, request_id)) - async def _all(request: WarmupRequest, index: int) -> None: - n = copies * request.concurrency - ids = [f"warmup-{index}"] + [f"warmup-{index}-{c}" for c in range(1, n)] + async def _all(request: WarmupRequest, base_id: str, n: int) -> None: + ids = [base_id] + [f"{base_id}-{c}" for c in range(1, n)] # every copy runs to its end before a failure is reported, so none is # cancelled mid-request and still in the server when the next spec starts results = await asyncio.gather(*(_one(request, rid) for rid in ids), return_exceptions=True) @@ -1461,21 +1462,23 @@ async def _all(request: WarmupRequest, index: int) -> None: except Exception as exc: # noqa: BLE001 — a repeat too large to build, say logger.warning("warmup_requests[%d] skipped: %s", index, exc) continue - t0 = time.perf_counter() - try: - asyncio.run(_all(request, index)) - except Exception as exc: # noqa: BLE001 — a cold server is better than none - logger.warning( - "warmup request %d/%d failed after %.1f s: %s", - index + 1, len(specs), time.perf_counter() - t0, - getattr(exc, "detail", exc), + for level, concurrency in enumerate(request.concurrency): + t0 = time.perf_counter() + try: + base_id = f"warmup-{index}" + (f".{level}" if level else "") + asyncio.run(_all(request, base_id, copies * concurrency)) + except Exception as exc: # noqa: BLE001 — a cold server is better than none + logger.warning( + "warmup request %d/%d failed after %.1f s at concurrency %d: %s", + index + 1, len(specs), time.perf_counter() - t0, concurrency, + getattr(exc, "detail", exc), + ) + continue + logger.info( + "warmup request %d/%d (%s) done in %.1f s at concurrency %d", + index + 1, len(specs), ",".join(request.output_modalities), + time.perf_counter() - t0, concurrency, ) - continue - logger.info( - "warmup request %d/%d (%s) done in %.1f s", - index + 1, len(specs), ",".join(request.output_modalities), - time.perf_counter() - t0, - ) # ------------------------------------------------------------------ diff --git a/test/modular/test_startup_warmup.py b/test/modular/test_startup_warmup.py index b948b84db0..b2e6988b23 100644 --- a/test/modular/test_startup_warmup.py +++ b/test/modular/test_startup_warmup.py @@ -171,8 +171,15 @@ def test_concurrency_runs_copies_together_for_batched_shapes() -> None: fake = _FakeServer() entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 3}], copies=2) assert [s["request_id"] for s in fake.submitted] == [f"warmup-0{s}" for s in ("", "-1", "-2", "-3", "-4", "-5")] + # a list runs one warmup per batch size, in order, each under its own ids fake = _FakeServer() - entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}]) + entrypoint._run_warmup_requests(fake, [{"text": "a", "concurrency": [1, 2]}]) + assert [s["request_id"] for s in fake.submitted] == ["warmup-0", "warmup-0.1", "warmup-0.1-1"] + fake = _FakeServer() + entrypoint._run_warmup_requests(fake, [ + {"text": "a", "concurrency": 0}, {"text": "b", "concurrency": True}, + {"text": "c", "concurrency": []}, {"text": "d", "concurrency": [2, 0]}, + ]) assert fake.submitted == [] From 4ccf8a86af473a1a6b89f9a62ba72c95c8845ae8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 52/59] Revert "Revert "fused_moe: one router top-k op for the softmax and sigmoid routers"" This reverts commit 4327e2238f2475a588030292abffa68989fecb87. --- mstar/model/components/moe.py | 8 +-- mstar/utils/fused_moe/decode.py | 49 ++++++++------ mstar/utils/fused_moe/router.py | 64 ------------------- test/modular/test_fused_moe_decode.py | 15 +++-- test/modular/test_moe_router.py | 6 +- test/modular/test_qwen3_omni_fused_kernels.py | 41 ------------ 6 files changed, 47 insertions(+), 136 deletions(-) delete mode 100644 mstar/utils/fused_moe/router.py diff --git a/mstar/model/components/moe.py b/mstar/model/components/moe.py index 17d864c63b..9a9fc24f4c 100644 --- a/mstar/model/components/moe.py +++ b/mstar/model/components/moe.py @@ -97,7 +97,7 @@ class TopKRouter(nn.Module): norm_topk_prob: if True, renormalize the top-k probabilities so they sum to 1. fused_softmax_topk: on CUDA, route through one Triton kernel - (``mstar::moe_softmax_topk``) instead of softmax, ``torch.topk`` + (``mstar::moe_router_topk``) instead of softmax, ``torch.topk`` and the renormalize. Equal probabilities go to the lower expert id, where ``torch.topk`` may pick another, so a model opts in once it has checked its outputs. @@ -138,10 +138,10 @@ def forward( hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) router_logits = F.linear(hidden_states, self.weight) if self.fused_softmax_topk and router_logits.is_cuda: - from mstar.utils.fused_moe.router import softmax_topk + from mstar.utils.fused_moe.decode import router_topk - routing_weights, router_indices = softmax_topk( - router_logits, self.top_k, self.norm_topk_prob, + routing_weights, router_indices = router_topk( + router_logits, None, self.top_k, 1.0, self.norm_topk_prob, softmax=True, ) return routing_weights, router_indices, None router_logits = F.softmax(router_logits, dtype=torch.float, dim=-1) diff --git a/mstar/utils/fused_moe/decode.py b/mstar/utils/fused_moe/decode.py index 16159ef3ac..eeafbc9ab4 100644 --- a/mstar/utils/fused_moe/decode.py +++ b/mstar/utils/fused_moe/decode.py @@ -4,6 +4,7 @@ Up to a per-GPU token count (``_Tuning.pair_max_tokens``) gate/up and down are GEMV kernels over (token, expert) pairs; above it experts repeat across the batch, and grouped kernels read each expert's weights once. ``swiglu_limit`` picks the clamped SwiGLU or the plain one. +The router's top-k program, as ``router_topk``, also serves softmax routers (``TopKRouter``). """ from __future__ import annotations @@ -202,9 +203,11 @@ def _router_logits_kernel(x, w, part, T, E: tl.constexpr, H: tl.constexpr, TB: t @triton.jit def _router_topk_kernel(part, bias, w_out, id_out, T, scale, E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, - S: tl.constexpr, NORM: tl.constexpr, PDL: tl.constexpr): - """One program per token: sum the S partials in a fixed order, then sigmoid -> + bias - -> top-K (ties to the lower expert) -> gather the unbiased scores -> normalize. The + S: tl.constexpr, NORM: tl.constexpr, SOFTMAX: tl.constexpr, + PDL: tl.constexpr): + """One program per token: sum the S partials in a fixed order and score them, then + top-K (ties to the lower expert) -> gather the scores -> normalize -> * scale. Sigmoid + scores are picked by score + bias, fp32 softmax scores (SOFTMAX) by themselves. The ``1e-20`` guards an all-zero top-K and rounds away for any sum above 2^-42.""" if PDL: gdc_wait() @@ -214,15 +217,21 @@ def _router_topk_kernel(part, bias, w_out, id_out, T, scale, valid = e < E logit = tl.zeros([EP], tl.float32) for s in tl.static_range(S): - logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0) - score = tl.sigmoid(logit) - biased = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) + logit += tl.load(part + (s * T + t) * E + e, mask=valid, other=0.0).to(tl.float32) + if SOFTMAX: + x = tl.where(valid, logit, float("-inf")) + num = tl.where(valid, tl.exp(x - tl.max(x, 0)), 0.0) + score = num / tl.sum(num, 0) + rank = tl.where(valid, score, float("-inf")) + else: + score = tl.sigmoid(logit) + rank = tl.where(valid, score + tl.load(bias + e, mask=valid, other=0.0), float("-inf")) kk = tl.arange(0, KP) ids = tl.zeros([KP], tl.int32) for j in tl.static_range(K): - idx = tl.argmax(biased, 0) + idx = tl.argmax(rank, 0) ids = tl.where(kk == j, idx, ids) - biased = tl.where(e == idx, float("-inf"), biased) + rank = tl.where(e == idx, float("-inf"), rank) ws = tl.sum(tl.where(e[None, :] == ids[:, None], score[None, :], 0.0), 1) ws = tl.where(kk < K, ws, 0.0) if NORM: @@ -489,8 +498,8 @@ def _launch_router(x, w, bias, part, topk_w, topk_ids, scale, normalize, pdl, tu num_warps=tuning.router["warps"]) _router_topk_kernel[(T,)]( part, bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), - K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, PDL=pdl, num_warps=1, - launch_pdl=pdl) + K=K, KP=triton.next_power_of_2(K), S=splits, NORM=normalize, SOFTMAX=False, PDL=pdl, + num_warps=1, launch_pdl=pdl) def _launch_gate_up(x, ids, w13, s13, sw13, act, block_size, limit, pdl, bj, bk, bjs, bks, warps): @@ -646,26 +655,28 @@ def route(x, w, bias, *, top_k, scale, normalize): @torch.library.custom_op("mstar::moe_router_topk", mutates_args=()) -def router_topk(logits: torch.Tensor, bias: torch.Tensor, top_k: int, scale: float, - normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: - """``route``'s top-k program on precomputed fp32 ``logits (T, E)``. An op opaque to - dynamo, so a compiled gate keeps it whole.""" +def router_topk(logits: torch.Tensor, bias: torch.Tensor | None, top_k: int, scale: float, + normalize: bool, softmax: bool = False) -> tuple[torch.Tensor, torch.Tensor]: + """``route``'s top-k program on precomputed ``logits (T, E)`` of any float dtype, read in + fp32. ``softmax``: rank and weight by the fp32 softmax, with no ``bias``; otherwise by + the sigmoid router. An op opaque to dynamo, so a compiled gate keeps it whole.""" T, E = logits.shape - assert logits.dtype == torch.float32 and logits.is_contiguous() + assert logits.is_contiguous() and (bias is None) == softmax topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) if T: _router_topk_kernel[(T,)]( - logits, bias, topk_w, topk_ids, T, float(scale), E=E, + logits, logits if bias is None else bias, topk_w, topk_ids, T, float(scale), E=E, EP=triton.next_power_of_2(E), K=top_k, KP=triton.next_power_of_2(top_k), S=1, - NORM=normalize, PDL=False, num_warps=1) + NORM=normalize, SOFTMAX=softmax, PDL=False, num_warps=1) return topk_w, topk_ids @router_topk.register_fake -def _(logits, bias, top_k, scale, normalize): +def _(logits, bias, top_k, scale, normalize, softmax=False): T = logits.shape[0] - return logits.new_empty(T, top_k), logits.new_empty(T, top_k, dtype=torch.int64) + return (logits.new_empty(T, top_k, dtype=torch.float32), + logits.new_empty(T, top_k, dtype=torch.int64)) def experts(x, w13, s13, w2, s2, shared_w13, shared_w2, topk_w, topk_ids, *, diff --git a/mstar/utils/fused_moe/router.py b/mstar/utils/fused_moe/router.py deleted file mode 100644 index b866614800..0000000000 --- a/mstar/utils/fused_moe/router.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Softmax top-k routing in one kernel. - -``softmax -> topk -> renormalize`` as separate ops is five launches per MoE layer -(``torch.topk`` alone is two, a radix select and a sort), which at decode batch sizes -costs more than the router GEMM. One program per token does it in registers. -""" -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _softmax_topk_kernel(logits, w_out, id_out, stride_t, - E: tl.constexpr, EP: tl.constexpr, K: tl.constexpr, KP: tl.constexpr, - NORM: tl.constexpr): - """fp32 softmax over the E logits, the K largest probabilities in descending order, - optionally divided by their sum. Equal probabilities (common: the logits are bf16) go - to the lower expert first; ``torch.topk`` orders such ties arbitrarily, so the two can - pick different experts at a tie, both equally valid.""" - t = tl.program_id(0) - e = tl.arange(0, EP) - valid = e < E - x = tl.load(logits + t * stride_t + e, mask=valid, other=float("-inf")).to(tl.float32) - x = x - tl.max(x, 0) - num = tl.where(valid, tl.exp(x), 0.0) - prob = num / tl.sum(num, 0) - cand = tl.where(valid, prob, -1.0) - kk = tl.arange(0, KP) - ids = tl.zeros([KP], tl.int32) - ws = tl.zeros([KP], tl.float32) - for j in tl.static_range(K): - idx = tl.argmax(cand, 0) - ws = tl.where(kk == j, tl.max(cand, 0), ws) - ids = tl.where(kk == j, idx, ids) - cand = tl.where(e == idx, -1.0, cand) - if NORM: - ws = ws / tl.sum(tl.where(kk < K, ws, 0.0), 0) - tl.store(w_out + t * K + kk, ws, mask=kk < K) - tl.store(id_out + t * K + kk, ids.to(tl.int64), mask=kk < K) - - -@torch.library.custom_op("mstar::moe_softmax_topk", mutates_args=()) -def softmax_topk(logits: torch.Tensor, top_k: int, normalize: bool) -> tuple[torch.Tensor, torch.Tensor]: - """Top-k of ``softmax(logits (T, E))`` as ``(fp32 weights, int64 expert ids)``, the - weights renormalized to sum to one when ``normalize``. Opaque to dynamo, so a compiled - router keeps it whole.""" - T, E = logits.shape - assert logits.stride(1) == 1 - topk_w = torch.empty(T, top_k, dtype=torch.float32, device=logits.device) - topk_ids = torch.empty(T, top_k, dtype=torch.int64, device=logits.device) - if T: - _softmax_topk_kernel[(T,)]( - logits, topk_w, topk_ids, logits.stride(0), E=E, EP=triton.next_power_of_2(E), - K=top_k, KP=triton.next_power_of_2(top_k), NORM=normalize, num_warps=1) - return topk_w, topk_ids - - -@softmax_topk.register_fake -def _(logits, top_k, normalize): - T = logits.shape[0] - return (logits.new_empty(T, top_k, dtype=torch.float32), - logits.new_empty(T, top_k, dtype=torch.int64)) diff --git a/test/modular/test_fused_moe_decode.py b/test/modular/test_fused_moe_decode.py index 26903521d0..67e1ace1da 100644 --- a/test/modular/test_fused_moe_decode.py +++ b/test/modular/test_fused_moe_decode.py @@ -144,16 +144,21 @@ def test_router_launch_follows_the_table(table, launches, monkeypatch): assert logits["BK"] == bk and logits["num_warps"] == tuning.router["warps"] -def test_router_topk_op(launches): - w, ids = decode.router_topk(torch.randn(5, 64), torch.zeros(64), 8, 2.5, True) +@pytest.mark.parametrize("softmax", [False, True]) +def test_router_topk_op(launches, softmax): + # the softmax router passes bf16 logits and no bias; the weights are fp32 either way + dtype, bias = (torch.bfloat16, None) if softmax else (torch.float32, torch.zeros(64)) + w, ids = decode.router_topk(torch.randn(5, 64, dtype=dtype), bias, 8, 2.5, True, softmax) assert w.shape == ids.shape == (5, 8) and ids.dtype == torch.int64 (_, grid, _, kw), = launches - assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True + assert grid == (5,) and kw["S"] == 1 and kw["NORM"] is True and kw["SOFTMAX"] is softmax from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): - w, ids = decode.router_topk(torch.empty(5, 64), torch.empty(64), 8, 2.5, True) - assert w.shape == (5, 8) and ids.dtype == torch.int64 + fake_bias = None if softmax else torch.empty(64) + w, ids = decode.router_topk(torch.empty(5, 64, dtype=dtype), fake_bias, 8, 2.5, True, + softmax) + assert w.shape == (5, 8) and w.dtype == torch.float32 and ids.dtype == torch.int64 # ---------------------------------------------------------------- CUDA diff --git a/test/modular/test_moe_router.py b/test/modular/test_moe_router.py index c914e7c8e8..88eb59cc54 100644 --- a/test/modular/test_moe_router.py +++ b/test/modular/test_moe_router.py @@ -1,4 +1,4 @@ -"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_softmax_topk``) +"""``TopKRouter``'s fused softmax -> top-k -> renormalize (``mstar::moe_router_topk``) against the op sequence it replaces, and where the swept MoE tiles apply.""" from __future__ import annotations @@ -12,7 +12,7 @@ @pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) @pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): - from mstar.utils.fused_moe.router import softmax_topk + from mstar.utils.fused_moe.decode import router_topk torch.manual_seed(1) logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 @@ -21,7 +21,7 @@ def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): if norm: ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) - w, ids = softmax_topk(logits, top_k, norm) + w, ids = router_topk(logits, None, top_k, 1.0, norm, softmax=True) assert ids.dtype == torch.int64 and w.dtype == torch.float32 # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the # selected probabilities, and that every pick is a distinct expert holding them diff --git a/test/modular/test_qwen3_omni_fused_kernels.py b/test/modular/test_qwen3_omni_fused_kernels.py index 522d25a9c4..ce37695554 100644 --- a/test/modular/test_qwen3_omni_fused_kernels.py +++ b/test/modular/test_qwen3_omni_fused_kernels.py @@ -3,7 +3,6 @@ * ``cp_qk_norm_rope_kv_write``: the code predictor's q/k norm + RoPE and its dense-KV write, bit-for-bit what two ``fused_qk_norm_rope`` calls and two strided copies produce. -* ``softmax_topk``: ``TopKRouter``'s softmax -> topk -> renormalize. """ from __future__ import annotations @@ -44,46 +43,6 @@ def test_cp_prologue_matches_the_unfused_ops(bs, seq_len, cache_pos): torch.testing.assert_close(cache, ref_cache, rtol=0, atol=0) -@pytest.mark.parametrize("tokens", [1, 7, 32, 300, 4096]) -@pytest.mark.parametrize("experts,top_k,norm", [(128, 8, True), (128, 6, True), (60, 4, False)]) -def test_softmax_topk_matches_the_torch_router(tokens, experts, top_k, norm): - from mstar.utils.fused_moe.router import softmax_topk - - torch.manual_seed(1) - logits = torch.randn(tokens, experts, device="cuda", dtype=torch.bfloat16) * 3 - probs = torch.softmax(logits, dim=-1, dtype=torch.float) - ref_w, ref_ids = torch.topk(probs, top_k, dim=-1) - if norm: - ref_w = ref_w / ref_w.sum(dim=-1, keepdim=True) - - w, ids = softmax_topk(logits, top_k, norm) - assert ids.dtype == torch.int64 and w.dtype == torch.float32 - # bf16 logits tie often and torch.topk orders ties arbitrarily: compare the - # selected probabilities, and that every pick is a distinct expert holding them - torch.testing.assert_close(w, ref_w, rtol=1e-6, atol=1e-7) - assert (ids.sort(dim=-1).values.diff(dim=-1) > 0).all() - picked = probs.gather(1, ids) - if norm: - picked = picked / picked.sum(dim=-1, keepdim=True) - torch.testing.assert_close(picked, w, rtol=1e-6, atol=1e-7) - # with no ties the choice and its order are torch's - distinct = (ref_w[:, :-1] != ref_w[:, 1:]).all(dim=-1) - assert torch.equal(ids[distinct], ref_ids[distinct]) - - -def test_router_compiles_with_the_fused_topk(): - from mstar.model.components.moe import TopKRouter - - router = TopKRouter(64, 128, 8, fused_softmax_topk=True).cuda().to(torch.bfloat16).requires_grad_(False) - torch.nn.init.normal_(router.weight) - x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) - with torch.inference_mode(): - eager = router(x) - compiled = torch.compile(router, fullgraph=True)(x) - assert torch.equal(eager[1], compiled[1]) - torch.testing.assert_close(eager[0], compiled[0], rtol=0, atol=0) - - @pytest.mark.parametrize("bs", [1, 8, 32]) def test_cp_embed_accumulate_matches_the_unfused_ops(bs): from mstar.model.qwen3_omni.components.cp_kernels import cp_embed_accumulate From 26932204aa4584a7475145eabd818ccc25f988c8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 53/59] Revert "Revert "Share the device loop-back decode input ids on ARNodeSubmodule"" This reverts commit e75ec283edf7dc84425aaa022a9e6f5bf5d3f0dc. --- mstar/model/higgs_audio/submodules.py | 12 +++--------- mstar/model/qwen3_5/submodules.py | 23 +++-------------------- mstar/model/submodule_base.py | 23 +++++++++++++++++++++++ 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 0ea8b5e33b..4a84991fd1 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -390,15 +390,9 @@ def preprocess( return {"input_ids": torch.cat([inp.input_ids for inp in inputs])} # device loop-back: each row's last sampled token, read by slot; # rows whose previous step was the prefill brought theirs - sampler: SamplerResource = engine_inputs.resources[SAMPLER] - ids = sampler.loopback_tokens(engine_inputs.request_ids) - routed = [i for i, inp in enumerate(inputs) if inp.input_ids is not None] - if routed: - ids.index_copy_( - 0, to_device_nonblocking(torch.tensor(routed), ids.device), - torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), - ) - return {"input_ids": ids} + return {"input_ids": self.loopback_input_ids( + engine_inputs.resources[SAMPLER], engine_inputs.request_ids, inputs, + )} return { "input_embeds": torch.cat( [inp.input_embeds for inp in inputs], dim=0 diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 10448bab06..c527e2a575 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -76,7 +76,6 @@ device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides -from mstar.utils.device import DeviceCapability from mstar.utils.h2d import to_device_nonblocking from mstar.utils.profiler import PHASE_PERIOD, phase_record @@ -438,33 +437,17 @@ def preprocess( # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] t0 = perf_counter() if PHASE_PERIOD else 0.0 - ids = sampler.loopback_tokens(engine_inputs.request_ids) - # rows whose last step ran elsewhere brought their token - routed = [ - i for i, inp in enumerate(inputs) - if inp.input_ids is not None and inp.input_seq_len - ] if graph_walk != "decode": # a mixed step packs its decode rows' single tokens beside the # prompts: each row takes its own ids, or its loop-back token + ids = sampler.loopback_tokens(engine_inputs.request_ids) ids = torch.cat([ inp.input_ids.to(ids.dtype) if inp.input_ids is not None else ids[i:i + 1] for i, inp in enumerate(inputs) ]) - elif routed: - # Through pinned memory, non-blocking: a pageable - # ``torch.tensor(..., device=cuda)`` is a synchronous copy - # that waits for the queued step, and every request's first - # decode step after its prefill takes this path. - index = torch.tensor( - routed, dtype=torch.long, - pin_memory=DeviceCapability.of(ids.device).pin_memory, - ).to(ids.device, non_blocking=True) - ids.index_copy_( - 0, index, - torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), - ) + else: + ids = self.loopback_input_ids(sampler, engine_inputs.request_ids, inputs) out["input_ids"] = ids if PHASE_PERIOD: phase_record("engine.preprocess.loopback", perf_counter() - t0) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 2cbc12330d..19a30efaea 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -13,6 +13,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep +from mstar.utils.h2d import to_device_nonblocking from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) @@ -31,6 +32,7 @@ class HostRows(NamedTuple): from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner from mstar.engine.engine import ExecutingBatch + from mstar.engine.resources.sampler.resource import SamplerResource from mstar.streaming.stream_buffer import StreamChunkInfo @@ -1058,3 +1060,24 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: # input name to tensor pass + + @staticmethod + def loopback_input_ids( + sampler: SamplerResource, request_ids: list[str], inputs: Sequence[ARNodeInputs], + ) -> torch.Tensor: + """A device loop-back decode step's input ids: each row's last sampled + token, read off the sampler on the device, except where a row brought + its own (its previous step ran elsewhere, e.g. the prefill).""" + ids = sampler.loopback_tokens(request_ids) + routed = [ + i for i, inp in enumerate(inputs) + if inp.input_ids is not None and inp.input_seq_len + ] + if routed: + # every request's first decode step after its prefill comes here, + # so the row index goes up without a blocking copy + ids.index_copy_( + 0, to_device_nonblocking(torch.tensor(routed, dtype=torch.long), ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + return ids From a25c4f0781563436dbbd3afd1f8d8131d0d69bf5 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 54/59] Revert "Revert "One token stop check for AR submodules, declared as TokenStop"" This reverts commit 38a7a01741393396f9ae558b1e25527d480aa43c. --- docs/adding_models.rst | 5 +- mstar/model/higgs_audio/submodules.py | 53 +---------- mstar/model/orpheus/submodules.py | 48 +--------- mstar/model/qwen3_5/submodules.py | 70 ++------------ mstar/model/qwen3_tts/submodules.py | 62 ++----------- mstar/model/submodule_base.py | 91 +++++++++++++++++++ test/modular/test_model_kv_layouts.py | 7 +- .../test_qwen3_5_check_stop_batched.py | 5 +- 8 files changed, 122 insertions(+), 219 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 3264c7e9cb..3780a6d44d 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -954,7 +954,10 @@ methods are: ``check_stop(...) -> set[str]`` (optional) Runs off the GPU thread, and may read tensor values. Return the names of the ``Loop`` objects to stop, for example after seeing the EOS token. This is how a decode loop - terminates. + terminates. An AR submodule whose loop ends on a stop token or the token budget sets + ``self.token_stop = TokenStop(...)`` in ``__init__`` instead of writing one; that also + gives it ``check_stop_batched``, which checks a whole step off one host copy of its + tokens. Qwen3.5's ``LLMSubmodule`` is an example. ``cleanup_request(self, request_id)`` (optional) Free per-request state held inside the submodule when a request finishes: buffers, diff --git a/mstar/model/higgs_audio/submodules.py b/mstar/model/higgs_audio/submodules.py index 4a84991fd1..2edec64990 100644 --- a/mstar/model/higgs_audio/submodules.py +++ b/mstar/model/higgs_audio/submodules.py @@ -37,10 +37,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, device_loopback_enabled, ) from mstar.utils.h2d import to_device_nonblocking @@ -269,6 +269,10 @@ def __init__(self, llm: nn.Module, config: HiggsAudioModelConfig): super().__init__() self.model = llm self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=frozenset(config.stop_token_ids), + sampler=SAMPLER, walk="decode", + ) def get_accelerator_graph_configs( self, device: torch.device, tp_world_size: int = 1, @@ -483,22 +487,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs or request_info.graph_walk != "decode": - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - if (not ignore_eos and token in self.config.stop_token_ids) or \ - decoded_tokens >= request_info.max_tokens: - return {"decode_loop"} - return set() - def step_is_input_free(self, graph_walk: str) -> bool: # a decode row is one token, like the capture template's return graph_walk == "decode" @@ -507,34 +495,3 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # the decode loop's client edge carries the token the stop check # already has on the host; prefill's client edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for the whole batch off one ``tolist``.""" - if step_walk is not None and step_walk != "decode": - return {} - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - stop_ids = self.config.stop_token_ids - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - if step_walk is None and info.graph_walk != "decode": - continue - if ( - values[i] in stop_ids - and not info.resource_configs[SAMPLER].ignore_eos - ) or info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens: - stops[rid] = {"decode_loop"} - return stops diff --git a/mstar/model/orpheus/submodules.py b/mstar/model/orpheus/submodules.py index 3194142765..75a10dfd3d 100644 --- a/mstar/model/orpheus/submodules.py +++ b/mstar/model/orpheus/submodules.py @@ -23,11 +23,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo @@ -53,6 +53,10 @@ def __init__( self.embed_tokens = language_model.model.embed_tokens self.lm_head = language_model.lm_head self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", + stop_ids=frozenset({config.stop_token_id}), sampler=SAMPLER, + ) # Orpheus only ever emits ```` ids (SNAC codes and the # speech/turn markers): the LM head computes the logits of rows # ``[custom_token_base_id, vocab)`` alone, about 18% of the head, and @@ -222,48 +226,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - eos_token_id = self.config.stop_token_id - ignore_eos= request_info.resource_configs[SAMPLER].ignore_eos - if (not ignore_eos and eos_token_id == token) or \ - (request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= request_info.max_tokens): - return {"decode_loop"} - return set() - - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every row off one host copy of the step's tokens.""" - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(-1).tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - eos = self.config.stop_token_id - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - if (values[i] == eos and not info.resource_configs[SAMPLER].ignore_eos) or ( - info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens - ): - stops[rid] = {"decode_loop"} - return stops - class SNACDecoderSubmodule(NodeSubmodule): """SNAC 24kHz streaming decoder submodule. diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index c527e2a575..eff201d40d 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -69,10 +69,10 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, device_loopback_enabled, ) from mstar.platforms.compile import with_inductor_overrides @@ -154,6 +154,11 @@ def __init__( self.model = model self.config = config self.vision_config = vision_config + # only a decode row's token ends decode_loop; a prefill is not in it yet + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=config.stop_token_ids, + sampler=SAMPLER, walk="decode", + ) self._vision_sentinels: tuple[torch.Tensor, torch.Tensor] | None = None # ------------------------------------------------------------------ @@ -652,24 +657,6 @@ def _token_is_real(request_info: CurrentForwardPassInfo) -> bool: "last_prefill", False, ) - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - # only a decode row's stop can end decode_loop; a prefill is not in it yet - if request_info.graph_walk != "decode" or "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - hit_eos = not ignore_eos and token in self.config.stop_token_ids - out_of_budget = ( - request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - >= request_info.max_tokens - ) - return {"decode_loop"} if hit_eos or out_of_budget else set() - def step_is_input_free(self, graph_walk: str) -> bool: # decode rows are one token each, like the template; the prefill walks # feed `prefill_tokens` from the inputs @@ -680,51 +667,6 @@ def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # check already has on the host; prefill's new_token edge is persisted return {"text_inputs": "new_token"} if graph_walk == "decode" else {} - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every request off one ``tolist``.""" - if step_walk is not None and step_walk != "decode": - return {} # as check_stop: only a decode row's stop ends decode_loop - tokens = host_rows.buffers.get("new_token") - if not torch.is_tensor(tokens) or tokens.dim() == 0: - return None - values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - rows = host_rows.request_ids - if len(rows) >= len(request_ids) and all( - a == b for a, b in zip(request_ids, rows, strict=False) - ): - # the common case: the batch in forward order, no row map needed - pairs = zip(request_ids, values, strict=False) - else: - row_of = {rid: i for i, rid in enumerate(rows)} - pairs = ( - (rid, values[i]) for rid in request_ids - if (i := row_of.get(rid)) is not None and i < len(values) - ) # no row, no output: the per-request path skips it too - stop_ids = self.config.stop_token_ids - stops: dict[str, set[str]] = {} - mixed = step_walk is None - for rid, value in pairs: - info = request_infos[rid] - # a mixed step reads each row's walk; a decode step need not - if mixed and info.graph_walk != "decode": - continue - # the set test first: a stop id is rare, the config read behind it - # goes through the wrapper - if ( - value in stop_ids - and not info.resource_configs[SAMPLER].ignore_eos - ) or ( - info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - >= info.max_tokens - ): - stops[rid] = {"decode_loop"} - return stops # The vision tower's block loop, as a piecewise capture region diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 35478346f1..c8f66ceb8e 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -85,11 +85,11 @@ ARNodeInputs, ARNodeSubmodule, BatchedModelOutput, - HostRows, InputMetadata, ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.streaming.stream_buffer import StreamChunkInfo from mstar.utils.h2d import PinnedStager @@ -162,6 +162,12 @@ def __init__( self.talker_config = config.talker self.cp_config = config.talker.code_predictor self.num_codes = config.talker.num_code_groups + # ``postprocess`` counts the frames, the prefill's included + self.token_stop = TokenStop( + loop="talker_decode_loop", buffer="layer0_codes", + stop_ids=frozenset({config.talker.codec_eos_token_id}), + sampler=TALKER_SAMPLER, generated_key="generated_frames", + ) self._suppress_mask: torch.Tensor | None = None self._decode_last_token_indices: torch.Tensor | None = None self._flag_stager = PinnedStager(torch.bool, numel=self.MAX_BATCH_SIZE) @@ -796,60 +802,6 @@ def postprocess( "generated_frames", int(state.get("generated_frames", 0)) + 1 ) - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - """Stop the graph loop on codec EOS or the request frame limit. - - This callback runs off the GPU execution thread, so reading the sampled - token with ``item`` is allowed here rather than in ``postprocess``. - """ - if "layer0_codes" not in outputs: - return set() - token = int(outputs["layer0_codes"][0].item()) - generated = int( - self.request_state(request_id).get("generated_frames", 0) - ) - max_tokens = request_info.max_tokens - sampling_config = request_info.resource_configs.get(TALKER_SAMPLER) - ignore_eos = bool(getattr(sampling_config, "ignore_eos", False)) - reached_eos = ( - not ignore_eos and token == self.talker_config.codec_eos_token_id - ) - if reached_eos or generated >= max_tokens: - logger.debug("talker %s: stop after %d frames (eos=%s)", request_id, generated, reached_eos) - return {"talker_decode_loop"} - return set() - - def check_stop_batched( - self, - request_ids: list[str], - request_infos: dict[str, CurrentForwardPassInfo], - host_rows: HostRows, - step_walk: str | None = None, - ) -> dict[str, set[str]] | None: - """``check_stop`` for every row off one host copy of the group-0 codes.""" - codes = host_rows.buffers.get("layer0_codes") - if not torch.is_tensor(codes) or codes.dim() == 0: - return None - values = codes.reshape(-1).tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} - eos = self.talker_config.codec_eos_token_id - stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too - info = request_infos[rid] - generated = int(self.request_state(rid).get("generated_frames", 0)) - ignore_eos = bool(getattr(info.resource_configs.get(TALKER_SAMPLER), "ignore_eos", False)) - if (values[i] == eos and not ignore_eos) or generated >= info.max_tokens: - stops[rid] = {"talker_decode_loop"} - return stops - def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Admit compatible prefill/decode requests to continuous batching. diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 19a30efaea..32bdb1f56a 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -28,6 +28,28 @@ class HostRows(NamedTuple): buffers: dict[str, torch.Tensor] +class TokenStop(NamedTuple): + """When an AR submodule's sampled token ends its decode loop: on a stop + id (unless the request's sampler config sets ``ignore_eos``), or once + the request has generated ``max_tokens`` tokens. + + Declared as ``ARNodeSubmodule.token_stop``; ``check_stop`` and + ``check_stop_batched`` follow from it. + """ + + loop: str + # the output (and ``check_stop_buffers`` entry) holding the step's token + buffer: str + stop_ids: frozenset[int] + # resource key of the sampler whose request config carries ``ignore_eos`` + sampler: str + # the only walk whose rows can stop the loop; None for every walk + walk: str | None = None + # request-state key counting the tokens generated so far; None counts + # the loop's iterations, plus the token this step sampled + generated_key: str | None = None + + if TYPE_CHECKING: from mstar.engine.accelerator_graph_config import AcceleratorGraphConfig, PiecewiseAcceleratorGraphConfig from mstar.engine.accelerator_graph_runner import PiecewiseAcceleratorGraphRunner @@ -997,6 +1019,75 @@ def cleanup_request(self, request_id: str): class ARNodeSubmodule(NodeSubmodule): + # Set (in ``__init__``, where the stop ids are known) by a submodule whose + # sampled token ends a dynamic loop; the stop checks below follow from it. + token_stop: TokenStop | None = None + + def check_stop( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + stop = self.token_stop + if stop is None: + return super().check_stop(request_id, request_info, outputs) + if (stop.walk is not None and request_info.graph_walk != stop.walk) or stop.buffer not in outputs: + return set() + token = outputs[stop.buffer][0].item() + return self._token_stops(((request_id, token),), {request_id: request_info}, False).get(request_id, set()) + + def check_stop_batched( + self, request_ids: list[str], + request_infos: dict[str, CurrentForwardPassInfo], + host_rows: HostRows, + step_walk: str | None = None, + ) -> dict[str, set[str]] | None: + stop = self.token_stop + if stop is None: + return super().check_stop_batched(request_ids, request_infos, host_rows, step_walk) + if stop.walk is not None and step_walk is not None and step_walk != stop.walk: + return {} + tokens = host_rows.buffers.get(stop.buffer) + if not torch.is_tensor(tokens) or tokens.dim() == 0: + return None + values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() + rows = host_rows.request_ids + if len(rows) >= len(request_ids) and all( + a == b for a, b in zip(request_ids, rows, strict=False) + ): + # the common case: the batch in forward order, no row map needed + pairs = zip(request_ids, values, strict=False) + else: + row_of = {rid: i for i, rid in enumerate(rows)} + pairs = ( + (rid, values[i]) for rid in request_ids + if (i := row_of.get(rid)) is not None and i < len(values) + ) # no row, no output: the per-request path skips it too + # a mixed step reads each row's walk; a step of one walk need not + return self._token_stops(pairs, request_infos, stop.walk is not None and step_walk is None) + + def _token_stops(self, pairs, request_infos, check_walk: bool) -> dict[str, set[str]]: + """Both stop checks' verdict on ``(rid, token)`` pairs.""" + stop = self.token_stop + stop_ids, loop, key = stop.stop_ids, stop.loop, stop.generated_key + stops: dict[str, set[str]] = {} + for rid, token in pairs: + info = request_infos[rid] + if check_walk and info.graph_walk != stop.walk: + continue + # the set test first: a stop id is rare, and the config read + # behind it goes through the info wrapper + if token in stop_ids and not info.resource_configs[stop.sampler].ignore_eos: + stops[rid] = {loop} + continue + generated = ( + info.dynamic_loop_iter_counts.get(loop, 0) + 1 if key is None + else int(self.request_state(rid).get(key, 0)) + ) + if generated >= info.max_tokens: + stops[rid] = {loop} + return stops + def split_inputs( self, graph_walk: str, diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index cfce625f9d..ed7f3cc397 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -462,15 +462,12 @@ def test_qwen3_5_caps_decode_where_its_stop_check_ends_it(max_tokens): from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel from mstar.model.qwen3_5.submodules import LLMSubmodule - llm = SimpleNamespace(config=SimpleNamespace(stop_token_ids=())) - - def check_stop(request_id, request_info, outputs): - return LLMSubmodule.check_stop(llm, request_id, request_info, outputs) + llm = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset())) cap = Qwen3_5DenseModel.get_loop_caps(object.__new__(Qwen3_5DenseModel), max_output_tokens=max_tokens) assert cap["decode_loop"] == _stops_after( - check_stop, "decode_loop", "decode", 2048, + llm.check_stop, "decode_loop", "decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, ), "the cap and the stop check end decode on different iterations" diff --git a/test/modular/test_qwen3_5_check_stop_batched.py b/test/modular/test_qwen3_5_check_stop_batched.py index 0441960842..f50b2d65bc 100644 --- a/test/modular/test_qwen3_5_check_stop_batched.py +++ b/test/modular/test_qwen3_5_check_stop_batched.py @@ -23,7 +23,7 @@ def _info(ignore_eos=False, done=0, max_tokens=100, walk="decode", last_prefill= def _sub(): - return SimpleNamespace(config=SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) + return LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({EOS, 9}))) @pytest.mark.parametrize("shape", ["flat", "column"]) @@ -118,8 +118,7 @@ def test_rows_out_of_batch_order_still_map_by_request(): from mstar.model.qwen3_5.submodules import LLMSubmodule from mstar.model.submodule_base import HostRows - sub = LLMSubmodule.__new__(LLMSubmodule) - sub.config = SimpleNamespace(stop_token_ids={7}) + sub = LLMSubmodule(None, SimpleNamespace(stop_token_ids=frozenset({7}))) rows = HostRows(("a", "b", "c"), {"new_token": torch.tensor([1, 7, 1])}) def info(max_tokens, iters=0): From 93ba258f17d9eff8f958fb4d131dd8dc59204732 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 55/59] Revert "Revert "OmniVoice: run the backbone's RMSNorms through the shared norm, not flashinfer"" This reverts commit c30a4c57dc992d66a9d44a23f687ce31f2f8fcfa. --- mstar/model/components/norm.py | 74 ++++++++++++-------- mstar/model/omnivoice/components/backbone.py | 37 +++++----- 2 files changed, 64 insertions(+), 47 deletions(-) diff --git a/mstar/model/components/norm.py b/mstar/model/components/norm.py index cfe9c33ad7..58b94ddc57 100644 --- a/mstar/model/components/norm.py +++ b/mstar/model/components/norm.py @@ -20,6 +20,50 @@ from mstar.engine.resources.rms_norm import run_rms_norm +def rms_norm( + hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, gemma_mode: bool = False, +) -> torch.Tensor: + """``RMSNorm``'s forward, for a norm whose weight lives elsewhere (e.g. a wrapped + Hugging Face module's): ``mstar::rms_norm``, or Gemma's fp32 ``(1 + weight)``.""" + if gemma_mode: + orig_dtype = hidden_states.dtype + x = hidden_states.to(torch.float32) + var = x.square().mean(dim=-1, keepdim=True) + normed = x * torch.rsqrt(var + eps) + normed = normed * (1.0 + weight.to(torch.float32)) + return normed.to(orig_dtype) + # Fused kernels expect 2D input. Reshape/restore so callers can pass + # arbitrary leading dims, such as [batch, sequence, hidden_size]. + orig_shape = hidden_states.shape + flat = hidden_states.reshape(-1, orig_shape[-1]) + return run_rms_norm(flat, weight, eps=eps).reshape(orig_shape) + + +def add_rms_norm( + hidden_states: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, + gemma_mode: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """``r = hidden_states + residual``; returns ``(rms_norm(r), r)``. + + On CUDA in half precision, one FlashInfer kernel, in place: ``residual += + hidden_states``, then ``hidden_states = norm(residual)``, so both inputs are + consumed. Inductor's own fusion of the pair is a reduction over a few rows that + takes ~3x as long at decode batch sizes. The kernel normalizes the sum before + rounding it to half precision, where the unfused pair rounds first. + """ + if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): + from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm + + orig_shape = hidden_states.shape + h = hidden_states.reshape(-1, orig_shape[-1]) + r = residual.reshape(-1, orig_shape[-1]) + fused = gemma_fused_add_rmsnorm if gemma_mode else fused_add_rmsnorm + fused(h, r, weight, eps) + return h.reshape(orig_shape), r.reshape(orig_shape) + r = hidden_states + residual + return rms_norm(r, weight, eps, gemma_mode), r + + class RMSNorm(nn.Module): """RMSNorm with optional Gemma-style ``(1 + weight)`` scaling. @@ -44,19 +88,7 @@ def __init__(self, hidden_size: int, eps: float = 1e-6, gemma_mode: bool = False self.weight = nn.Parameter(init) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.gemma_mode: - orig_dtype = hidden_states.dtype - x = hidden_states.to(torch.float32) - var = x.square().mean(dim=-1, keepdim=True) - normed = x * torch.rsqrt(var + self.variance_epsilon) - normed = normed * (1.0 + self.weight.to(torch.float32)) - return normed.to(orig_dtype) - # Fused kernels expect 2D input. Reshape/restore so callers can pass - # arbitrary leading dims, such as [batch, sequence, hidden_size]. - orig_shape = hidden_states.shape - flat = hidden_states.reshape(-1, orig_shape[-1]) - out = run_rms_norm(flat, self.weight, eps=self.variance_epsilon) - return out.reshape(orig_shape) + return rms_norm(hidden_states, self.weight, self.variance_epsilon, self.gemma_mode) def forward_residual( self, hidden_states: torch.Tensor, residual: torch.Tensor, @@ -78,21 +110,7 @@ def forward_residual( weight_bias=1.0 if self.gemma_mode else 0.0, ) return normed.reshape(orig_shape), r.reshape(orig_shape) - if hidden_states.is_cuda and hidden_states.dtype in (torch.bfloat16, torch.float16): - # One FlashInfer kernel, in place: `residual += hidden_states`, then - # `hidden_states = norm(residual)`. Inductor's own fusion of the - # pair is a reduction over a few rows that takes ~3x as long at - # decode batch sizes. Both inputs are consumed by the caller. - from flashinfer.norm import fused_add_rmsnorm, gemma_fused_add_rmsnorm - - orig_shape = hidden_states.shape - h = hidden_states.reshape(-1, orig_shape[-1]) - r = residual.reshape(-1, orig_shape[-1]) - fused = gemma_fused_add_rmsnorm if self.gemma_mode else fused_add_rmsnorm - fused(h, r, self.weight, self.variance_epsilon) - return h.reshape(orig_shape), r.reshape(orig_shape) - r = hidden_states + residual - return self(r), r + return add_rms_norm(hidden_states, residual, self.weight, self.variance_epsilon, self.gemma_mode) def extra_repr(self) -> str: return f"{self.hidden_size}, eps={self.variance_epsilon}, gemma_mode={self.gemma_mode}" diff --git a/mstar/model/omnivoice/components/backbone.py b/mstar/model/omnivoice/components/backbone.py index bb227d1faa..f77ba81c4a 100644 --- a/mstar/model/omnivoice/components/backbone.py +++ b/mstar/model/omnivoice/components/backbone.py @@ -7,11 +7,12 @@ nothing, and the head GEMM runs only at the positions whose logits are read. The per-step kernels are the reference's own fastest mode: ``apply_flashinfer`` -fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE, -RMSNorm and ``silu_and_mul``. ``OmniVoiceBackbone.encode`` drives those fused -weights with its own layer loop rather than through the patched ``llm``: it -norms q and k straight off the fused QKV output (the reference copies each out -first) and folds every residual add into the next RMSNorm. At 16 requests a +fuses QKV and gate/up, and the layer runs in ``(S, H, D)`` with flashinfer RoPE +and ``silu_and_mul`` (its RMSNorms go through M*'s shared ones). +``OmniVoiceBackbone.encode`` drives those fused weights with its own layer loop +rather than through the patched ``llm``: it norms q and k straight off the fused +QKV output (the reference copies each out first) and folds every residual add +into the next RMSNorm. At 16 requests a step those copies and adds were a tenth of the step's GPU time. What M* adds on top is what the reference cannot do from inside one @@ -34,6 +35,8 @@ import torch.nn.functional as F from torch import nn +from mstar.engine.resources.rms_norm import run_rms_norm +from mstar.model.components.norm import add_rms_norm from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) @@ -293,30 +296,26 @@ def encode( norm = layer.input_layernorm if residual is None: residual = hidden - x = flashinfer.norm.rmsnorm(hidden, norm.weight, eps=norm.variance_epsilon) + x = run_rms_norm(hidden, norm.weight, eps=norm.variance_epsilon) else: - # residual += hidden; hidden = norm(residual), in place. The kernel - # normalizes the sum before rounding it to half precision, where - # the reference rounds first: logits match it to that rounding, - # not bit for bit. - flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) - x = hidden + # The fused add-norm normalizes the sum before rounding it to half + # precision, where the reference rounds first: logits match it to + # that rounding, not bit for bit. + x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) n, d = x.shape[0], sa.head_dim heads_q, heads_k, _ = (width // d for width in sa._fi_qkv_split) qkv = F.linear(x, sa._fi_w_qkv).view(n, -1, d) # the q/k norms read their heads straight off the fused output - q = flashinfer.norm.rmsnorm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) - k = flashinfer.norm.rmsnorm( - qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon, - ) + q = run_rms_norm(qkv[:, :heads_q], sa.q_norm.weight, eps=sa.q_norm.variance_epsilon) + k = run_rms_norm(qkv[:, heads_q:heads_q + heads_k], sa.k_norm.weight, eps=sa.k_norm.variance_epsilon) flashinfer.rope.apply_rope_pos_ids_inplace( q, k, pos_ids, rope_theta=sa._fi_rope_theta, interleave=False, ) hidden = sa.o_proj(attn.run(q, k, qkv[:, heads_q + heads_k:]).reshape(n, -1)) norm = layer.post_attention_layernorm - flashinfer.norm.fused_add_rmsnorm(hidden, residual, norm.weight, eps=norm.variance_epsilon) - hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(hidden, mlp._fi_w_gate_up))) - flashinfer.norm.fused_add_rmsnorm(hidden, residual, llm.norm.weight, eps=llm.norm.variance_epsilon) + x, residual = add_rms_norm(hidden, residual, norm.weight, norm.variance_epsilon) + hidden = mlp.down_proj(flashinfer.activation.silu_and_mul(F.linear(x, mlp._fi_w_gate_up))) + hidden, _ = add_rms_norm(hidden, residual, llm.norm.weight, llm.norm.variance_epsilon) return hidden def head(self, hidden: torch.Tensor, tgt_index: torch.Tensor) -> torch.Tensor: From f8ae22eadfd59bfed97f74622ebb34942dd46b43 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 56/59] Revert "Revert "Delete RaggedAttentionCallable: callers bind the label on the resource's run"" This reverts commit 4bf4ea0fcead564f820721b53d6fccc9809486a8. --- .claude/skills/diffusion-scaffold/SKILL.md | 2 +- docs/adding_models.rst | 2 +- mstar/engine/resources/attn/ragged/config.py | 2 +- mstar/engine/resources/convenience.py | 46 +------------------ .../components/diffusion/denoise_loop.py | 16 +++---- mstar/model/ltx2_5/submodules.py | 7 ++- mstar/model/omnilingual_asr/submodules.py | 10 ++-- mstar/model/wan22/submodules.py | 7 ++- test/modular/test_diffusion_denoise_loop.py | 16 +++---- test/modular/test_ragged_attention.py | 11 ++--- 10 files changed, 34 insertions(+), 85 deletions(-) diff --git a/.claude/skills/diffusion-scaffold/SKILL.md b/.claude/skills/diffusion-scaffold/SKILL.md index b921d45181..28416973f1 100644 --- a/.claude/skills/diffusion-scaffold/SKILL.md +++ b/.claude/skills/diffusion-scaffold/SKILL.md @@ -41,7 +41,7 @@ A recipe such as LTX-2.5's two-stage distilled one (half-resolution denoise → Every attention in the DiT should go through a ragged resource, so the model gets the engine's kernels and its backend can change in config. - `attention_segments(bucket_key)` declares every span a request's layers attend over, as `(label, length)`; the default is one `"main"` span of all its tokens. A layer that attends over a sub-span (a refiner over the image tokens alone, a joint model's per-modality self-attention) gets its own label, and `ragged_for(label)` gives that layer its callable. -- `ragged_for` memoizes one callable per label for the life of a binding. The compiled transformer guards on the identity of the callables it is handed; a fresh one per step recompiles every step. +- `ragged_for` memoizes one callable per label (the resource's `run` with the label bound) for the life of a binding, and a rebinding drops them so none attends through the old resource. - A model with more than one head geometry (LTX-2.5's video 32×128 and audio 32×64 heads) declares one resource per geometry. - Cross-attention between spans of the same request (video queries over text keys) is a separate cross-attention resource kind (`RaggedCrossAttentionSpec`, added in #372) whose step names `(q_label, kv_label)` pairs; see "Cross-attention between spans" in `docs/adding_models.rst`. Don't fall back to SDPA for it because the shapes are fixed inside a bucket. - `joint_attention` in `attention.py` packs per-row q/k/v into the ragged layout (or runs SDPA when no resource is bound, which is what a parity backend uses). diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 3780a6d44d..80fe3221d8 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -772,7 +772,7 @@ step names the ``(q_label, kv_label)`` pairs: Each pair is planned as its own layout, pairing each request's ``q_label`` span with that request's ``kv_label`` span; cross-attention is never causal. A layer attends through -``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``, passing the key +``resource.run(q, k, v, label=cross_label(q_label, kv_label))``, passing the key span's tokens as ``k`` and ``v``. One resource has one head geometry, so a model whose attentions differ in head count or head dim declares one spec per (kind, geometry). LTX-2.5 is the reference (``mstar/model/ltx2_5/submodules.py``). diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 1fa1fdeb10..e258a0756a 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -159,7 +159,7 @@ def resource_class(self) -> "type[Resource]": class RaggedCrossAttentionSpec(_RaggedSpecBase): """Cacheless cross-attention between two spans of each request: a forward's ``(q_label, kv_label)`` pairs (``RaggedCrossAttentionStep``), each run through - ``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``. + ``resource.run(q, k, v, label=cross_label(q_label, kv_label))``. Separate from ``RaggedAttentionSpec`` so a resource is one kind of attention: a model with both declares one of each per head geometry. diff --git a/mstar/engine/resources/convenience.py b/mstar/engine/resources/convenience.py index ed46c83b15..7a26f06c2a 100644 --- a/mstar/engine/resources/convenience.py +++ b/mstar/engine/resources/convenience.py @@ -1,11 +1,6 @@ import torch from mstar.engine.resources.attn.base import AttentionManager -from mstar.engine.resources.attn.ragged.base import ( - RaggedAttnManager, - RaggedBlockCausalAttnManager, - RaggedCrossAttnManager, -) from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.linear_attn.base import LinearAttnManager from mstar.engine.resources.recurrent.pool import RecurrentStatePool @@ -26,7 +21,8 @@ class AttentionCallable: varies its label per layer today; one that needs to should thread the label explicitly rather than use this. - For a cacheless attention, see :class:`RaggedAttentionCallable`. + A cacheless (ragged) attention needs no callable: call the resource's + ``run(q, k, v, label=...)`` with the span's label bound. """ def __init__(self, kv: KVManager, attn: AttentionManager | None=None): @@ -128,44 +124,6 @@ def __call__( ) -class RaggedAttentionCallable: - """`AttentionCallable` for a cacheless attention: ``(q, k, v) -> out`` over - one declared span of a `RaggedAttnManager`. - - No ``kv`` and no layer cursor, because there is neither to carry: a ragged - plan is keyed by the segment label, and the wrapper behind it is per - (bucket, slot, label) rather than per layer. The label is therefore bound - here at construction instead of being a cursor on the resource -- which is - what lets one stack attend over several spans at once (a refiner over the - image tokens next to a pass over all of them), the case `AttentionCallable` - says to thread explicitly. A cross-attention pair of a - ``RaggedCrossAttentionSpec`` resource is bound the same way, under - ``cross_label(q_label, kv_label)``; its ``k`` / ``v`` are then the key span's - tokens, packed in the same request order as ``q``. A - ``RaggedBlockCausalAttentionSpec`` resource binds like a self-attention one. - - One instance per label, held for the life of the resource binding: a - compiled transformer region guards on the identity of the callables it is - handed, so handing it a fresh one per step retraces that frame every step - until Dynamo gives up. - """ - - def __init__( - self, - attn: RaggedAttnManager | RaggedCrossAttnManager | RaggedBlockCausalAttnManager, - label: str, - ): - self.attn = attn - self.label = label - - # Not compiler-disabled: a compiled region traces through to the resource's - # ``run``, which keeps the attention in its graph as one custom op. - def __call__( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - ) -> torch.Tensor: - return self.attn.run(q, k, v, label=self.label) - - class Mamba2Callable: """``LinearAttnCallable`` for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor. diff --git a/mstar/model/components/diffusion/denoise_loop.py b/mstar/model/components/diffusion/denoise_loop.py index e120a70007..2b7dea73fc 100644 --- a/mstar/model/components/diffusion/denoise_loop.py +++ b/mstar/model/components/diffusion/denoise_loop.py @@ -44,6 +44,7 @@ from __future__ import annotations +import functools import logging from collections import OrderedDict from collections.abc import Hashable, Mapping, Sequence @@ -54,7 +55,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep -from mstar.engine.resources.convenience import RaggedAttentionCallable +from mstar.model.components.diffusion.attention import RaggedAttentionFn from mstar.model.components.diffusion.flow_match import FlowMatchSchedule from mstar.model.components.diffusion.noise import NoiseStager from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule @@ -142,7 +143,7 @@ def __init__( self._noise_stagers: dict[torch.dtype, NoiseStager] = {} # One callable per attention label for the life of a binding; see # :meth:`ragged_for`. - self._ragged_fns: dict[str, RaggedAttentionCallable] = {} + self._ragged_fns: dict[str, RaggedAttentionFn] = {} @property def loop_back_names(self) -> tuple[str, ...]: @@ -279,7 +280,7 @@ def layout(self, bucket_key: Hashable, device: torch.device) -> Any: def step_index(self, fwd_info: CurrentForwardPassInfo) -> int: return int(fwd_info.dynamic_loop_iter_counts.get(self.loop_name, 0)) - def _ragged(self) -> RaggedAttentionCallable | None: + def _ragged(self) -> RaggedAttentionFn | None: """The attention callable for the ``"main"`` span, or None (SDPA).""" return self.ragged_for(self.MAIN_SPAN) @@ -287,20 +288,19 @@ def bind_node_resources(self, resources) -> None: super().bind_node_resources(resources) self._ragged_fns.clear() # the memoized callables hold the old resource - def ragged_for(self, label: str) -> RaggedAttentionCallable | None: + def ragged_for(self, label: str) -> RaggedAttentionFn | None: """``(q, k, v) -> out`` over the segments declared under ``label``, or None when no ragged resource is bound (layers then fall back to SDPA). - One callable per label for the life of the binding: the compiled transformer - region guards on the identity of the callables it is handed, and a fresh one - per step would recompile it every step until dynamo gave up.""" + The resource's ``run`` with the label bound, built once per label for the life + of the binding rather than per step.""" if self.attn_resource_key is None: return None if label not in self._ragged_fns: resource = self.node_resources.get(self.attn_resource_key) if resource is None: return None - self._ragged_fns[label] = RaggedAttentionCallable(resource, label) + self._ragged_fns[label] = functools.partial(resource.run, label=label) return self._ragged_fns[label] # --------------------------------------------------------- engine contract diff --git a/mstar/model/ltx2_5/submodules.py b/mstar/model/ltx2_5/submodules.py index 9205c32a31..626223a24f 100644 --- a/mstar/model/ltx2_5/submodules.py +++ b/mstar/model/ltx2_5/submodules.py @@ -31,7 +31,6 @@ SubmoduleStep, cross_label, ) -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.components.batched_rows import BatchedRows from mstar.model.components.diffusion.attention import joint_attention, sdpa_attention from mstar.model.components.diffusion.compile_utils import compile_transformer_forward @@ -303,15 +302,15 @@ def bind_node_resources(self, resources) -> None: self._attends = None # the memoized callables hold the old resources def attends(self) -> LTX2Attends: - """The six attentions' callables, built once per resource binding: a compiled - region guards on their identity (see ``DenoiseLoopSubmodule.ragged_for``).""" + """The six attentions' callables, built once per resource binding rather than + per step.""" if self._attends is None: def on(key, label): resource = self.node_resources.get(key) if self.use_ragged_attention else None if resource is None: return sdpa_attention # every row has the bucket's spans, so [B, L, H, D] packs per request - return functools.partial(joint_attention, ragged=RaggedAttentionCallable(resource, label)) + return functools.partial(joint_attention, ragged=functools.partial(resource.run, label=label)) self._attends = LTX2Attends( video_self=on(VIDEO_ATTN, VIDEO_SPAN), audio_self=on(AUDIO_ATTN, AUDIO_SPAN), diff --git a/mstar/model/omnilingual_asr/submodules.py b/mstar/model/omnilingual_asr/submodules.py index fae5b48cad..221003f49b 100644 --- a/mstar/model/omnilingual_asr/submodules.py +++ b/mstar/model/omnilingual_asr/submodules.py @@ -19,7 +19,8 @@ from __future__ import annotations -from collections.abc import Mapping +import functools +from collections.abc import Callable, Mapping from contextlib import contextmanager from typing import Any @@ -35,7 +36,6 @@ ) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources.attn.config import AttentionStep -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep from mstar.model.omnilingual_asr.components.wav2vec2 import Wav2Vec2CTC from mstar.model.omnilingual_asr.config import ENCODER_ATTN, SAMPLE_RATE, TOKEN_IDS, WAVEFORM @@ -92,7 +92,7 @@ def __init__(self, model: Wav2Vec2CTC, max_batch_size: int, max_batch_seconds: f self._compiled_features = with_inductor_overrides( torch.compile(model.encoder_frontend.feature_extractor, dynamic=True) ) - self._attn_fns: dict[int, RaggedAttentionCallable] = {} + self._attn_fns: dict[int, Callable] = {} def warm_compile(self, device: str) -> None: """Compile the feature extractor at load rather than on a live request: @@ -179,12 +179,12 @@ def _capture_encoder(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: attn = self._attn_fn(inp.resources[ENCODER_ATTN]) return {"ids": self.model.encode_packed(inp.static_inputs["hidden"], attn).argmax(dim=-1)} - def _attn_fn(self, resource) -> RaggedAttentionCallable: + def _attn_fn(self, resource) -> Callable: """The resource's attention over the step's segments, one callable per bound resource.""" fn = self._attn_fns.get(id(resource)) if fn is None: - fn = self._attn_fns[id(resource)] = RaggedAttentionCallable(resource, "main") + fn = self._attn_fns[id(resource)] = functools.partial(resource.run, label="main") return fn def max_batch_size(self, graph_walk: str) -> int: diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 1a09d2da18..f171a2c8cd 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -17,6 +17,7 @@ the v1 engine supplies neither the autocast nor the no-grad wrapper. """ +import functools import logging import os from contextlib import contextmanager @@ -28,7 +29,6 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep -from mstar.engine.resources.convenience import RaggedAttentionCallable from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule from mstar.model.wan22.components.dit import Wan22AttentionFns from mstar.model.wan22.components.unipc import ( @@ -497,14 +497,13 @@ def bind_node_resources(self, resources) -> None: self._attn_fns = None # the memoized callables hold the old resources def _attention_fns(self) -> Wan22AttentionFns | None: - """The ragged callables, built once per resource binding (the compiled - transformer guards on their identity), or None to run SDPA.""" + """The ragged callables, built once per resource binding, or None to run SDPA.""" if self._attn_fns is None: self_attn = self.node_resources.get(self.SELF_ATTN) if self_attn is None: return None self._attn_fns = Wan22AttentionFns( - self_attn=RaggedAttentionCallable(self_attn, self.VIDEO_SPAN), cross_attn=None, + self_attn=functools.partial(self_attn.run, label=self.VIDEO_SPAN), cross_attn=None, ) logger.info("Wan2.2 DiT: self-attention through the engine's ragged attention resource") return self._attn_fns diff --git a/test/modular/test_diffusion_denoise_loop.py b/test/modular/test_diffusion_denoise_loop.py index 38897e3a1b..4222bd7a00 100644 --- a/test/modular/test_diffusion_denoise_loop.py +++ b/test/modular/test_diffusion_denoise_loop.py @@ -17,7 +17,6 @@ sys.path.insert(0, ".") from mstar.conductor.request_info import CurrentForwardPassInfo # noqa: E402 -from mstar.engine.resources.convenience import RaggedAttentionCallable # noqa: E402 from mstar.model.components.diffusion.denoise_loop import ( # noqa: E402 LATENTS, DenoiseLoopSubmodule, @@ -397,25 +396,22 @@ class OffsetPastTheEnd(ToyDenoise): def test_ragged_for_binds_the_label(): class Resource: def run(self, q, k, v, label=None): - return label + return self, label sub = TwoSpanDenoise(attn_resource_key="dit_attn") resource = Resource() sub.bind_node_resources({"dit_attn": resource}) image = sub.ragged_for("image") - assert isinstance(image, RaggedAttentionCallable) - assert image.attn is resource and image.label == "image" - assert image(None, None, None) == "image" - assert sub._ragged()(None, None, None) == "main" - # one callable object per label, so a compiled region's identity guards hold across steps + assert image(None, None, None) == (resource, "image") + assert sub._ragged()(None, None, None) == (resource, "main") + # built once per label for the binding, not per step assert sub.ragged_for("image") is image and sub._ragged() is sub._ragged() # a rebinding starts over: a callable holding the old resource would attend # through a plan nothing writes to any more new_resource = Resource() sub.bind_node_resources({"dit_attn": new_resource}) - assert sub.ragged_for("image") is not image - assert sub.ragged_for("image").attn is new_resource - assert sub.ragged_for("main")(None, None, None) == "main" + assert sub.ragged_for("image")(None, None, None) == (new_resource, "image") + assert sub.ragged_for("main")(None, None, None) == (new_resource, "main") assert ToyDenoise().ragged_for("image") is None # no resource declared -> SDPA unbound = TwoSpanDenoise(attn_resource_key="dit_attn") assert unbound.ragged_for("image") is None # declared but not bound yet -> SDPA diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 5d1bb9e052..92bf27a513 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -12,6 +12,7 @@ from __future__ import annotations +import functools import sys sys.path.insert(0, ".") @@ -325,13 +326,11 @@ class _Group: def test_compiled_caller_keeps_attention_in_its_graph(): """Traced by ``torch.compile``, the resource's ``run`` is one custom op: the region keeps one graph around it (no break) and computes what eager does.""" - from mstar.engine.resources.convenience import RaggedAttentionCallable - mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fn = RaggedAttentionCallable(mgr, "main") + fn = functools.partial(mgr.run, label="main") def region(q, k, v): return fn(q * 1.0, k, v) + 1.0 @@ -346,18 +345,16 @@ def region(q, k, v): def test_callable_built_inside_a_compiled_region(): """A node that builds its callables lazily on first forward does so inside the compiled region; later calls still reach the resource.""" - from mstar.engine.resources.convenience import RaggedAttentionCallable - mgr = manager() spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) - fns: dict[int, RaggedAttentionCallable] = {} + fns: dict[int, functools.partial] = {} def region(q, k, v): fn = fns.get(id(mgr)) if fn is None: - fn = fns[id(mgr)] = RaggedAttentionCallable(mgr, "main") + fn = fns[id(mgr)] = functools.partial(mgr.run, label="main") return fn(q, k, v) + 1.0 torch._dynamo.reset() From 0a5790140e859f677aed7ddbf124a1251ab268e4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:34 -0700 Subject: [PATCH 57/59] Revert "Revert "Whisper, Qwen3-Omni Thinker, Chatterbox T3: declare their stop as a TokenStop"" This reverts commit ab3db58f91610049802449d0b7fcc8e009f00f1e. --- mstar/model/chatterbox/submodules.py | 20 +++++--------------- mstar/model/qwen3_omni/submodules.py | 21 ++++++--------------- mstar/model/whisper/submodules.py | 20 ++++---------------- test/modular/test_model_kv_layouts.py | 2 +- 4 files changed, 16 insertions(+), 47 deletions(-) diff --git a/mstar/model/chatterbox/submodules.py b/mstar/model/chatterbox/submodules.py index 96491b40cb..f874df5bf5 100644 --- a/mstar/model/chatterbox/submodules.py +++ b/mstar/model/chatterbox/submodules.py @@ -60,6 +60,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) logger = logging.getLogger(__name__) @@ -220,6 +221,10 @@ def __init__( self.t3 = config.t3 self.builtin_voice = builtin_voice self.supports_cfg = not config.is_turbo + self.token_stop = TokenStop( + loop=DECODE_LOOP, buffer=SPEECH_TOKENS, stop_ids=frozenset({config.t3.stop_speech_token}), + sampler=T3_SAMPLER, generated_key="generated", + ) # -- request knobs ------------------------------------------------------- @@ -441,21 +446,6 @@ def postprocess( state = self.request_state(request_id) state.add("generated", int(state.get("generated", 0)) + 1) - def check_stop( - self, request_id: str, request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if SPEECH_TOKENS not in outputs: - return set() - token = int(outputs[SPEECH_TOKENS][0].reshape(-1)[0].item()) - generated = int(self.request_state(request_id).get("generated", 0)) - max_new = request_info.max_tokens - sampling = request_info.resource_configs.get(T3_SAMPLER) - ignore_eos = bool(getattr(sampling, "ignore_eos", False)) - if (not ignore_eos and token == self.t3.stop_speech_token) or generated >= max_new: - return {DECODE_LOOP} - return set() - # -- batching and capture --------------------------------------------------- def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 68334778c7..112abb7110 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -64,6 +64,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) logger = logging.getLogger(__name__) @@ -351,6 +352,11 @@ def __init__( super().__init__() self.model = thinker_model # Qwen3OmniThinkerModel self.config = config + # only a decode row's token ends the loop; a prefill is not in it yet + self.token_stop = TokenStop( + loop="thinker_decode_loop", buffer="new_token", stop_ids=frozenset({config.im_end_token_id}), + sampler=THINKER_SAMPLER, walk="thinker_decode", + ) # Pre-compute inverse frequencies for 3D MRoPE self._inv_freq: torch.Tensor | None = None @@ -1143,21 +1149,6 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] - def check_stop( - self, request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs or request_info.graph_walk != "thinker_decode": - return set() - token = outputs["new_token"][0].item() - ignore_eos= request_info.resource_configs[THINKER_SAMPLER].ignore_eos - eos_token_id = self.config.im_end_token_id - if (not ignore_eos and eos_token_id == token) or \ - (request_info.dynamic_loop_iter_counts.get("thinker_decode_loop", 0) + 1 >= request_info.max_tokens): - return {"thinker_decode_loop"} - return set() - # =================================================================== # 4. TalkerSubmodule (ar engine) -- SECOND MOST COMPLEX diff --git a/mstar/model/whisper/submodules.py b/mstar/model/whisper/submodules.py index e8969d625c..4de63cad6e 100644 --- a/mstar/model/whisper/submodules.py +++ b/mstar/model/whisper/submodules.py @@ -30,6 +30,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + TokenStop, ) from mstar.model.whisper.components.decoder import WhisperDecoderModel from mstar.model.whisper.config import ( @@ -124,6 +125,9 @@ def __init__(self, decoder: WhisperDecoderModel, config: WhisperModelConfig): super().__init__() self.decoder = decoder self.config = config + self.token_stop = TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=frozenset({config.eos_token_id}), sampler=SAMPLER, + ) self._suppress_ids: torch.Tensor | None = None self._begin_suppress_ids: torch.Tensor | None = None @@ -339,19 +343,3 @@ def postprocess( if "new_token" not in outputs: return outputs["text_inputs"] = outputs["new_token"] - - def check_stop( - self, - request_id: str, - request_info: CurrentForwardPassInfo, - outputs: dict[str, list[torch.Tensor]], - ) -> set[str]: - if "new_token" not in outputs: - return set() - token = outputs["new_token"][0].item() - ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos - decoded_tokens = request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 - if (not ignore_eos and token == self.config.eos_token_id) or \ - decoded_tokens >= request_info.max_tokens: - return {"decode_loop"} - return set() diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py index ed7f3cc397..d66858861a 100644 --- a/test/modular/test_model_kv_layouts.py +++ b/test/modular/test_model_kv_layouts.py @@ -486,7 +486,7 @@ def test_qwen3_omni_caps_the_thinkers_decode_where_its_stop_check_ends_it(max_to cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), max_output_tokens=max_tokens) assert cap["thinker_decode_loop"] == _stops_after( - lambda *args: ThinkerSubmodule.check_stop(_omni_submodule(), *args), "thinker_decode_loop", "thinker_decode", + ThinkerSubmodule(None, _omni_submodule().config).check_stop, "thinker_decode_loop", "thinker_decode", 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, resource_configs={THINKER_SAMPLER: SamplingReqConfig(ignore_eos=True)}, ), "the cap and the stop check end decode on different iterations" From 4f2deb3358210aeafaeb3a385e876f35ea9b4740 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 19:22:46 -0700 Subject: [PATCH 58/59] test: the ragged rebuild test binds the label on the resource's run Co-Authored-By: Claude Opus 5.5 --- test/modular/test_ragged_attention.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 92bf27a513..4ffc5162ab 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -368,8 +368,6 @@ def test_a_compiled_region_attends_through_a_rebuilt_resource(): """The op names its manager by resource key: after the resource is rebuilt under the same key, a region compiled against the old one attends through the new one's plan, not the old one's.""" - from mstar.engine.resources.convenience import RaggedAttentionCallable - q, k, v = qkv(224) def region(fn, q, k, v): @@ -379,11 +377,11 @@ def region(fn, q, k, v): compiled = torch.compile(region) old = manager() old.plan(step(128, 96), ctx(["r0", "r1"])) - close(compiled(RaggedAttentionCallable(old, "main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) + close(compiled(functools.partial(old.run, label="main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) rebuilt = manager() rebuilt.plan(step(64, 160), ctx(["r0", "r1"])) - close(compiled(RaggedAttentionCallable(rebuilt, "main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0) + close(compiled(functools.partial(rebuilt.run, label="main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0) def test_plan_then_run_matches_reference(): From 6782504c578d7ea3d14795856ec1f591a08af705 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Fri, 23 Oct 2026 21:57:00 -0700 Subject: [PATCH 59/59] ragged attention: the compile op is opt-in per resource (compile_op), on for Wan2.2 On H200, routing every compiled ragged caller through the op made Z-Image (attention_backend: flashinfer) trace its transformer as one graph per capture bucket: +2.6 GiB of capture memory and 1.7-2.7x slower startup, with no step-time gain since the step is captured. Wan2.2's DiT, compiled and uncaptured, is where keeping attention in the graph pays, so it declares compile_op; everyone else breaks at the attention as on main. Also initialize the Qwen3 encoder test's parameters: FusedColumnLinear allocates with torch.empty, so the test read whatever memory it reused and failed after some other modules. Co-Authored-By: Claude Opus 5.5 --- .claude/skills/diffusion-scaffold/SKILL.md | 2 +- docs/adding_models.rst | 8 ++++--- mstar/engine/resources/attn/ragged/config.py | 7 +++++++ .../resources/attn/ragged/flashinfer.py | 15 +++++++------ mstar/model/wan22/wan22_model.py | 2 ++ test/modular/test_diffusion_numerics.py | 3 +++ test/modular/test_ragged_attention.py | 21 ++++++++++--------- 7 files changed, 38 insertions(+), 20 deletions(-) diff --git a/.claude/skills/diffusion-scaffold/SKILL.md b/.claude/skills/diffusion-scaffold/SKILL.md index 28416973f1..9464972e45 100644 --- a/.claude/skills/diffusion-scaffold/SKILL.md +++ b/.claude/skills/diffusion-scaffold/SKILL.md @@ -50,7 +50,7 @@ Every attention in the DiT should go through a ragged resource, so the model get - The engine's blanket compile is off for the node (`disable_torch_compile = True`); the model compiles its transformer region with `compile_utils.compile_transformer_forward`, which keeps the exact-op exclusions that preserve eager rounding. Compiling the DiT mostly buys the memory-bound elementwise work (modulation, fp32 norms, RoPE, gating), which on LTX-2.5 was ~45% of an eager step. - `capture_buckets` lists `(graph_walk, bucket_key)` pairs to capture; `capture_request_inputs` supplies one placeholder row per bucket. A bucket is captured per batch size in `capture_batch_sizes`, and each one costs startup, so capture the shapes you serve, not every shape you accept. `replay_walks` lets another walk reuse a captured bucket. -- The captured graph removes the launch gaps of a compiled-but-eager step. Ragged attention does not break that compiled graph (the resource's `run` traces as the `mstar::ragged_attention` custom op), so measure both before deciding a model needs capture. +- The captured graph removes the launch gaps of a compiled-but-eager step. Ragged attention breaks that compiled graph per layer unless the resource sets `compile_op=True` (then `run` traces as the `mstar::ragged_attention` custom op), so measure both before deciding a model needs capture. ## The overshoot veto diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 80fe3221d8..cab6b0dc49 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -742,9 +742,11 @@ encoder node only, and names no cache: A layer binds the key as it binds any other resource, then attends through the resource's ``run(q, k, v, label=None)``. There is no KV write to pair it with, so no -``AttentionCallable`` is involved. Inside a ``torch.compile`` region, ``run`` is the -``mstar::ragged_attention`` custom op, so the region keeps one graph across the attention -instead of breaking per layer; an eager call goes straight to the kernel. +``AttentionCallable`` is involved. A compiled caller breaks at ``run``; with +``compile_op=True`` in the config, ``run`` traces as the ``mstar::ragged_attention`` +custom op instead, so the region keeps one graph across the attention (worth it for a +compiled, uncaptured region; it costs startup time and capture memory where the region +is also captured). An eager call goes straight to the kernel. The step is an ordinary ``AttentionStep``, and its ``segments`` are the entire layout. There is no cache to read the layout off, so a label that carries no segment in the diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index e258a0756a..f7e241c401 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -80,6 +80,13 @@ class RaggedAttentionConfig: # guidance branch per rank, is the one user. shard_heads: bool = True + # Whether a compiled caller keeps the attention in its graph, as the + # ``mstar::ragged_attention`` op, rather than breaking there per layer. On for + # a compiled, uncaptured region whose fusions around attention pay (Wan2.2's + # DiT); off by default, since tracing a whole transformer as one graph per + # capture bucket costs startup time and capture memory (Z-Image: +2.6 GiB). + compile_op: bool = False + # Activation dtype of the q / k / v the node hands the kernel (and of its # output). A cacheless attention doesn't need to adhere to the dtype of its # KV cache, so it is free to declare a separate dtype here; when None, the diff --git a/mstar/engine/resources/attn/ragged/flashinfer.py b/mstar/engine/resources/attn/ragged/flashinfer.py index 0d04277195..9d727bed85 100644 --- a/mstar/engine/resources/attn/ragged/flashinfer.py +++ b/mstar/engine/resources/attn/ragged/flashinfer.py @@ -252,14 +252,17 @@ def run( ) -> torch.Tensor: """One layer's attention through the plan for ``label``. - Traced by ``torch.compile``, this is the ``mstar::ragged_attention`` op: - the region keeps one graph, and Inductor fuses the pointwise work on - either side of the attention instead of breaking per layer. Called - eagerly, including in a captured region that is not compiled, it is the - plain kernel call, with no op dispatch. + Traced by ``torch.compile`` with ``config.compile_op``, this is the + ``mstar::ragged_attention`` op: the region keeps one graph, and Inductor + fuses the pointwise work on either side of the attention. Otherwise a + compiled caller breaks here, and an eager one calls the kernel directly. """ - if torch.compiler.is_compiling(): + if self._config.compile_op and torch.compiler.is_compiling(): return ragged_attention(q, k, v, self._key, label) + return self._run(q, k, v, label) + + @torch.compiler.disable + def _run(self, q, k, v, label): return self._wrapper(label).run(q, k, v) diff --git a/mstar/model/wan22/wan22_model.py b/mstar/model/wan22/wan22_model.py index 9e09814e9e..db91b8c47c 100644 --- a/mstar/model/wan22/wan22_model.py +++ b/mstar/model/wan22/wan22_model.py @@ -146,6 +146,8 @@ def get_node_resources(self) -> list[NodeResourceSpec]: max_segments_per_request=2, # CFG-parallel: the dit's two ranks split guidance branches, not heads shard_heads=False, + # the DiT is compiled, not captured: keep attention in its graph + compile_op=True, dtype=torch.bfloat16, ), ), diff --git a/test/modular/test_diffusion_numerics.py b/test/modular/test_diffusion_numerics.py index 5b1bf8e5cd..5b228fe910 100644 --- a/test/modular/test_diffusion_numerics.py +++ b/test/modular/test_diffusion_numerics.py @@ -88,6 +88,9 @@ def test_qwen3_encoder_builds_one_rotary_table_across_prompt_lengths(monkeypatch head_dim=8, rms_norm_eps=1e-6, rope_theta=10_000.0, tap_layers=(1,), ) encoder = qwen3_encoder.Qwen3HiddenStateEncoder(**dims, max_positions=16).eval() + with torch.no_grad(): + for p in encoder.parameters(): # some layers allocate with torch.empty + p.normal_(std=0.1) builds = [] build = qwen3_encoder.qwen3_rotary_tables monkeypatch.setattr(qwen3_encoder, "qwen3_rotary_tables", lambda *a: builds.append(a) or build(*a)) diff --git a/test/modular/test_ragged_attention.py b/test/modular/test_ragged_attention.py index 4ffc5162ab..64f1c1de55 100644 --- a/test/modular/test_ragged_attention.py +++ b/test/modular/test_ragged_attention.py @@ -323,10 +323,12 @@ class _Group: assert mgr._kwargs["head_dim"] == HEAD_DIM -def test_compiled_caller_keeps_attention_in_its_graph(): - """Traced by ``torch.compile``, the resource's ``run`` is one custom op: the - region keeps one graph around it (no break) and computes what eager does.""" - mgr = manager() +@pytest.mark.parametrize("compile_op", [True, False]) +def test_compiled_caller_keeps_attention_in_its_graph(compile_op): + """With ``compile_op``, the resource's ``run`` traces as one custom op: the region + keeps one graph around it and computes what eager does. Without it (the default) + the region breaks at the attention.""" + mgr = manager(compile_op=compile_op) spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) @@ -337,15 +339,14 @@ def region(q, k, v): torch._dynamo.reset() explain = torch._dynamo.explain(region)(q, k, v) - assert explain.graph_break_count == 0 - compiled = torch.compile(region, fullgraph=True) - close(compiled(q, k, v), region(q, k, v)) + assert (explain.graph_break_count == 0) is compile_op + close(torch.compile(region)(q, k, v), region(q, k, v)) def test_callable_built_inside_a_compiled_region(): """A node that builds its callables lazily on first forward does so inside the compiled region; later calls still reach the resource.""" - mgr = manager() + mgr = manager(compile_op=True) spans = [128, 96] mgr.plan(step(*spans), ctx(["r0", "r1"])) q, k, v = qkv(sum(spans)) @@ -375,11 +376,11 @@ def region(fn, q, k, v): torch._dynamo.reset() compiled = torch.compile(region) - old = manager() + old = manager(compile_op=True) old.plan(step(128, 96), ctx(["r0", "r1"])) close(compiled(functools.partial(old.run, label="main"), q, k, v), ref(q, k, v, cu([128, 96])) + 1.0) - rebuilt = manager() + rebuilt = manager(compile_op=True) rebuilt.plan(step(64, 160), ctx(["r0", "r1"])) close(compiled(functools.partial(rebuilt.run, label="main"), q, k, v), ref(q, k, v, cu([64, 160])) + 1.0)