From 0c899a6742c6801be368daf1dafceeef9dadfe51 Mon Sep 17 00:00:00 2001 From: merceod Date: Sat, 31 Oct 2026 02:35:53 +0000 Subject: [PATCH 1/3] share the AuT encoder node and the Qwen3 ViT blocks Qwen3-Omni is about to need both, so move them out of qwen3_asr and qwen3_5 into mstar/model/components. Code moves as is; the ViT attention takes its resource key as an argument. --- mstar/model/components/aut_submodule.py | 220 +++++++++++++++++ mstar/model/components/qwen3_vit.py | 286 +++++++++++++++++++++++ mstar/model/qwen3_5/components/vision.py | 261 +-------------------- mstar/model/qwen3_5/submodules.py | 18 +- mstar/model/qwen3_asr/submodules.py | 195 +--------------- test/modular/test_qwen3_asr_model.py | 2 +- 6 files changed, 527 insertions(+), 455 deletions(-) create mode 100644 mstar/model/components/aut_submodule.py create mode 100644 mstar/model/components/qwen3_vit.py diff --git a/mstar/model/components/aut_submodule.py b/mstar/model/components/aut_submodule.py new file mode 100644 index 0000000000..ee057a8421 --- /dev/null +++ b/mstar/model/components/aut_submodule.py @@ -0,0 +1,220 @@ +"""The shared AuT encoder node (Qwen3-ASR, Qwen3-Omni): requests packed into +one forward, every 8 s attention window of every request one segment of the +node's ragged-attention step, and the transformer stack captured by packed +token count. + +A model subclasses it for ``prepare_inputs`` only, which turns its request's +audio into ``(num_mel_bins, frames)`` features (``feature_inputs``). +""" +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +import torch.nn.functional as F + +from mstar.communication.tensors import NameToTensorList +from mstar.engine.accelerator_graph_config import ( + PiecewiseAcceleratorGraphConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, + PiecewisePackedConfig, +) +from mstar.engine.engine import ExecutingBatch +from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep +from mstar.model.components.aut_encoder import AuTEncoder, AuTEncoderConfig +from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule + +WINDOW_LABEL = "main" +# The AuT encoder's transformer stack and projector, as a piecewise capture region +AUT_ENCODE = "aut_encode" + + +class AuTEncoderSubmodule(NodeSubmodule): + """Packed AuT encoder, emitting ``audio_embeds`` of shape + ``(num_audio_tokens, output_dim)`` per request. The runner plans the + windows' kernel layout before the forward, so nothing is recomputed + inside it.""" + + MAX_BATCH_SIZE = 8 + # The engine's blanket torch.compile of ``forward_batched`` traces the + # per-request Python around the tensor work and guards on request ids + # and batch composition, so batches kept recompiling (1-2 s stalls each, + # measured on Whisper's decoder). The encoder's kernels are the ragged + # attention and cuBLAS GEMMs, there is nothing for inductor to fuse. + disable_torch_compile = True + + # Packed token counts the stack is captured at, and batch sizes (up to + # MAX_BATCH_SIZE). 4096 tokens is 8 requests of ~39 s. + ENCODE_TOKEN_BUCKETS = [64, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096] + ENCODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8] + + @classmethod + def max_windows_per_replay(cls, config: AuTEncoderConfig) -> int: + """Attention windows per request a captured replay can be handed, to + size the node's ragged resource: the largest bucket's tokens make at + most this many full windows, and each request adds one partial.""" + return cls.ENCODE_TOKEN_BUCKETS[-1] // config.window_tokens + 1 + + def __init__(self, encoder: AuTEncoder, attn_key: str): + super().__init__() + self.encoder = encoder + self.attn_key = attn_key + + def _param_dtype(self) -> torch.dtype: + return self.encoder.conv2d1.weight.dtype + + def feature_inputs(self, features: torch.Tensor) -> NodeInputs: + """This node's inputs for one request's ``(num_mel_bins, frames)`` + features, already on the encoder's device and in its dtype.""" + num_frames = int(features.shape[-1]) + return NodeInputs( + tensor_inputs={"audio_features": features}, + kwargs={"num_frames": num_frames}, + input_seq_len=self.encoder.config.tokens_for_frames(num_frames), + ) + + 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: + # leased, the captured stack plans its own attention per replay + if (piecewise_leases or {}).get(AUT_ENCODE): + return None + # one segment per attention window, in the packed order the forward + # lays the requests out in + segments = [ + Segment(request_id=rid, label=WINDOW_LABEL, span=window) + for rid, inp in zip(request_ids, inputs, strict=True) + for window in self.encoder.config.window_lengths(inp.input_seq_len) + ] + return SubmoduleStep(segments=segments, steps={self.attn_key: AttentionStep(causal=False)}) + + def preprocess( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + inputs: list[NodeInputs], + ) -> dict[str, torch.Tensor | Any]: + lens = [inp.kwargs["num_frames"] for inp in inputs] + longest = max(lens) + padded = torch.stack([ + F.pad(inp.tensor_inputs["audio_features"], (0, longest - n)) + for inp, n in zip(inputs, lens, strict=True) + ]) + return {"audio_features": padded, "feature_lens": lens} + + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, + tp_world_size: int = 1, **kwargs, + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + """The encoder's stack and projector, by batch size and packed token + count. At ASR clip lengths (13 tokens a second) its GEMMs are a few + dozen rows, so eager it is launch-bound: a 7 s clip's 24-layer stack + is ~6.3 ms eager and ~1.9 ms replayed. The conv frontend before it + runs eagerly. Batches past the largest bucket run eagerly too.""" + d_model = self.encoder.config.d_model + attn_key = self.attn_key + + def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return {"hidden": torch.zeros(shape.total_tokens, d_model, dtype=autocast_dtype, device=device)} + + def declare_step(request_ids: list[str], seq_lens: list[int]) -> SubmoduleStep: + # Every window of every request is a segment; the kernel needs + # only the boundaries, so they all name the first row's id. + # Zero-length rows pad the batch to the bucket. + return SubmoduleStep( + segments=[Segment(request_id=request_ids[0], label=WINDOW_LABEL, span=n) for n in seq_lens], + steps={attn_key: AttentionStep(causal=False)}, + ) + + return { + AUT_ENCODE: PiecewisePackedConfig( + capture_fn=self._capture_encode, + make_static_inputs=make_static_inputs, + declare_step=declare_step, + lease_before_step=True, + total_tokens=self.ENCODE_TOKEN_BUCKETS, + capture_batch_sizes=self.ENCODE_CAPTURE_BATCH_SIZES, + ) + } + + def _capture_encode(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + """The captured region: reads the runner's buffer, never rebinds.""" + return {"embeds": self.encoder.encode(inp.static_inputs["hidden"], None)} + + @torch.compiler.disable + def _replay_encode( + self, engine_inputs: ModelInputsFromEngine, hidden: torch.Tensor, windows: list[int], + ) -> torch.Tensor | None: + """The stack replayed for the whole packed batch, or None if not leased. + + A leased step declares nothing (`declare_step`), so ``engine_inputs.step`` + is None; a bucket that fits is not enough, as ``can_lease`` may have + refused the batch.""" + runner = engine_inputs.piecewise_runners.get(AUT_ENCODE) + bs = len(engine_inputs.request_ids) + if runner is None or engine_inputs.step is not None or not runner.can_run(bs, int(hidden.shape[0])): + return None + # an owned copy: the encoder states leave this worker after the step, + # and the buffer is the next replay's + return runner.run( + static_inputs={"hidden": hidden}, + request_ids=list(engine_inputs.request_ids), + seq_lens=windows, + real_bs=bs, + ).get("embeds") + + def _encode( + self, audio_features: torch.Tensor, feature_lens: list[int], + engine_inputs: ModelInputsFromEngine | None = None, + ) -> list[torch.Tensor]: + enc = self.encoder + layout = enc.layout(feature_lens) + hidden = enc.frontend(audio_features, feature_lens) + embeds = None + if engine_inputs is not None: + embeds = self._replay_encode(engine_inputs, hidden, layout.window_lengths) + if embeds is None: + cu = torch.tensor(layout.cu_seqlens, dtype=torch.int32, device=hidden.device) + embeds = enc.encode(hidden, cu) + return list(embeds.split(layout.tokens_per_request, dim=0)) + + def forward( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + audio_features: torch.Tensor, + feature_lens: list[int], + **kwargs, + ) -> NameToTensorList: + if audio_features.dim() == 2: + audio_features = audio_features.unsqueeze(0) + return {"audio_embeds": [self._encode(audio_features, feature_lens, engine_inputs)[0]]} + + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: + return True + + def max_batch_size(self, graph_walk: str) -> int | None: + del graph_walk + return self.MAX_BATCH_SIZE + + def forward_batched( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + audio_features: torch.Tensor, + feature_lens: list[int], + **kwargs, + ) -> dict[str, NameToTensorList]: + per_request = self._encode(audio_features, feature_lens, engine_inputs) + return { + rid: {"audio_embeds": [emb]} + for rid, emb in zip(engine_inputs.request_ids, per_request, strict=True) + } diff --git a/mstar/model/components/qwen3_vit.py b/mstar/model/components/qwen3_vit.py new file mode 100644 index 0000000000..ebe5e97b29 --- /dev/null +++ b/mstar/model/components/qwen3_vit.py @@ -0,0 +1,286 @@ +"""The Qwen3-VL family's ViT, the parts every tower shares: a ViT over packed +image patches with a learned position grid resampled to each image, axial 2D +rope, and attention within each frame through a ragged attention resource. + +Ported from transformers' ``modeling_qwen3_5.py`` and ``vision_utils.py`` +(Apache-2.0); the grid helpers are copied because ``transformers.vision_utils`` +is internal. Not here, as the towers differ in them: the patch mergers and +deepstack. Bicubic resampling and the temporal rotary axis are dropped, as no +tower using this sets them. + +Images arrive packed: ``pixel_values`` is ``[total_patches, patch_numel]`` and +``grid_thw`` gives each image's ``(t, h, w)`` patch grid. +""" +from __future__ import annotations + +from typing import Protocol + +import torch +import torch.nn.functional as F +from torch import nn + + +class ViTConfig(Protocol): + """What the shared modules read off a tower's config.""" + + hidden_size: int + intermediate_size: int + num_heads: int + in_channels: int + patch_size: int + temporal_patch_size: int + hidden_act: str + rope_theta: float + + @property + def head_dim(self) -> int: ... + + +# ---------------------------------------------------------------------------- +# Grid helpers +# ---------------------------------------------------------------------------- + + +def vision_seq_lengths(grid: list[tuple[int, int, int]]) -> tuple[int, ...]: + """One attending segment per frame, on the host: ``t`` segments of + ``h * w`` per entry, planned by the ragged attention resource.""" + return tuple(h * w for t, h, w in grid for _ in range(t)) + + +def vision_position_ids( + grid: list[tuple[int, int, int]], + spatial_merge_size: int, + device: torch.device, +) -> torch.Tensor: + """``[total_patches, 2]`` (h, w) indices in spatial-merge-block order. + + Block-major, so the rotary sees the order the merger will fold. + """ + out = [] + for t, h, w in grid: + hpos, wpos = torch.meshgrid( + torch.arange(h, device=device), + torch.arange(w, device=device), + indexing="ij", + ) + block = ( + h // spatial_merge_size, spatial_merge_size, + w // spatial_merge_size, spatial_merge_size, + ) + hpos = hpos.reshape(block).transpose(1, 2).flatten() + wpos = wpos.reshape(block).transpose(1, 2).flatten() + out.append(torch.stack([hpos, wpos], dim=-1).repeat(t, 1)) + return torch.cat(out, dim=0) + + +def _axis_taps_weights( + index: torch.Tensor, size: torch.Tensor, side: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Bilinear taps into a ``side``-long table, for positions on an axis of + length ``size``. ``align_corners=True`` and border padding, matching the + ``F.interpolate`` call upstream reproduces.""" + index = index.to(torch.float32) + # closed form of `linspace(0, side - 1, size)[index]`; clamp avoids a + # divide-by-zero at size == 1, where index is 0 anyway + src = index * (side - 1) / torch.clamp(size - 1, min=1) + floor = torch.floor(src) + offsets = torch.arange(0, 2, device=index.device) + taps = (floor.long()[:, None] + offsets).clamp(0, side - 1) + distance = (src[:, None] - floor[:, None] - offsets).abs() + return taps, (1 - distance).clamp(min=0) + + +def vision_interpolation( + grid: list[tuple[int, int, int]], + num_grid_per_side: int, + spatial_merge_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather indices and weights that resample the square learned position + table onto each image's ``(h, w)`` grid. + + Returns ``[total_patches, 4]`` each — the 2x2 bilinear neighbourhood per + patch, as the outer product of the two axes' taps. + """ + side, merge = num_grid_per_side, spatial_merge_size + + grid_thw = torch.tensor(grid, dtype=torch.long, device=device) + counts = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] + heights = torch.repeat_interleave(grid_thw[:, 1], counts) + widths = torch.repeat_interleave(grid_thw[:, 2], counts) + starts = torch.repeat_interleave( + F.pad(counts.cumsum(0)[:-1], (1, 0)), counts, + ) + # position within one frame's flat patch run, repeating across frames; + # the total comes off the host grid to avoid a device read + total_patches = sum(t * h * w for t, h, w in grid) + within = ( + torch.arange(total_patches, device=device) - starts + ) % (heights * widths) + + # undo the spatial-merge-block ordering to recover (row, col) + blocks_w = widths // merge + in_col = within % merge + in_row = (within // merge) % merge + block_col = (within // (merge * merge)) % blocks_w + block_row = within // (merge * merge * blocks_w) + row = block_row * merge + in_row + col = block_col * merge + in_col + + h_taps, h_w = _axis_taps_weights(row, heights, side) + w_taps, w_w = _axis_taps_weights(col, widths, side) + indices = (h_taps[:, :, None] * side + w_taps[:, None, :]).reshape(-1, 4) + weights = (h_w[:, :, None] * w_w[:, None, :]).reshape(-1, 4) + return indices, weights + + +# ---------------------------------------------------------------------------- +# Modules +# ---------------------------------------------------------------------------- + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +class VisionRotaryEmbedding(nn.Module): + """Axial 2D rope: the same frequencies drive the H and W axes, and the two + halves are concatenated to cover the whole head — no partial rotation.""" + + def __init__(self, config: ViTConfig): + super().__init__() + spatial_dim = config.head_dim // 2 + inv_freq = 1.0 / ( + config.rope_theta + ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward( + self, position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # position_ids: [N, 2] -> freqs [N, 2, dim/4] + freqs = position_ids[..., None].float() * self.inv_freq.float() + return self._recompose(freqs.cos()), self._recompose(freqs.sin()) + + @staticmethod + def _recompose(freq: torch.Tensor) -> torch.Tensor: + freq_hw = torch.cat([freq[:, 0], freq[:, 1]], dim=-1) + return torch.cat([freq_hw, freq_hw], dim=-1) + + +class VisionPatchEmbed(nn.Module): + def __init__(self, config: ViTConfig): + super().__init__() + self.in_channels = config.in_channels + self.temporal_patch_size = config.temporal_patch_size + self.patch_size = config.patch_size + self.embed_dim = config.hidden_size + kernel = (config.temporal_patch_size, config.patch_size, config.patch_size) + self.proj = nn.Conv3d( + config.in_channels, config.hidden_size, + kernel_size=kernel, stride=kernel, bias=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.view( + -1, self.in_channels, self.temporal_patch_size, + self.patch_size, self.patch_size, + ) + return self.proj(x.to(self.proj.weight.dtype)).view(-1, self.embed_dim) + + +class VisionMLP(nn.Module): + """Plain, not gated: one up projection, an activation, one down.""" + + def __init__(self, config: ViTConfig): + super().__init__() + self.linear_fc1 = nn.Linear( + config.hidden_size, config.intermediate_size, bias=True, + ) + self.linear_fc2 = nn.Linear( + config.intermediate_size, config.hidden_size, bias=True, + ) + if config.hidden_act not in ("gelu_pytorch_tanh", "gelu_new"): + raise NotImplementedError( + f"vision activation {config.hidden_act!r} is not wired up" + ) + self.act_fn = nn.GELU(approximate="tanh") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear_fc2(self.act_fn(self.linear_fc1(x))) + + +class VisionAttention(nn.Module): + """Full attention within each frame, over packed patches. + + The ragged attention resource plans the per-frame layout, so this forward + never sees the boundaries. That is what lets the tower compile: splitting + per frame here gave each frame its own symint, and inductor then refused + (``CantSplit``) to codegen packed runs of three or more images. + """ + + def __init__(self, config: ViTConfig, attn_key: str): + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=True) + self.proj = nn.Linear(config.hidden_size, config.hidden_size) + self.attn_key = attn_key + self.ragged_attn = None + + def bind_resources(self, resources: dict) -> None: + """See ``NodeSubmodule.bind_node_resources``. ``.get``: a deployment + that never builds the vision node binds nothing here.""" + self.ragged_attn = resources.get(self.attn_key) + + def forward( + self, + hidden_states: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + if self.ragged_attn is None: + raise RuntimeError( + "the vision tower has no ragged attention resource; declare a " + f"RaggedAttentionSpec under {self.attn_key!r} for its node" + ) + seq_len = hidden_states.shape[0] + # [tokens, heads, head_dim]: the ragged kernel's layout, no transposes + q, k, v = ( + self.qkv(hidden_states) + .reshape(seq_len, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) + # rope in fp32, as upstream does, then back + dtype = q.dtype + c, s = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + qf, kf = q.float(), k.float() + q = ((qf * c) + (_rotate_half(qf) * s)).to(dtype) + k = ((kf * c) + (_rotate_half(kf) * s)).to(dtype) + + out = self.ragged_attn.run(q, k, v) + return self.proj(out.reshape(seq_len, -1).contiguous()) + + +class VisionBlock(nn.Module): + def __init__(self, config: ViTConfig, attn_key: str): + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = VisionAttention(config, attn_key) + self.mlp = VisionMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), cos, sin, + ) + return hidden_states + self.mlp(self.norm2(hidden_states)) diff --git a/mstar/model/qwen3_5/components/vision.py b/mstar/model/qwen3_5/components/vision.py index 3aa0f7f7a6..4aadf54098 100644 --- a/mstar/model/qwen3_5/components/vision.py +++ b/mstar/model/qwen3_5/components/vision.py @@ -1,9 +1,5 @@ -"""Qwen3.5's vision tower: a ViT over packed image patches. - -Ported from transformers' ``modeling_qwen3_5.py`` and ``vision_utils.py`` -(Apache-2.0); the grid helpers are copied because ``transformers.vision_utils`` -is internal. Dropped, as Qwen3.5 does not use them: bicubic resampling, the -temporal rotary axis, and deepstack. +"""Qwen3.5's vision tower: the shared Qwen3-VL ViT (``components.qwen3_vit``) +with Qwen3.5's patch merger and no deepstack. Images arrive packed: ``pixel_values`` is ``[total_patches, patch_numel]`` and ``grid_thw`` gives each image's ``(t, h, w)`` patch grid. @@ -11,166 +7,11 @@ from __future__ import annotations import torch -import torch.nn.functional as F from torch import nn +from mstar.model.components.qwen3_vit import VisionBlock, VisionPatchEmbed, VisionRotaryEmbedding from mstar.model.qwen3_5.config import VISION_ATTN, Qwen3_5VisionConfig -# ---------------------------------------------------------------------------- -# Grid helpers -# ---------------------------------------------------------------------------- - - -def vision_seq_lengths(grid: list[tuple[int, int, int]]) -> tuple[int, ...]: - """One attending segment per frame, on the host: ``t`` segments of - ``h * w`` per entry, planned by the ragged attention resource.""" - return tuple(h * w for t, h, w in grid for _ in range(t)) - - -def vision_position_ids( - grid: list[tuple[int, int, int]], - spatial_merge_size: int, - device: torch.device, -) -> torch.Tensor: - """``[total_patches, 2]`` (h, w) indices in spatial-merge-block order. - - Block-major, so the rotary sees the order the merger will fold. - """ - out = [] - for t, h, w in grid: - hpos, wpos = torch.meshgrid( - torch.arange(h, device=device), - torch.arange(w, device=device), - indexing="ij", - ) - block = ( - h // spatial_merge_size, spatial_merge_size, - w // spatial_merge_size, spatial_merge_size, - ) - hpos = hpos.reshape(block).transpose(1, 2).flatten() - wpos = wpos.reshape(block).transpose(1, 2).flatten() - out.append(torch.stack([hpos, wpos], dim=-1).repeat(t, 1)) - return torch.cat(out, dim=0) - - -def _axis_taps_weights( - index: torch.Tensor, size: torch.Tensor, side: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Bilinear taps into a ``side``-long table, for positions on an axis of - length ``size``. ``align_corners=True`` and border padding, matching the - ``F.interpolate`` call upstream reproduces.""" - index = index.to(torch.float32) - # closed form of `linspace(0, side - 1, size)[index]`; clamp avoids a - # divide-by-zero at size == 1, where index is 0 anyway - src = index * (side - 1) / torch.clamp(size - 1, min=1) - floor = torch.floor(src) - offsets = torch.arange(0, 2, device=index.device) - taps = (floor.long()[:, None] + offsets).clamp(0, side - 1) - distance = (src[:, None] - floor[:, None] - offsets).abs() - return taps, (1 - distance).clamp(min=0) - - -def vision_interpolation( - grid: list[tuple[int, int, int]], - num_grid_per_side: int, - spatial_merge_size: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - """Gather indices and weights that resample the square learned position - table onto each image's ``(h, w)`` grid. - - Returns ``[total_patches, 4]`` each — the 2x2 bilinear neighbourhood per - patch, as the outer product of the two axes' taps. - """ - side, merge = num_grid_per_side, spatial_merge_size - - grid_thw = torch.tensor(grid, dtype=torch.long, device=device) - counts = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] - heights = torch.repeat_interleave(grid_thw[:, 1], counts) - widths = torch.repeat_interleave(grid_thw[:, 2], counts) - starts = torch.repeat_interleave( - F.pad(counts.cumsum(0)[:-1], (1, 0)), counts, - ) - # position within one frame's flat patch run, repeating across frames; - # the total comes off the host grid to avoid a device read - total_patches = sum(t * h * w for t, h, w in grid) - within = ( - torch.arange(total_patches, device=device) - starts - ) % (heights * widths) - - # undo the spatial-merge-block ordering to recover (row, col) - blocks_w = widths // merge - in_col = within % merge - in_row = (within // merge) % merge - block_col = (within // (merge * merge)) % blocks_w - block_row = within // (merge * merge * blocks_w) - row = block_row * merge + in_row - col = block_col * merge + in_col - - h_taps, h_w = _axis_taps_weights(row, heights, side) - w_taps, w_w = _axis_taps_weights(col, widths, side) - indices = (h_taps[:, :, None] * side + w_taps[:, None, :]).reshape(-1, 4) - weights = (h_w[:, :, None] * w_w[:, None, :]).reshape(-1, 4) - return indices, weights - - -# ---------------------------------------------------------------------------- -# Modules -# ---------------------------------------------------------------------------- - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - half = x.shape[-1] // 2 - return torch.cat((-x[..., half:], x[..., :half]), dim=-1) - - -class VisionRotaryEmbedding(nn.Module): - """Axial 2D rope: the same frequencies drive the H and W axes, and the two - halves are concatenated to cover the whole head — no partial rotation.""" - - def __init__(self, config: Qwen3_5VisionConfig): - super().__init__() - spatial_dim = config.head_dim // 2 - inv_freq = 1.0 / ( - config.rope_theta - ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim) - ) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - @torch.no_grad() - def forward( - self, position_ids: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # position_ids: [N, 2] -> freqs [N, 2, dim/4] - freqs = position_ids[..., None].float() * self.inv_freq.float() - return self._recompose(freqs.cos()), self._recompose(freqs.sin()) - - @staticmethod - def _recompose(freq: torch.Tensor) -> torch.Tensor: - freq_hw = torch.cat([freq[:, 0], freq[:, 1]], dim=-1) - return torch.cat([freq_hw, freq_hw], dim=-1) - - -class VisionPatchEmbed(nn.Module): - def __init__(self, config: Qwen3_5VisionConfig): - super().__init__() - self.in_channels = config.in_channels - self.temporal_patch_size = config.temporal_patch_size - self.patch_size = config.patch_size - self.embed_dim = config.hidden_size - kernel = (config.temporal_patch_size, config.patch_size, config.patch_size) - self.proj = nn.Conv3d( - config.in_channels, config.hidden_size, - kernel_size=kernel, stride=kernel, bias=True, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x.view( - -1, self.in_channels, self.temporal_patch_size, - self.patch_size, self.patch_size, - ) - return self.proj(x.to(self.proj.weight.dtype)).view(-1, self.embed_dim) - class VisionPatchMerger(nn.Module): """Folds each ``m x m`` block of patches into one LLM-width token.""" @@ -188,100 +29,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear_fc2(self.act_fn(self.linear_fc1(x))) -class VisionMLP(nn.Module): - """Plain, not gated: one up projection, an activation, one down.""" - - def __init__(self, config: Qwen3_5VisionConfig): - super().__init__() - self.linear_fc1 = nn.Linear( - config.hidden_size, config.intermediate_size, bias=True, - ) - self.linear_fc2 = nn.Linear( - config.intermediate_size, config.hidden_size, bias=True, - ) - if config.hidden_act not in ("gelu_pytorch_tanh", "gelu_new"): - raise NotImplementedError( - f"vision activation {config.hidden_act!r} is not wired up" - ) - self.act_fn = nn.GELU(approximate="tanh") - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.linear_fc2(self.act_fn(self.linear_fc1(x))) - - -class VisionAttention(nn.Module): - """Full attention within each frame, over packed patches. - - The ragged attention resource plans the per-frame layout, so this forward - never sees the boundaries. That is what lets the tower compile: splitting - per frame here gave each frame its own symint, and inductor then refused - (``CantSplit``) to codegen packed runs of three or more images. - """ - - def __init__(self, config: Qwen3_5VisionConfig): - super().__init__() - self.num_heads = config.num_heads - self.head_dim = config.head_dim - self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=True) - self.proj = nn.Linear(config.hidden_size, config.hidden_size) - self.ragged_attn = None - - def bind_resources(self, resources: dict) -> None: - """See ``NodeSubmodule.bind_node_resources``. ``.get``: a deployment - that never builds the vision node binds nothing here.""" - self.ragged_attn = resources.get(VISION_ATTN) - - def forward( - self, - hidden_states: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - ) -> torch.Tensor: - if self.ragged_attn is None: - raise RuntimeError( - "Qwen3.5's vision tower has no ragged attention resource; " - f"declare a RaggedAttentionSpec under {VISION_ATTN!r} for the " - "vision_encoder node" - ) - seq_len = hidden_states.shape[0] - # [tokens, heads, head_dim]: the ragged kernel's layout, no transposes - q, k, v = ( - self.qkv(hidden_states) - .reshape(seq_len, 3, self.num_heads, -1) - .permute(1, 0, 2, 3) - .unbind(0) - ) - # rope in fp32, as upstream does, then back - dtype = q.dtype - c, s = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() - qf, kf = q.float(), k.float() - q = ((qf * c) + (_rotate_half(qf) * s)).to(dtype) - k = ((kf * c) + (_rotate_half(kf) * s)).to(dtype) - - out = self.ragged_attn.run(q, k, v) - return self.proj(out.reshape(seq_len, -1).contiguous()) - - -class VisionBlock(nn.Module): - def __init__(self, config: Qwen3_5VisionConfig): - super().__init__() - self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) - self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) - self.attn = VisionAttention(config) - self.mlp = VisionMLP(config) - - def forward( - self, - hidden_states: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - ) -> torch.Tensor: - hidden_states = hidden_states + self.attn( - self.norm1(hidden_states), cos, sin, - ) - return hidden_states + self.mlp(self.norm2(hidden_states)) - - class Qwen3_5VisionModel(nn.Module): def __init__(self, config: Qwen3_5VisionConfig): super().__init__() @@ -298,7 +45,7 @@ def __init__(self, config: Qwen3_5VisionConfig): ) self.rotary_pos_emb = VisionRotaryEmbedding(config) self.blocks = nn.ModuleList( - VisionBlock(config) for _ in range(config.depth) + VisionBlock(config, VISION_ATTN) for _ in range(config.depth) ) self.merger = VisionPatchMerger(config) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 584895923d..670a01c89d 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -37,21 +37,21 @@ from mstar.engine.resources.sampler.resource import SamplerResource from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep from mstar.model.components.mrope import mrope_position_ids -from mstar.model.qwen3_5.components.language_model import Qwen3_5ForCausalLM -from mstar.model.qwen3_5.components.rope import ( - text_position_ids, - vision_position_advance, - vision_position_ids, -) -from mstar.model.qwen3_5.components.vision import ( - Qwen3_5VisionModel, +from mstar.model.components.qwen3_vit import ( vision_interpolation, vision_seq_lengths, ) -from mstar.model.qwen3_5.components.vision import ( +from mstar.model.components.qwen3_vit import ( # the tower's 2D patch grid, not `rope`'s 3D MRoPE ids of the same name vision_position_ids as vision_grid_position_ids, ) +from mstar.model.qwen3_5.components.language_model import Qwen3_5ForCausalLM +from mstar.model.qwen3_5.components.rope import ( + text_position_ids, + vision_position_advance, + vision_position_ids, +) +from mstar.model.qwen3_5.components.vision import Qwen3_5VisionModel from mstar.model.qwen3_5.config import ( ATTN, GDN_STATE, diff --git a/mstar/model/qwen3_asr/submodules.py b/mstar/model/qwen3_asr/submodules.py index d0dd3c4b97..a03c0b2d5f 100644 --- a/mstar/model/qwen3_asr/submodules.py +++ b/mstar/model/qwen3_asr/submodules.py @@ -16,24 +16,19 @@ from typing import Any import torch -import torch.nn.functional as F from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.accelerator_graph_config import ( AcceleratorGraphConfig, PackedAcceleratorGraphConfig, - PiecewiseAcceleratorGraphConfig, - PiecewiseCallInputs, - PiecewiseCaptureShape, - PiecewisePackedConfig, ) -from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep from mstar.engine.resources.attn.base import AttentionManager from mstar.engine.resources.sampler.resource import SamplerResource from mstar.model.components.audio_features import LogMelSpectrogram from mstar.model.components.aut_encoder import AuTEncoder +from mstar.model.components.aut_submodule import AuTEncoderSubmodule from mstar.model.components.qwen3_lm import Qwen3DenseLM from mstar.model.components.token_step import TokenStepNode from mstar.model.qwen3_asr.config import ( @@ -53,45 +48,26 @@ BatchedModelOutput, ModelInputsFromEngine, NodeInputs, - NodeSubmodule, device_loopback_enabled, ) logger = logging.getLogger(__name__) -WINDOW_LABEL = "main" -# The AuT encoder's transformer stack and projector, as a piecewise capture region -AUT_ENCODE = "aut_encode" - # =================================================================== # 1. Qwen3ASREncoderSubmodule # =================================================================== -class Qwen3ASREncoderSubmodule(NodeSubmodule): - """Packed AuT encoder. +class Qwen3ASREncoderSubmodule(AuTEncoderSubmodule): + """Packed AuT encoder (see ``AuTEncoderSubmodule``). - Consumes one clip of samples per request (from ``process_prompt``), - turns it into unpadded log-mel features on the GPU (``prepare_inputs``) - and emits ``audio_embeds`` of shape - ``(num_audio_tokens, hidden)`` in LLM space. Requests are packed into - one forward; every 8 s attention window of every request is one - segment of the node's ragged-attention step, so the runner plans the - kernel's layout before the forward and nothing is recomputed inside. + Consumes one clip of samples per request (from ``process_prompt``) and + turns it into unpadded log-mel features on the GPU (``prepare_inputs``). """ - MAX_BATCH_SIZE = 8 - # The engine's blanket torch.compile of ``forward_batched`` traces the - # per-request Python around the tensor work and guards on request ids - # and batch composition, so batches kept recompiling (1-2 s stalls each, - # measured on Whisper's decoder). The encoder's kernels are the ragged - # attention and cuBLAS GEMMs, there is nothing for inductor to fuse. - disable_torch_compile = True - def __init__(self, encoder: AuTEncoder, config: Qwen3ASRModelConfig): - super().__init__() - self.encoder = encoder + super().__init__(encoder, attn_key=AUT_ATTN) self.config = config # built here, after the encoder was materialized, so its filter bank # and window are real tensors on the encoder's device @@ -110,9 +86,6 @@ def _mel(self, audio: torch.Tensor) -> torch.Tensor: self.log_mel.float() return self.log_mel(audio) - def _param_dtype(self) -> torch.dtype: - return self.encoder.conv2d1.weight.dtype - def prepare_inputs( self, graph_walk: str, @@ -121,161 +94,7 @@ def prepare_inputs( **kwargs, ) -> NodeInputs: audio = inputs["audio"][0].reshape(-1).to(device=self.get_device(), dtype=torch.float32) - feats = self._mel(audio).to(self._param_dtype()) # (num_mel_bins, T) - num_frames = int(feats.shape[-1]) - return NodeInputs( - tensor_inputs={"audio_features": feats}, - kwargs={"num_frames": num_frames}, - input_seq_len=self.encoder.config.tokens_for_frames(num_frames), - ) - - 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: - # leased, the captured stack plans its own attention per replay - if (piecewise_leases or {}).get(AUT_ENCODE): - return None - # one segment per attention window, in the packed order the forward - # lays the requests out in - segments = [ - Segment(request_id=rid, label=WINDOW_LABEL, span=window) - for rid, inp in zip(request_ids, inputs, strict=True) - for window in self.encoder.config.window_lengths(inp.input_seq_len) - ] - return SubmoduleStep(segments=segments, steps={AUT_ATTN: AttentionStep(causal=False)}) - - def preprocess( - self, - graph_walk: str, - engine_inputs: ModelInputsFromEngine, - inputs: list[NodeInputs], - ) -> dict[str, torch.Tensor | Any]: - lens = [inp.kwargs["num_frames"] for inp in inputs] - longest = max(lens) - padded = torch.stack([ - F.pad(inp.tensor_inputs["audio_features"], (0, longest - n)) - for inp, n in zip(inputs, lens, strict=True) - ]) - return {"audio_features": padded, "feature_lens": lens} - - def get_piecewise_accelerator_graph_configs( - self, device: torch.device, autocast_dtype: torch.dtype, - tp_world_size: int = 1, **kwargs, - ) -> dict[str, PiecewiseAcceleratorGraphConfig]: - """The encoder's 24-layer stack and projector, by batch size and packed - token count. At ASR clip lengths (13 tokens a second) its GEMMs are a - few dozen rows, so eager it is launch-bound: a 7 s clip's stack is - ~6.3 ms eager and ~1.9 ms replayed. The conv frontend before it runs - eagerly. Batches past the largest bucket run eagerly too.""" - d_model = self.encoder.config.d_model - - def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: - return {"hidden": torch.zeros(shape.total_tokens, d_model, dtype=autocast_dtype, device=device)} - - def declare_step(request_ids: list[str], seq_lens: list[int]) -> SubmoduleStep: - # Every window of every request is a segment; the kernel needs - # only the boundaries, so they all name the first row's id. - # Zero-length rows pad the batch to the bucket. - return SubmoduleStep( - segments=[Segment(request_id=request_ids[0], label=WINDOW_LABEL, span=n) for n in seq_lens], - steps={AUT_ATTN: AttentionStep(causal=False)}, - ) - - return { - AUT_ENCODE: PiecewisePackedConfig( - capture_fn=self._capture_encode, - make_static_inputs=make_static_inputs, - declare_step=declare_step, - lease_before_step=True, - total_tokens=self.ENCODE_TOKEN_BUCKETS, - capture_batch_sizes=self.ENCODE_CAPTURE_BATCH_SIZES, - ) - } - - # Packed token counts the stack is captured at, and batch sizes (up to - # MAX_BATCH_SIZE). 4096 tokens is 8 requests of ~39 s. - ENCODE_TOKEN_BUCKETS = [64, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096] - ENCODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8] - - def _capture_encode(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]: - """The captured region: reads the runner's buffer, never rebinds.""" - return {"embeds": self.encoder.encode(inp.static_inputs["hidden"], None)} - - @torch.compiler.disable - def _replay_encode( - self, engine_inputs: ModelInputsFromEngine, hidden: torch.Tensor, windows: list[int], - ) -> torch.Tensor | None: - """The stack replayed for the whole packed batch, or None if not leased. - - A leased step declares nothing (`declare_step`), so ``engine_inputs.step`` - is None; a bucket that fits is not enough, as ``can_lease`` may have - refused the batch.""" - runner = engine_inputs.piecewise_runners.get(AUT_ENCODE) - bs = len(engine_inputs.request_ids) - if runner is None or engine_inputs.step is not None or not runner.can_run(bs, int(hidden.shape[0])): - return None - # an owned copy: the encoder states leave this worker after the step, - # and the buffer is the next replay's - return runner.run( - static_inputs={"hidden": hidden}, - request_ids=list(engine_inputs.request_ids), - seq_lens=windows, - real_bs=bs, - ).get("embeds") - - def _encode( - self, audio_features: torch.Tensor, feature_lens: list[int], - engine_inputs: ModelInputsFromEngine | None = None, - ) -> list[torch.Tensor]: - enc = self.encoder - layout = enc.layout(feature_lens) - hidden = enc.frontend(audio_features, feature_lens) - embeds = None - if engine_inputs is not None: - embeds = self._replay_encode(engine_inputs, hidden, layout.window_lengths) - if embeds is None: - cu = torch.tensor(layout.cu_seqlens, dtype=torch.int32, device=hidden.device) - embeds = enc.encode(hidden, cu) - return list(embeds.split(layout.tokens_per_request, dim=0)) - - def forward( - self, - graph_walk: str, - engine_inputs: ModelInputsFromEngine, - audio_features: torch.Tensor, - feature_lens: list[int], - **kwargs, - ) -> NameToTensorList: - if audio_features.dim() == 2: - audio_features = audio_features.unsqueeze(0) - return {"audio_embeds": [self._encode(audio_features, feature_lens, engine_inputs)[0]]} - - def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: - return True - - def max_batch_size(self, graph_walk: str) -> int | None: - del graph_walk - return self.MAX_BATCH_SIZE - - def forward_batched( - self, - graph_walk: str, - engine_inputs: ModelInputsFromEngine, - audio_features: torch.Tensor, - feature_lens: list[int], - **kwargs, - ) -> dict[str, NameToTensorList]: - per_request = self._encode(audio_features, feature_lens, engine_inputs) - return { - rid: {"audio_embeds": [emb]} - for rid, emb in zip(engine_inputs.request_ids, per_request, strict=True) - } + return self.feature_inputs(self._mel(audio).to(self._param_dtype())) # (num_mel_bins, T) # =================================================================== diff --git a/test/modular/test_qwen3_asr_model.py b/test/modular/test_qwen3_asr_model.py index 0d2964c75b..8d111d713d 100644 --- a/test/modular/test_qwen3_asr_model.py +++ b/test/modular/test_qwen3_asr_model.py @@ -311,7 +311,7 @@ def test_encoder_replays_exactly_when_leased(): then and runs eagerly otherwise, whatever the runner's buckets allow.""" from types import SimpleNamespace - from mstar.model.qwen3_asr.submodules import AUT_ENCODE + from mstar.model.components.aut_submodule import AUT_ENCODE class Runner: has_bucket = True From 87f3d9fb93ac1a2ee87147394d4e090162de5011 Mon Sep 17 00:00:00 2001 From: merceod Date: Sat, 31 Oct 2026 02:35:54 +0000 Subject: [PATCH 2/3] qwen3-omni: run the encoders on ragged attention, drop flash_attn The HF audio and vision encoders were built with FlashAttention-2, which has no sm_103 kernels and isn't installed on Blackwell boxes, so the server couldn't start. Use the native AuT encoder and a ViT tower with deepstack, both attending through declared ragged resources, with their block stacks captured. --- README.md | 2 +- benchmark/eval/e2e/cases.yaml | 16 +- docs/installation.rst | 18 +- docs/models.rst | 7 + mstar/model/components/aut_submodule.py | 4 +- .../model/qwen3_omni/components/attention.py | 34 -- mstar/model/qwen3_omni/components/vision.py | 120 ++++++ mstar/model/qwen3_omni/config.py | 49 ++- mstar/model/qwen3_omni/qwen3_omni_model.py | 159 ++++--- mstar/model/qwen3_omni/submodules.py | 395 ++++++++++-------- pyproject.toml | 3 - test/modular/test_qwen3_omni_encoders.py | 240 +++++++++++ test/modular/test_qwen3_omni_prefix_cache.py | 141 +------ test/modular/test_trust_remote_code.py | 15 - 14 files changed, 717 insertions(+), 486 deletions(-) create mode 100644 mstar/model/qwen3_omni/components/vision.py create mode 100644 test/modular/test_qwen3_omni_encoders.py diff --git a/README.md b/README.md index 6fb16fd59e..67f7b11fde 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ The PyPI package is `mstar-ai` (`mstar-project` and `mstar-serve` install the sa and the CLI are `mstar`. Working from a checkout instead: `uv pip install --torch-backend=auto -e .[all]`. Qwen3-TTS also needs `uv pip install --no-deps "qwen-tts==0.1.1"` (see `docs/installation.rst`). -To enable flash-attn support (required for Qwen3-Omni, recommended for BAGEL), +To enable flash-attn support (recommended for BAGEL), ```bash # torch 2.10 built for CUDA 13.x (cu13); on a cu12 torch, build from source (docs/installation.rst) uv pip install \ diff --git a/benchmark/eval/e2e/cases.yaml b/benchmark/eval/e2e/cases.yaml index 3b6ec79b86..26d875f287 100644 --- a/benchmark/eval/e2e/cases.yaml +++ b/benchmark/eval/e2e/cases.yaml @@ -221,7 +221,7 @@ models: config: configs/qwen3omni_colocated.yaml family: Omni tags: [omni, llm, text, image-in, audio-in, video-in, speech-out, sse, big] - requires: [flash_attn, qwen_omni_utils, audioread] + requires: [qwen_omni_utils, audioread] est_min: 16 cases: - id: text @@ -313,7 +313,7 @@ models: env: {MSTAR_ALLOW_REMOTE: "1"} family: Omni tags: [omni, rust-frontend, image-in, sse, big] - requires: [flash_attn, qwen_omni_utils] + requires: [qwen_omni_utils] est_min: 14 cases: - id: text @@ -458,7 +458,7 @@ models: optional: true family: Omni tags: [mooncake, transport, omni, image-in, speech-out, big] - requires: [flash_attn, qwen_omni_utils, mooncake.engine] + requires: [qwen_omni_utils, mooncake.engine] est_min: 10 cases: - id: text @@ -540,7 +540,7 @@ models: case_timeout_s: 900 family: Omni tags: [multi-gpu, tp, omni, image-in, speech-out, big] - requires: [flash_attn, qwen_omni_utils, audioread] + requires: [qwen_omni_utils, audioread] est_min: 16 conc: [1, 16] cases: @@ -634,7 +634,7 @@ models: case_timeout_s: 900 family: Omni tags: [multi-gpu, disaggregated, omni, speech-out, big] - requires: [flash_attn, qwen_omni_utils, audioread] + requires: [qwen_omni_utils, audioread] est_min: 16 cases: - id: text @@ -879,7 +879,7 @@ models: health: true family: Omni tags: [w2w, transport, shm, disaggregated, omni, speech-out, big] - requires: [flash_attn, qwen_omni_utils, audioread] + requires: [qwen_omni_utils, audioread] est_min: 18 cases: &qwen3omni_pd_w2w - id: text @@ -917,7 +917,7 @@ models: health: true family: Omni tags: [w2w, transport, disaggregated, omni, speech-out, big] - requires: [flash_attn, qwen_omni_utils, audioread, mooncake.engine] + requires: [qwen_omni_utils, audioread, mooncake.engine] est_min: 18 cases: *qwen3omni_pd_w2w @@ -931,7 +931,7 @@ models: health: true family: Omni tags: [rdma, disaggregated, omni, speech-out, big] - requires: [flash_attn, qwen_omni_utils, audioread, mooncake.engine, rdma_nic] + requires: [qwen_omni_utils, audioread, mooncake.engine, rdma_nic] est_min: 18 cases: *qwen3omni_pd_w2w diff --git a/docs/installation.rst b/docs/installation.rst index c3ecbd9b6e..32d7a7dab8 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -135,9 +135,8 @@ Model families and some output formats need extra packages, exposed as pip *extr * - ``.[qwen3_omni]`` - Qwen3-Omni runtime: the BAGEL set plus ``qwen-omni-utils``, ``datasets``, and ``ninja`` (speeds up the JIT build of the vendored MoE align kernel). It needs - ``transformers>=5.18``, whose ``transformers.vision_utils`` the vision encoder uses. - **Also needs** ``flash-attn``, which is installed separately — - see `flash-attn (Qwen3-Omni, LingBot-World)`_. + ``transformers>=5.18``. Its encoders attend through FlashInfer, so it needs no + ``flash-attn``. * - ``.[orpheus]`` - Orpheus TTS runtime: ``transformers``, ``flashinfer-python``, ``safetensors``, ``einops``, ``huggingface-hub``, ``mooncake-transfer-engine``. @@ -191,7 +190,7 @@ Model families and some output formats need extra packages, exposed as pip *extr Convenient for a machine that serves multiple models; heavier and slower to install than a single family's extra. It excludes the separately installed TAEHV, ``flash-attn`` and ``qwen-tts`` packages; see below, - `flash-attn (Qwen3-Omni, LingBot-World)`_ and `qwen-tts (Qwen3-TTS)`_. + `flash-attn (LingBot-World)`_ and `qwen-tts (Qwen3-TTS)`_. Combine extras as needed (keep ``--torch-backend=auto`` on every install): @@ -225,8 +224,7 @@ Or, in an existing Python 3.12 environment: If you're just getting started or have the disk/time to spare, ``.[all]`` installs all index-hosted model dependencies in one shot. Waypoint still needs the pinned TAEHV - command above, Qwen3-Omni still needs ``flash-attn``, and Qwen3-TTS still needs - ``qwen-tts`` (see `qwen-tts (Qwen3-TTS)`_): + command above, and Qwen3-TTS still needs ``qwen-tts`` (see `qwen-tts (Qwen3-TTS)`_): .. code-block:: bash @@ -246,8 +244,8 @@ The GPU model families depend on: - **FlashInfer** (``flashinfer-python``) — paged attention and continuous batching for the autoregressive backbones (every model with a ``KV_CACHE`` node runs attention through it). -- **flash-attn** — used by Qwen3-Omni and by LingBot-World's ``reference_numerics`` mode. **Not installed by any extra**; install it separately - (see `flash-attn (Qwen3-Omni, LingBot-World)`_). +- **flash-attn** — used by LingBot-World's ``reference_numerics`` mode. **Not installed by any extra**; install it separately + (see `flash-attn (LingBot-World)`_). - **flash-attn-4** — Waypoint's flex-attention ``FLASH`` backend. **Not installed by any extra**; install it separately (see `flash-attn-4 (Waypoint, FA3 sm90 kernel)`_). - **mooncake-transfer-engine** (0.3.13.post1 or newer) — RDMA tensor transport for @@ -307,8 +305,8 @@ others) as unmet for ``qwen-tts``; that is expected. flash-attn (Qwen3-Omni, LingBot-World) -------------------------------------- -``flash-attn`` is only needed for **Qwen3-Omni** and **LingBot-World** with ``reference_numerics: true``, and it is **not** pulled in by -``.[qwen3_omni]``, ``.[lingbot_world]`` or ``.[all]`` — you install it as a separate step. The reason: flash-attn +``flash-attn`` is only needed for **LingBot-World** with ``reference_numerics: true``, and it is **not** pulled in by +``.[lingbot_world]`` or ``.[all]`` — you install it as a separate step. The reason: flash-attn publishes no wheels on PyPI, so ``pip``/``uv`` fall back to compiling it from source, which is slow. When a prebuilt GitHub wheel matches your stack, use it to skip the build; otherwise build from source (see `No matching wheel — build from source`_ below). diff --git a/docs/models.rst b/docs/models.rst index 062da7be90..3a0c722c60 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -478,6 +478,13 @@ Kokoro notes Qwen3-Omni notes ---------------- +- The audio and vision encoders attend through ragged attention resources + (``audio_attn`` within each 8 s window, ``vision_attn`` within each image or + video frame group), FlashInfer by default, so no ``flash-attn`` install is + needed. A deployment can pick the kernel under ``resources:`` in its YAML, + e.g. ``vision_attn: {flashinfer_backend: cutlass}``. Each encoder's block stack + is captured as a CUDA graph per token-count bucket; a request past the + largest bucket (32768 patches, about 5 minutes of audio) runs it eagerly. - Streamed audio starts after 12 codec frames (0.96 s of audio), then 13, then 25-frame chunks behind 25 frames of left context. Playback stays gapless while the Talker produces frames at about 1.5x real time over the first 50 diff --git a/mstar/model/components/aut_submodule.py b/mstar/model/components/aut_submodule.py index ee057a8421..50d77f9f06 100644 --- a/mstar/model/components/aut_submodule.py +++ b/mstar/model/components/aut_submodule.py @@ -32,8 +32,8 @@ class AuTEncoderSubmodule(NodeSubmodule): - """Packed AuT encoder, emitting ``audio_embeds`` of shape - ``(num_audio_tokens, output_dim)`` per request. The runner plans the + """Packed AuT encoder, emitting one ``(num_audio_tokens, output_dim)`` + ``audio_embeds`` tensor per request. The runner plans the windows' kernel layout before the forward, so nothing is recomputed inside it.""" diff --git a/mstar/model/qwen3_omni/components/attention.py b/mstar/model/qwen3_omni/components/attention.py index 9248c047c5..e6683cbaa6 100644 --- a/mstar/model/qwen3_omni/components/attention.py +++ b/mstar/model/qwen3_omni/components/attention.py @@ -22,40 +22,6 @@ from mstar.distributed.communication import CommGroup from mstar.model.components.distributed import ParallelAttention -_fa2_maxlen_patched = False - - -def patch_hf_fa2_int_maxlen() -> None: - """Make HF's FA2 encoder path torch.compile-able. - - ``Qwen3OmniMoeAudioAttention`` / ``Qwen3OmniMoeVisionAttention`` compute - ``max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()`` and hand that - 0-dim tensor to ``flash_attn_varlen_func``, whose schema declares - ``max_seqlen_q/k`` as ``SymInt``. Eager tolerates it (the arg parser - unwraps a 0-dim tensor); dynamo's fake-tensor propagation does not, and - the encoder forward dies with "Expected a value of type 'int'". - - Wrap the registered ``flash_attention_2`` interface so those two kwargs - arrive as ints. Idempotent, and a no-op for callers already passing ints. - """ - global _fa2_maxlen_patched - if _fa2_maxlen_patched: - return - - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS - - inner = ALL_ATTENTION_FUNCTIONS["flash_attention_2"] - - def fa2_int_maxlen(*args, max_length_q=None, max_length_k=None, **kwargs): - if isinstance(max_length_q, torch.Tensor): - max_length_q = max_length_q.item() - if isinstance(max_length_k, torch.Tensor): - max_length_k = max_length_k.item() - return inner(*args, max_length_q=max_length_q, max_length_k=max_length_k, **kwargs) - - ALL_ATTENTION_FUNCTIONS.register("flash_attention_2", fa2_int_maxlen) - _fa2_maxlen_patched = True - class Qwen3OmniAttention(ParallelAttention): """TP-aware attention with QK-norm and pluggable 1D / 3D RoPE. diff --git a/mstar/model/qwen3_omni/components/vision.py b/mstar/model/qwen3_omni/components/vision.py new file mode 100644 index 0000000000..d2b28eceb8 --- /dev/null +++ b/mstar/model/qwen3_omni/components/vision.py @@ -0,0 +1,120 @@ +"""Qwen3-Omni's vision tower: the shared Qwen3-VL ViT (``components.qwen3_vit``) +plus deepstack, i.e. the hidden states after a few chosen blocks, each folded +by its own patch merger, which the Thinker adds to its first layers' vision +rows. + +Parameter names follow the checkpoint's ``thinker.visual.*`` with the prefix +stripped, so it loads by name: ``merger`` and ``merger_list.{i}`` are +``Qwen3OmniMoeVisionPatchMerger`` (``ln_q``, ``mlp.0``, ``mlp.2``). +""" +from __future__ import annotations + +import torch +from torch import nn + +from mstar.model.components.qwen3_vit import VisionBlock, VisionPatchEmbed, VisionRotaryEmbedding +from mstar.model.qwen3_omni.config import VisionEncoderConfig + + +class OmniPatchMerger(nn.Module): + """Folds each ``m x m`` block of patches into one Thinker-width token. + + The main merger normalizes each patch before folding; the deepstack ones + (``postshuffle``) normalize the folded vector. + """ + + def __init__(self, config: VisionEncoderConfig, postshuffle: bool): + super().__init__() + self.merged_size = config.hidden_size * config.merge_unit + self.postshuffle = postshuffle + self.ln_q = nn.LayerNorm(self.merged_size if postshuffle else config.hidden_size, eps=1e-6) + self.mlp = nn.ModuleList([ + nn.Linear(self.merged_size, self.merged_size), + nn.GELU(), + nn.Linear(self.merged_size, config.out_hidden_size), + ]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.view(-1, self.merged_size) if self.postshuffle else x + x = self.ln_q(x).view(-1, self.merged_size) + for layer in self.mlp: + x = layer(x) + return x + + +class Qwen3OmniVisionTower(nn.Module): + def __init__(self, config: VisionEncoderConfig, attn_key: str): + super().__init__() + self.config = config + self.patch_embed = VisionPatchEmbed(config) + self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size) + self.rotary_pos_emb = VisionRotaryEmbedding(config) + self.blocks = nn.ModuleList(VisionBlock(config, attn_key) for _ in range(config.depth)) + self.merger = OmniPatchMerger(config, postshuffle=False) + self.merger_list = nn.ModuleList( + OmniPatchMerger(config, postshuffle=True) for _ in config.deepstack_visual_indexes + ) + + def _apply(self, fn, recurse=True): + """Keep the rotary frequencies fp32 through a dtype cast (the worker + casts every submodule to its node's dtype), as the reference does: in + bf16 the far patches of a large image rotate by visibly wrong angles.""" + out = super()._apply(fn, recurse) + rotary = self.rotary_pos_emb + if rotary.inv_freq.dtype != torch.float32 and rotary.inv_freq.device.type != "meta": + rotary.inv_freq = VisionRotaryEmbedding(self.config).inv_freq.to(rotary.inv_freq.device) + return out + + def embed( + self, + pixel_values: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Patch embed plus the resampled position table, and the rope for + the blocks: everything before the block loop, per patch. Grid-derived + inputs come from the caller, so the tower has no host reads.""" + n = pixel_values.shape[0] + # one symbol for every leading dim, not one each + torch._check(indices.shape[0] == n) + torch._check(weights.shape[0] == n) + torch._check(position_ids.shape[0] == n) + + hidden = self.patch_embed(pixel_values) + # bilinear resample of the learned table, as a weighted gather + pos = (self.pos_embed(indices) * weights[:, :, None]).sum(1) + hidden = hidden + pos.to(hidden.dtype) + + cos, sin = self.rotary_pos_emb(position_ids) + return hidden, cos.to(hidden.dtype), sin.to(hidden.dtype) + + def encode( + self, hidden: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """The block loop, ``[patches, hidden]`` in and out, with the hidden + states after each deepstack block (unmerged).""" + taps = [] + for i, block in enumerate(self.blocks): + hidden = block(hidden, cos, sin) + if i in self.config.deepstack_visual_indexes: + taps.append(hidden) + return hidden, taps + + def merge( + self, hidden: torch.Tensor, taps: list[torch.Tensor], + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """The merged embeddings and the merged deepstack features.""" + return self.merger(hidden), [m(t) for m, t in zip(self.merger_list, taps, strict=True)] + + def forward( + self, + pixel_values: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """``[total_patches, patch_numel]`` in; ``[merged_tokens, out_hidden]`` + and one such deepstack tensor per deepstack block out.""" + hidden, cos, sin = self.embed(pixel_values, indices, weights, position_ids) + return self.merge(*self.encode(hidden, cos, sin)) diff --git a/mstar/model/qwen3_omni/config.py b/mstar/model/qwen3_omni/config.py index d03ca6dea2..b61e7306a5 100644 --- a/mstar/model/qwen3_omni/config.py +++ b/mstar/model/qwen3_omni/config.py @@ -16,6 +16,8 @@ from pathlib import Path from typing import Any, Tuple +from mstar.model.components.aut_encoder import AuTEncoderConfig + # --------------------------------------------------------------------------- # Resource label constants # --------------------------------------------------------------------------- @@ -24,6 +26,9 @@ THINKER_ATTN = "thinker_attn" TALKER_ATTN = "talker_attn" +# the encoders' cacheless attention (ragged: within each window / frame) +AUDIO_ATTN = "audio_attn" +VISION_ATTN = "vision_attn" THINKER_POS = "thinker_pos" TALKER_POS = "talker_pos" @@ -126,14 +131,35 @@ def from_dict(cls, d: dict[str, Any]) -> ThinkerConfig: class VisionEncoderConfig: depth: int = 27 hidden_size: int = 1152 + intermediate_size: int = 4304 num_heads: int = 16 + in_channels: int = 3 + patch_size: int = 16 + temporal_patch_size: int = 2 spatial_merge_size: int = 2 + num_position_embeddings: int = 2304 out_hidden_size: int = 3584 + hidden_act: str = "gelu_pytorch_tanh" + rope_theta: float = 10_000.0 deepstack_visual_indexes: tuple[int, ...] = (8, 16, 24) + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_heads + + @property + def num_grid_per_side(self) -> int: + """The learned position table is a square grid.""" + return int(self.num_position_embeddings**0.5) + + @property + def merge_unit(self) -> int: + return self.spatial_merge_size**2 + @classmethod def from_dict(cls, d: dict[str, Any]) -> VisionEncoderConfig: fnames = {f.name for f in cls.__dataclass_fields__.values()} + d = {**d, **d.get("rope_parameters", {})} filtered = {k: v for k, v in d.items() if k in fnames} # HF may store tuple fields as lists if "deepstack_visual_indexes" in filtered and isinstance( @@ -145,25 +171,6 @@ def from_dict(cls, d: dict[str, Any]) -> VisionEncoderConfig: return cls(**filtered) -# --------------------------------------------------------------------------- -# Audio encoder -# --------------------------------------------------------------------------- - -@dataclass -class AudioEncoderConfig: - d_model: int = 1280 - encoder_layers: int = 32 - encoder_attention_heads: int = 20 - output_dim: int = 3584 - max_source_positions: int = 1500 - n_window: int = 50 - - @classmethod - def from_dict(cls, d: dict[str, Any]) -> AudioEncoderConfig: - fnames = {f.name for f in cls.__dataclass_fields__.values()} - return cls(**{k: v for k, v in d.items() if k in fnames}) - - # --------------------------------------------------------------------------- # Talker text backbone # --------------------------------------------------------------------------- @@ -438,7 +445,7 @@ class Qwen3OmniModelConfig: thinker_text: ThinkerTextConfig = field(default_factory=ThinkerTextConfig) thinker: ThinkerConfig = field(default_factory=ThinkerConfig) vision: VisionEncoderConfig = field(default_factory=VisionEncoderConfig) - audio_encoder: AudioEncoderConfig = field(default_factory=AudioEncoderConfig) + audio_encoder: AuTEncoderConfig = field(default_factory=AuTEncoderConfig) talker_text: TalkerTextConfig = field(default_factory=TalkerTextConfig) talker: TalkerConfig = field(default_factory=TalkerConfig) code_predictor: CodePredictorConfig = field(default_factory=CodePredictorConfig) @@ -570,7 +577,7 @@ def _from_raw_dict( thinker_text = ThinkerTextConfig.from_dict(thinker_text_raw) thinker = ThinkerConfig.from_dict(thinker_raw) vision = VisionEncoderConfig.from_dict(vision_raw) - audio_encoder = AudioEncoderConfig.from_dict(audio_enc_raw) + audio_encoder = AuTEncoderConfig.from_hf(audio_enc_raw) # --- Talker --------------------------------------------------- talker_raw = raw.get("talker_config", {}) diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index 2cf17e2828..d1ea5739e7 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/mstar/model/qwen3_omni/qwen3_omni_model.py @@ -50,6 +50,8 @@ PagedKVConfig, PositionConfig, PositionSpec, + RaggedAttentionConfig, + RaggedAttentionSpec, ResourceReqConfig, SamplerSpec, SamplingReqConfig, @@ -83,6 +85,7 @@ from mstar.model.qwen3_omni.components.rope import get_rope_index_vision from mstar.model.qwen3_omni.components.talker import Qwen3OmniCodePredictor from mstar.model.qwen3_omni.config import ( + AUDIO_ATTN, CODE_PRED_SAMPLER, TALKER_ATTN, TALKER_KV, @@ -93,6 +96,7 @@ THINKER_MIXED, THINKER_POS, THINKER_SAMPLER, + VISION_ATTN, ) from mstar.model.qwen3_omni.video_reader import header_charged, install_budgeted_readers from mstar.model.submodule_base import NodeSubmodule @@ -319,7 +323,39 @@ def get_node_resources(self) -> list[NodeResourceSpec]: nodes={"Talker"}, vocab_size=self.config.code_predictor.vocab_size, enable_repetion_penalty=False, - ) + ), + *self._encoder_resources(), + ] + + def _encoder_resources(self) -> list[NodeResourceSpec]: + """The encoders' attention: within each 8 s audio window, and within + each image or video frame group. No cache, so nothing else.""" + from mstar.model.components.aut_submodule import AuTEncoderSubmodule + from mstar.model.qwen3_omni.submodules import VisionEncoderSubmodule + + audio, vision = self.config.audio_encoder, self.config.vision + return [ + RaggedAttentionSpec( + resource_key=AUDIO_ATTN, nodes={"audio_encoder"}, + config=RaggedAttentionConfig( + num_qo_heads=audio.encoder_attention_heads, + num_kv_heads=audio.encoder_attention_heads, + head_dim=audio.head_dim, + max_segments_per_request=AuTEncoderSubmodule.max_windows_per_replay(audio), + ), + ), + RaggedAttentionSpec( + resource_key=VISION_ATTN, nodes={"vision_encoder"}, + config=RaggedAttentionConfig( + num_qo_heads=vision.num_heads, + num_kv_heads=vision.num_heads, + head_dim=vision.head_dim, + # Frames of 256 patches or more fill the largest captured + # bucket at this many; a request with more frames runs the + # block loop eagerly (`VisionEncoderSubmodule._fits_capture`). + max_segments_per_request=VisionEncoderSubmodule.BLOCK_LOOP_PATCH_BUCKETS[-1] // 256, + ), + ), ] def get_request_resource_configs( @@ -1467,9 +1503,9 @@ def _create_submodule( elif node_name == "Code2Wav": return self._create_code2wav_submodule(device) elif node_name == "audio_encoder": - return self._create_audio_encoder_submodule(device) + return self._create_audio_encoder_submodule(device, autocast_dtype) elif node_name == "vision_encoder": - return self._create_vision_encoder_submodule(device) + return self._create_vision_encoder_submodule(device, autocast_dtype) return None @staticmethod @@ -1699,88 +1735,41 @@ def _create_code2wav_submodule(self, device: str) -> NodeSubmodule: config=self.config, ) - def _create_audio_encoder_submodule(self, device: str) -> NodeSubmodule: - """Load the audio encoder (AuT) from HF weights.""" - # Reuse HF audio encoder directly (Whisper-style, not perf-critical) - from transformers import AutoConfig - from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( - Qwen3OmniMoeAudioEncoder, - ) - - from mstar.model.qwen3_omni.components.attention import patch_hf_fa2_int_maxlen - from mstar.model.utils import ModuleAndPrefix, load_weights_from_hf_shards - - # Load config only (no weights) - config = AutoConfig.from_pretrained( - self.local_dir, - trust_remote_code=False, - ) - - # This should be a Qwen3OmniMoeConfig - audio_config = config.thinker_config.audio_config - - # Build the audio encoder from config. - # IMPORTANT: pass attn_implementation="flash_attention_2" so the - # encoder uses the cu_seqlens FA2 path. With the HF default - # (which resolves to "sdpa"), Qwen3OmniMoeAudioAttention runs - # SDPA on the full packed sequence (no per-segment fusion), - # which is significantly slower than FA2's varlen path. - patch_hf_fa2_int_maxlen() - audio_encoder = Qwen3OmniMoeAudioEncoder._from_config( - audio_config, attn_implementation="flash_attention_2" - ) - - load_weights_from_hf_shards( - repo_dir=self.local_dir, - modules=[ModuleAndPrefix(audio_encoder, prefix="thinker.audio_tower")], - device=device, - ) - audio_encoder.eval() - + def _create_audio_encoder_submodule( + self, device: str, autocast_dtype: torch.dtype | None, + ) -> NodeSubmodule: + """The AuT audio encoder (``thinker.audio_tower``).""" + from mstar.model.components.aut_encoder import AuTEncoder + from mstar.model.loader.iterators import iter_safetensors_shards from mstar.model.qwen3_omni.submodules import AudioEncoderSubmodule - return AudioEncoderSubmodule(audio_encoder=audio_encoder, config=self.config) - - def _create_vision_encoder_submodule(self, device: str) -> NodeSubmodule: - """Load the vision encoder (SigLIP2 ViT) from HF weights.""" - # Reuse HF vision encoder directly - from transformers import AutoConfig - from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( - Qwen3OmniMoeVisionEncoder, - ) - - from mstar.model.qwen3_omni.components.attention import patch_hf_fa2_int_maxlen - from mstar.model.utils import ModuleAndPrefix, load_weights_from_hf_shards - - # Load full config (no weights) - config = AutoConfig.from_pretrained( - self.local_dir, - trust_remote_code=False, - ) - - # Extract the vision sub-config - vision_config = config.thinker_config.vision_config - - # Build the vision encoder. - # CRITICAL: pass attn_implementation="flash_attention_2". Without - # this, vision_config._attn_implementation defaults to None and is - # resolved to "sdpa" at runtime (modeling_utils.py:1889). With - # "sdpa", Qwen3OmniMoeVisionAttention.forward falls into the - # per-segment Python loop (modeling_qwen3_omni_moe.py:892-913), - # which issues N sequential attention calls per layer for an - # N-frame video and makes V2T/V2S TTFT far slower. With - # "flash_attention_2", a single varlen FA2 call - # per layer handles all frames at once via cu_seqlens. - patch_hf_fa2_int_maxlen() - vision_encoder = Qwen3OmniMoeVisionEncoder._from_config( - vision_config, attn_implementation="flash_attention_2" - ) - - load_weights_from_hf_shards( - repo_dir=self.local_dir, - modules=[ModuleAndPrefix(vision_encoder, prefix="thinker.visual")], - device=device, - ) - vision_encoder.eval() + with torch.device("meta"): + encoder = AuTEncoder(self.config.audio_encoder, attn_key=AUDIO_ATTN) + if autocast_dtype is not None: + encoder = encoder.to(autocast_dtype) + encoder.to_empty(device=device) + prefix = "thinker.audio_tower." + weights = iter_safetensors_shards(self.local_dir, device=device, prefix=prefix) + encoder.load_weights((k.removeprefix(prefix), v) for k, v in weights) + return AudioEncoderSubmodule(encoder.eval(), self.config) + + def _create_vision_encoder_submodule( + self, device: str, autocast_dtype: torch.dtype | None, + ) -> NodeSubmodule: + """The ViT (``thinker.visual``) with its deepstack mergers.""" + from mstar.model.loader import load_hf_weights + from mstar.model.loader.iterators import iter_safetensors_shards + from mstar.model.qwen3_omni.components.vision import Qwen3OmniVisionTower from mstar.model.qwen3_omni.submodules import VisionEncoderSubmodule - return VisionEncoderSubmodule(vision_encoder=vision_encoder, config=self.config) + from mstar.model.utils import check_loaded + + # built on the device, not on meta: the rotary table is a computed buffer + with torch.device(device): + tower = Qwen3OmniVisionTower(self.config.vision, attn_key=VISION_ATTN) + if autocast_dtype is not None: + tower = tower.to(autocast_dtype) + prefix = "thinker.visual." + weights = iter_safetensors_shards(self.local_dir, device=device, prefix=prefix) + loaded = load_hf_weights(tower, ((k.removeprefix(prefix), v) for k, v in weights)) + check_loaded(tower, loaded, "Qwen3-Omni vision tower") + return VisionEncoderSubmodule(tower.eval(), self.config) diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index a333e5454a..1edd67d00c 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -13,7 +13,6 @@ from __future__ import annotations import functools -import inspect import logging from collections.abc import Mapping from dataclasses import replace @@ -21,17 +20,6 @@ import torch from torch import nn -from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( - _get_feat_extract_output_lengths, - chunk_and_pad_features, - get_audio_cu_seqlens, - get_valid_indices, -) -from transformers.vision_utils import ( - get_vision_cu_seqlens, - get_vision_interpolation_indices_and_weights, - get_vision_position_ids, -) from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo @@ -39,12 +27,23 @@ AcceleratorGraphConfig, BatchedAcceleratorGraphConfig, PackedAcceleratorGraphConfig, + PiecewiseAcceleratorGraphConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, + PiecewisePackedConfig, ) from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep from mstar.engine.resources.attn.flashinfer import FlashInferManager from mstar.engine.resources.sampler.resource import SamplerResource +from mstar.model.components.aut_encoder import AuTEncoder +from mstar.model.components.aut_submodule import AuTEncoderSubmodule from mstar.model.components.mrope import mrope_position_ids +from mstar.model.components.qwen3_vit import vision_interpolation, vision_seq_lengths +from mstar.model.components.qwen3_vit import ( + # the tower's 2D patch grid, not the Thinker's 3D MRoPE ids + vision_position_ids as vision_grid_position_ids, +) from mstar.model.qwen3_omni.components.code2wav import Qwen3OmniMoeCode2Wav from mstar.model.qwen3_omni.components.rope import ( compute_3d_cos_sin, @@ -54,7 +53,9 @@ get_rope_index_vision, ) from mstar.model.qwen3_omni.components.talker import Qwen3OmniCodePredictor, Qwen3OmniTalkerModel +from mstar.model.qwen3_omni.components.vision import Qwen3OmniVisionTower from mstar.model.qwen3_omni.config import ( + AUDIO_ATTN, CODE_PRED_SAMPLER, TALKER_ATTN, TALKER_KV, @@ -65,6 +66,7 @@ THINKER_MIXED, THINKER_POS, THINKER_SAMPLER, + VISION_ATTN, Qwen3OmniModelConfig, ) from mstar.model.qwen3_omni.qwen3_omni_model import image_span @@ -80,20 +82,8 @@ logger = logging.getLogger(__name__) - -def _accepts_kwarg(fn: Any, name: str) -> bool: - """True if fn takes `name` as a keyword (explicitly or via **kwargs).""" - try: - params = inspect.signature(fn).parameters - except (TypeError, ValueError): - return False - if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): - return True - param = params.get(name) - return param is not None and param.kind in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) +# The vision tower's block loop, as a piecewise capture region +VISION_BLOCK_LOOP = "vision_block_loop" def talker_text_mask(input_ids: torch.Tensor, config: Qwen3OmniModelConfig) -> torch.Tensor: @@ -118,33 +108,17 @@ def talker_text_mask(input_ids: torch.Tensor, config: Qwen3OmniModelConfig) -> t # =================================================================== -# 1. AudioEncoderSubmodule (enc_dec engine) +# 1. AudioEncoderSubmodule # =================================================================== -class AudioEncoderSubmodule(NodeSubmodule): - """Thin wrapper around the HF Whisper-style audio encoder. - - Extracts mel spectrograms from raw audio inputs, pads for batching, - and runs the encoder to produce audio embeddings that the Thinker - will splice into its input sequence. +class AudioEncoderSubmodule(AuTEncoderSubmodule): + """The AuT encoder (see ``AuTEncoderSubmodule``) over the unpadded + ``(num_mel_bins, frames)`` log-mel features ``process_prompt`` made.""" - Runs once per request (not batched across requests). - """ - - marks_dynamic_dims = True - - def __init__(self, audio_encoder: nn.Module, config: Qwen3OmniModelConfig): - super().__init__() - self.audio_encoder = audio_encoder + def __init__(self, audio_encoder: AuTEncoder, config: Qwen3OmniModelConfig): + super().__init__(audio_encoder, attn_key=AUDIO_ATTN) self.config = config - # Older transformers take return_dict; newer ones dropped it. - self._encoder_takes_return_dict = _accepts_kwarg(audio_encoder.forward, "return_dict") - n_window = audio_encoder.n_window - # the length of a full attention window, since no clip's window is longer - self._max_seqlen = _get_feat_extract_output_lengths(2 * n_window, n_window) * ( - audio_encoder.n_window_infer // (2 * n_window) - ) def prepare_inputs( self, @@ -153,73 +127,38 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs ) -> NodeInputs: - # Edge name from graph walk is "audio_features" - audio_features = inputs["audio_features"][0] - audio_seqlens = inputs["audio_seqlens"][0] - - # Built outside the compiled forward, since these helpers read the clip's length on the - # host, which would compile a graph for every length. - encoder = self.audio_encoder - padded_feature, chunk_lengths = chunk_and_pad_features(audio_features, audio_seqlens, encoder.n_window) - tensor_inputs = { - "padded_feature": padded_feature, - "chunk_lengths": chunk_lengths, - "valid_indices": get_valid_indices(chunk_lengths, encoder.n_window), - "cu_seqlens": get_audio_cu_seqlens(chunk_lengths, audio_seqlens, encoder.n_window_infer, encoder.n_window), - } - for name in ("chunk_lengths", "valid_indices", "cu_seqlens"): - torch._dynamo.mark_dynamic(tensor_inputs[name], 0) - # A clip shorter than one chunk is padded only to its own length, so the frame dimension of - # `padded_feature` varies too. Some clips give a dimension of size 1 inside the encoder, which torch - # specializes, so `mark_dynamic` would raise for them. - torch._dynamo.maybe_mark_dynamic(padded_feature, 0) - torch._dynamo.maybe_mark_dynamic(padded_feature, 2) - return NodeInputs(tensor_inputs=tensor_inputs) - - def forward( - self, - graph_walk: str, - engine_inputs: ModelInputsFromEngine, - **kwargs, - ) -> NameToTensorList: - """Run the audio encoder and return embeddings. - - Returns: - {"audio_embeds": [tensor of shape (audio_tokens, hidden_size)]} - """ - extra_kwargs = {"return_dict": True} if self._encoder_takes_return_dict else {} - audio_out = self.audio_encoder( - max_seqlen=self._max_seqlen, - **kwargs, - **extra_kwargs, - ) - audio_embeds = getattr(audio_out, "last_hidden_state", audio_out) - - # Flatten to (num_audio_tokens, hidden_size) if needed - if audio_embeds.dim() == 3: - audio_embeds = audio_embeds.squeeze(0) - - return {"audio_embeds": [audio_embeds]} + features = inputs["audio_features"][0].to(device=self.get_device(), dtype=self._param_dtype()) + return self.feature_inputs(features) # =================================================================== -# 2. VisionEncoderSubmodule (enc_dec engine) +# 2. VisionEncoderSubmodule # =================================================================== class VisionEncoderSubmodule(NodeSubmodule): - """Thin wrapper around the HF vision encoder (ViT + spatial merge). - - Extracts pixel_values and grid_thw from inputs, computes cu_seqlens - for FlashAttention, runs the encoder, and returns vision embeddings - plus DeepStack intermediate features for the Thinker. - - Runs once per request (not batched across requests). + """The ViT, which turns a request's pixel patches into Thinker-width + embeddings and the deepstack features. + + A request's images (or one video's frames) go through packed, in one + call. Grid maths runs in ``prepare_inputs`` and each frame is one segment + of the node's ragged-attention step, so the only size the tower sees + change is its patch count. + + The block loop is captured one request per replay (eager, the tower is + launch-bound), so the encoder takes one request a step. A replay's + attention plan holds at most the resource's ``max_segments_per_request`` + frames; a request with more runs the block loop eagerly. """ - marks_dynamic_dims = True + # Patch-count buckets, one request per replay. A larger request runs eagerly. + BLOCK_LOOP_PATCH_BUCKETS = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768] + # Not compiled: the block loop is captured, and around it the forward + # takes the frame lengths, so a compiled wrapper would recompile per + # frame count while the prefix cache holds what earlier graphs computed. + disable_torch_compile = True - def __init__(self, vision_encoder: nn.Module, config: Qwen3OmniModelConfig): + def __init__(self, vision_encoder: Qwen3OmniVisionTower, config: Qwen3OmniModelConfig): super().__init__() self.vision_encoder = vision_encoder self.config = config @@ -231,92 +170,192 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs ) -> NodeInputs: - """Extract pixel_values, grid_thw, and compute cu_seqlens. - - ``pixel_values`` and ``image_grid_thw`` are produced by - ``Qwen3OmniModel.process_prompt`` from the raw ``image_inputs`` - loaded by the data worker. - """ - # Edge name from graph walk is "pixel_values" - pixel_values = inputs["pixel_values"][0] # (N_patches, C, patch_H, patch_W) - grid_thw = inputs.get("image_grid_thw", inputs.get("grid_thw", [None]))[0] - - # Normalize grid_thw to shape (num_images, 3). Single-image requests - # store grid_thw as a 1-D tensor [T, H, W] (after process_prompt - # indexes proc_out["image_grid_thw"][0] to strip the batch dim); - # the per-image iteration logic below requires 2-D. - if grid_thw is None: - raise ValueError( - "VisionEncoder: 'image_grid_thw' input is None. " - "Make sure process_prompt is producing image_grid_thw via the " - "HF AutoImageProcessor." - ) - if grid_thw.dim() == 1: - grid_thw = grid_thw.unsqueeze(0) # (1, 3) - - # Built outside the compiled forward, since these helpers call `.tolist()` on the - # grid, which would compile a graph for every image size. - encoder = self.vision_encoder - interp_indices, interp_weights = get_vision_interpolation_indices_and_weights( - grid_thw, num_grid_per_side=encoder.num_grid_per_side, mode=encoder.interpolation_mode, - align_corners=encoder.interpolation_align_corners, spatial_merge_size=encoder.spatial_merge_size, + """``pixel_values`` and ``image_grid_thw`` come from + ``Qwen3OmniModel.process_prompt`` (a video walk maps its + ``pixel_values_videos`` / ``video_grid_thw`` onto them).""" + device = self.get_device() + vision = self.config.vision + merge = vision.spatial_merge_size + # (num_images, 3): a single image's grid arrives as a bare [t, h, w] + grid = [ + (int(t), int(h), int(w)) + for g in inputs["image_grid_thw"] + for t, h, w in g.reshape(-1, 3).tolist() + ] + # Built here, not in the forward, so the tower reads nothing off the host. + indices, weights = vision_interpolation(grid, vision.num_grid_per_side, merge, device) + return NodeInputs( + # what the engine leases the captured block loop's bucket by + input_seq_len=sum(t * h * w for t, h, w in grid), + tensor_inputs={ + "pixel_values": torch.cat(list(inputs["pixel_values"]), dim=0), + "indices": indices, + "weights": weights, + "position_ids": vision_grid_position_ids(grid, merge, device), + # not a forward arg: `declare_step` makes a segment of each + "seq_lengths": vision_seq_lengths(grid), + }, ) - tensor_inputs = { - "pixel_values": pixel_values, - "grid_thw": grid_thw, - "interp_indices": interp_indices, - "interp_weights": interp_weights, - "position_ids": get_vision_position_ids(grid_thw, encoder.spatial_merge_size), - "cu_seqlens": get_vision_cu_seqlens(grid_thw), - } - for name in ("pixel_values", "interp_indices", "interp_weights", "position_ids", "cu_seqlens"): - torch._dynamo.mark_dynamic(tensor_inputs[name], 0) - return NodeInputs(tensor_inputs=tensor_inputs) - def forward( + def declare_step( self, graph_walk: str, - engine_inputs: ModelInputsFromEngine, - pixel_values: torch.Tensor, - grid_thw: torch.Tensor, + request_ids: list[str], + inputs: list[NodeInputs], + piecewise_leases: Mapping[str, SlotLease] | None = None, **kwargs, - ) -> NameToTensorList: - """Run vision encoder, return embeddings and DeepStack features. + ) -> SubmoduleStep | None: + # leased, the captured block loop plans its own attention per replay + if (piecewise_leases or {}).get(VISION_BLOCK_LOOP): + return None + # one segment per frame: a frame attends to itself alone + segments = [ + Segment(request_id=rid, label="main", span=span) + for rid, inp in zip(request_ids, inputs, strict=True) + for span in inp.tensor_inputs["seq_lengths"] + ] + return SubmoduleStep(segments=segments, steps={VISION_ATTN: AttentionStep(causal=False)}) + + def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: + 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 + + 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 request per replay, by patch-count bucket. - Returns: - { - "vision_embeds": [tensor of shape (vision_tokens, hidden_size)], - "deepstack": [list of intermediate layer features], + Patch embed, position resample and rope run eagerly before it, the + mergers after (they quarter the row count, which a region's output + view cannot express).""" + vision = self.config.vision + hidden, head_dim = vision.hidden_size, vision.head_dim + + def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + n = shape.total_tokens + return { + "hidden": torch.zeros(n, hidden, dtype=autocast_dtype, device=device), + "cos": torch.zeros(n, head_dim, dtype=autocast_dtype, device=device), + "sin": torch.zeros(n, head_dim, dtype=autocast_dtype, device=device), } - """ - logger.debug( - "Running VisionEncoder with pixel_values shape=%s, grid_thw shape=%s", - pixel_values.shape, grid_thw.shape, - ) - # HF vision encoder returns (hidden_states, deepstack_features) - # depending on the model variant; handle both cases - encoder_output = self.vision_encoder( - pixel_values, - grid_thw=grid_thw, - # each video frame is its own segment, and a call encodes one item, so this is one frame's patches - max_seqlen=pixel_values.shape[0] // (kwargs["cu_seqlens"].shape[0] - 1), - **kwargs, - ) - if isinstance(encoder_output, tuple): - vision_embeds, deepstack = encoder_output - else: - vision_embeds = encoder_output.pooler_output - deepstack = encoder_output.deepstack_features + 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={VISION_ATTN: AttentionStep(causal=False)}, + ) + + return { + VISION_BLOCK_LOOP: PiecewisePackedConfig( + capture_fn=self._capture_block_loop, + make_static_inputs=make_static_inputs, + declare_step=declare_step, + lease_before_step=True, + can_lease=self._fits_capture, + total_tokens=self.BLOCK_LOOP_PATCH_BUCKETS, + capture_batch_sizes=[1], + ) + } - if isinstance(deepstack, torch.Tensor): - deepstack = [deepstack] + def _fits_capture(self, inputs: list[NodeInputs]) -> bool: + """Whether the batch's frames fit a replay's attention plan: the + resource's per-request ceiling, read from it so a YAML override holds.""" + 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) + + 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 + hidden, taps = self.vision_encoder.encode(s["hidden"], s["cos"], s["sin"]) + return {"hidden": hidden, **{f"deepstack_{i}": t for i, t in enumerate(taps)}} + + @torch.compiler.disable + def _replay_block_loop( + self, + engine_inputs: ModelInputsFromEngine, + hidden: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + seq_lengths: tuple[int, ...], + ) -> tuple[torch.Tensor, list[torch.Tensor]] | None: + """The block loop replayed once per request, 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 + ``can_lease`` may have refused the batch. The bucket check keeps a + forward driven outside the engine (no declaration either) eager. + """ + runner = engine_inputs.piecewise_runners.get(VISION_BLOCK_LOOP) + 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]), + ): + return None + # one replay for every frame: a second plan overwrites the schedule a queued replay still reads + out = runner.run( + static_inputs={"hidden": hidden, "cos": cos, "sin": sin}, + request_ids=[engine_inputs.request_ids[0]], + seq_lens=list(seq_lengths), + real_bs=1, + ) + # views of the runner's buffers, merged into new tensors within this step + taps = [out.get_view(f"deepstack_{i}") for i in range(len(self.config.vision.deepstack_visual_indexes))] + return out.get_view("hidden"), taps + def preprocess( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + inputs: list[NodeInputs], + ) -> dict[str, Any]: + """The batch's one request (see `max_batch_size`), its frames packed + in `declare_step`'s segment order.""" + (inp,) = inputs + device = self.get_device() + tensors = inp.tensor_inputs return { - "vision_embeds": [vision_embeds], - "deepstack": deepstack if deepstack is not None else [torch.tensor([])], + "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"]), } + 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] + + def forward_batched( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + pixel_values: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + position_ids: torch.Tensor, + seq_lengths: tuple[int, ...] = (), + **kwargs, + ) -> dict[str, NameToTensorList]: + """Returns ``vision_embeds`` ``[merged_tokens, out_hidden]`` and one + ``deepstack`` tensor of that size per deepstack block.""" + tower = self.vision_encoder + hidden, cos, sin = tower.embed(pixel_values, indices, weights, position_ids) + encoded = self._replay_block_loop(engine_inputs, hidden, cos, sin, seq_lengths) + if encoded is None: + encoded = tower.encode(hidden, cos, sin) + embeds, deepstack = tower.merge(*encoded) + return {engine_inputs.request_ids[0]: {"vision_embeds": [embeds], "deepstack": deepstack}} + # =================================================================== # 3. ThinkerSubmodule (ar engine) -- MOST COMPLEX diff --git a/pyproject.toml b/pyproject.toml index 656bc382f8..7773d12f00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -188,9 +188,6 @@ cosmos3 = [ ] qwen3_omni = [ - # NOTE: flash-attn is required by Qwen3-Omni but is NOT listed here — its - # prebuilt wheels aren't on PyPI, so a plain install compiles from source - # (and fails on CUDA 13). Install it separately; see docs/installation.rst. "einops", "flashinfer-python==0.7.2", "huggingface-hub", diff --git a/test/modular/test_qwen3_omni_encoders.py b/test/modular/test_qwen3_omni_encoders.py new file mode 100644 index 0000000000..dec860cb82 --- /dev/null +++ b/test/modular/test_qwen3_omni_encoders.py @@ -0,0 +1,240 @@ +"""Qwen3-Omni's encoders declare their attention as ragged resources and run +natively against them. + +CPU, tiny random weights. A stand-in resource attends within the segments the +node's step declared (what the planned kernel does), and transformers' +``Qwen3OmniMoe{Audio,Vision}Encoder`` on SDPA are the oracle: the same +windows, frames and deepstack outputs, with no flash_attn anywhere. +""" +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +transformers = pytest.importorskip("transformers") + +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( # noqa: E402 + Qwen3OmniMoeAudioEncoderConfig, + Qwen3OmniMoeVisionEncoderConfig, +) +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( # noqa: E402 + Qwen3OmniMoeAudioEncoder, + Qwen3OmniMoeVisionEncoder, +) + +from mstar.engine.resources import RaggedAttentionSpec # noqa: E402 +from mstar.engine.resources.base import EngineResourceInfo, build_resource # noqa: E402 +from mstar.model.components.aut_encoder import AuTEncoder, AuTEncoderConfig # noqa: E402 +from mstar.model.components.aut_submodule import AuTEncoderSubmodule # noqa: E402 +from mstar.model.loader import load_hf_weights # noqa: E402 +from mstar.model.qwen3_omni.components.vision import Qwen3OmniVisionTower # noqa: E402 +from mstar.model.qwen3_omni.config import ( # noqa: E402 + AUDIO_ATTN, + VISION_ATTN, + Qwen3OmniModelConfig, + VisionEncoderConfig, +) +from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel # noqa: E402 +from mstar.model.qwen3_omni.submodules import AudioEncoderSubmodule, VisionEncoderSubmodule # noqa: E402 +from mstar.model.utils import check_loaded # noqa: E402 + +AUDIO = dict( + d_model=64, encoder_layers=2, encoder_attention_heads=4, encoder_ffn_dim=128, + num_mel_bins=16, n_window=50, n_window_infer=200, conv_chunksize=3, + downsample_hidden_size=8, output_dim=48, max_source_positions=1500, +) +VISION = dict( + depth=3, hidden_size=32, intermediate_size=64, num_heads=4, patch_size=4, + temporal_patch_size=2, spatial_merge_size=2, out_hidden_size=24, + num_position_embeddings=16, deepstack_visual_indexes=[0, 2], +) + + +class StandInRagged: + """The ragged resource as a node sees it: the step's segments in, and + ``run`` attending within each of them, as the planned kernel does.""" + + max_segments_per_request = 64 + + def __init__(self): + self.spans: list[int] = [] + self.calls = 0 + + def plan(self, step): + self.spans = [s.span for s in step.segments] + + def run(self, q, k, v): + self.calls += 1 + assert sum(self.spans) == q.shape[0], "the forward's rows are not the planned segments" + out, start = torch.empty_like(q), 0 + for n in self.spans: + o = F.scaled_dot_product_attention(*(t[start:start + n].transpose(0, 1) for t in (q, k, v))) + out[start:start + n] = o.transpose(0, 1) + start += n + return out + + +def _engine_inputs(rids): + # no capture runners and a declared step: the eager path + return SimpleNamespace(request_ids=list(rids), piecewise_runners={}, step=object()) + + +def _fwd(rid="r"): + return SimpleNamespace(request_id=rid) + + +@pytest.fixture(scope="module") +def audio_pair(): + torch.manual_seed(0) + hf = Qwen3OmniMoeAudioEncoder._from_config( + Qwen3OmniMoeAudioEncoderConfig(**AUDIO), attn_implementation="sdpa", + ).eval() + encoder = AuTEncoder(AuTEncoderConfig.from_hf(AUDIO), attn_key=AUDIO_ATTN) + encoder.load_weights(list(hf.state_dict().items())) + sub = AudioEncoderSubmodule(encoder.eval(), Qwen3OmniModelConfig(audio_encoder=encoder.config)) + ragged = StandInRagged() + sub.bind_node_resources({AUDIO_ATTN: ragged}) + return hf, sub, ragged + + +@pytest.fixture(scope="module") +def vision_pair(): + torch.manual_seed(0) + hf = Qwen3OmniMoeVisionEncoder._from_config( + Qwen3OmniMoeVisionEncoderConfig(**VISION), attn_implementation="sdpa", + ).eval() + with torch.no_grad(): + for p in hf.parameters(): + p.copy_(torch.randn_like(p) * 0.1) + config = VisionEncoderConfig.from_dict(VISION) + tower = Qwen3OmniVisionTower(config, attn_key=VISION_ATTN) + check_loaded(tower, load_hf_weights(tower, hf.state_dict().items()), "tower") + sub = VisionEncoderSubmodule(tower.eval(), Qwen3OmniModelConfig(vision=config)) + ragged = StandInRagged() + sub.bind_node_resources({VISION_ATTN: ragged}) + return hf, sub, ragged + + +def test_encoders_declare_ragged_attention(): + """Each encoder node gets one ragged resource with its tower's heads, and + the Thinker and Talker keep their paged attention.""" + model = object.__new__(Qwen3OmniModel) + model.config = Qwen3OmniModelConfig( + audio_encoder=AuTEncoderConfig.from_hf(AUDIO), vision=VisionEncoderConfig.from_dict(VISION), + ) + specs = {s.resource_key: s for s in model.get_node_resources()} + + audio, vision = specs[AUDIO_ATTN], specs[VISION_ATTN] + assert isinstance(audio, RaggedAttentionSpec) and isinstance(vision, RaggedAttentionSpec) + assert audio.nodes == {"audio_encoder"} and vision.nodes == {"vision_encoder"} + assert (audio.config.num_qo_heads, audio.config.num_kv_heads, audio.config.head_dim) == (4, 4, 16) + assert (vision.config.num_qo_heads, vision.config.num_kv_heads, vision.config.head_dim) == (4, 4, 8) + # a capture replay can always be handed the windows of what fits its largest bucket + window = AuTEncoderConfig.from_hf(AUDIO).window_tokens + biggest = AuTEncoderSubmodule.ENCODE_TOKEN_BUCKETS[-1] + assert audio.config.max_segments_per_request * window >= biggest + others = {k for k, s in specs.items() if s.nodes & {"audio_encoder", "vision_encoder"}} + assert others == {AUDIO_ATTN, VISION_ATTN} + + +def test_ragged_spec_builds_its_backend_without_flash_attn(monkeypatch): + """The declaration resolves to the platform's ragged backend (FlashInfer by + default); the backend module is stubbed, so nothing here needs a GPU.""" + built = {} + + class FakeManager: + def __init__(self, device, dtype, config): + built.update(device=device, dtype=dtype, config=config) + + fake = SimpleNamespace(FlashInferRaggedManager=FakeManager) + monkeypatch.setitem(sys.modules, "mstar.engine.resources.attn.ragged.flashinfer", fake) + model = object.__new__(Qwen3OmniModel) + model.config = Qwen3OmniModelConfig(vision=VisionEncoderConfig.from_dict(VISION)) + spec = next(s for s in model.get_node_resources() if s.resource_key == VISION_ATTN) + + res = build_resource(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.bfloat16)) + + assert isinstance(res, FakeManager) + assert built["dtype"] == torch.bfloat16 and built["config"].head_dim == 8 + assert "flash_attn" not in sys.modules + + +def test_audio_windows_match_hf(audio_pair): + """Three clips packed in one forward, every window a declared segment: + each clip's embeddings match HF's encoder on that clip alone.""" + hf, sub, ragged = audio_pair + # under one window, about one, and several (a window is 200 frames here) + feats = [torch.randn(AUDIO["num_mel_bins"], n) for n in (150, 230, 905)] + rows = [ + sub.prepare_inputs("prefill_audio", _fwd(), {"audio_features": [f], "audio_seqlens": [torch.tensor([n])]}) + for f, n in zip(feats, (150, 230, 905), strict=True) + ] + step = sub.declare_step("prefill_audio", ["a", "b", "c"], rows) + window = sub.encoder.config + assert step.steps[AUDIO_ATTN].causal is False + assert [(s.request_id, s.span) for s in step.segments] == [ + (rid, w) for rid, r in zip("abc", rows, strict=True) for w in window.window_lengths(r.input_seq_len) + ] + ragged.plan(step) + with torch.no_grad(): + out = sub.forward_batched("prefill_audio", _engine_inputs("abc"), **sub.preprocess("prefill_audio", None, rows)) + for rid, f in zip("abc", feats, strict=True): + want = hf(f, feature_lens=torch.tensor([f.shape[1]])).last_hidden_state + got = out[rid]["audio_embeds"][0] + assert got.shape == want.shape + torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4) + assert ragged.calls == AUDIO["encoder_layers"] + + +@pytest.mark.parametrize( + "grid", [[[1, 4, 6]], [[3, 4, 4]], [[1, 2, 2], [1, 4, 4]]], ids=["image", "video", "two_images"], +) +def test_vision_frames_and_deepstack_match_hf(vision_pair, grid): + """One segment per frame, and the merged embeddings and every deepstack + feature match HF's encoder.""" + hf, sub, ragged = vision_pair + grid_thw = torch.tensor(grid) + patch_numel = 3 * VISION["temporal_patch_size"] * VISION["patch_size"] ** 2 + pixels = torch.randn(int(grid_thw.prod(-1).sum()), patch_numel) + row = sub.prepare_inputs("prefill_vision", _fwd(), {"pixel_values": [pixels], "image_grid_thw": [grid_thw]}) + assert row.input_seq_len == pixels.shape[0] + step = sub.declare_step("prefill_vision", ["v"], [row]) + assert [s.span for s in step.segments] == [h * w for t, h, w in grid for _ in range(t)] + assert step.steps[VISION_ATTN].causal is False + ragged.plan(step) + with torch.no_grad(): + out = sub.forward("prefill_vision", _engine_inputs("v"), **sub.preprocess("prefill_vision", None, [row])) + want = hf(pixels, grid_thw=grid_thw) + torch.testing.assert_close(out["vision_embeds"][0], want.pooler_output, atol=1e-4, rtol=1e-4) + assert len(out["deepstack"]) == len(VISION["deepstack_visual_indexes"]) + for got, ref in zip(out["deepstack"], want.deepstack_features, strict=True): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +def test_vision_capture_refuses_more_frames_than_a_replay_plans(vision_pair): + """A request with more frames than the resource's per-request ceiling + takes no lease, so its block loop runs eagerly.""" + _, sub, ragged = vision_pair + rows = [SimpleNamespace(tensor_inputs={"seq_lengths": (4,) * n}) for n in (1, ragged.max_segments_per_request + 1)] + assert sub._fits_capture(rows[:1]) + assert not sub._fits_capture(rows[1:]) + + +def test_rotary_table_stays_fp32_through_a_cast(): + tower = Qwen3OmniVisionTower(VisionEncoderConfig.from_dict(VISION), attn_key=VISION_ATTN) + want = tower.rotary_pos_emb.inv_freq.clone() + tower = tower.to(torch.bfloat16) + assert tower.patch_embed.proj.weight.dtype == torch.bfloat16 + assert tower.rotary_pos_emb.inv_freq.dtype == torch.float32 + assert torch.equal(tower.rotary_pos_emb.inv_freq, want) + + +def test_tower_parameters_are_the_checkpoints(): + """Names line up with ``thinker.visual.*``, so the checkpoint loads by name.""" + hf = Qwen3OmniMoeVisionEncoder._from_config(Qwen3OmniMoeVisionEncoderConfig(**VISION), attn_implementation="sdpa") + tower = Qwen3OmniVisionTower(VisionEncoderConfig.from_dict(VISION), attn_key=VISION_ATTN) + assert dict(tower.named_parameters()).keys() == dict(hf.named_parameters()).keys() diff --git a/test/modular/test_qwen3_omni_prefix_cache.py b/test/modular/test_qwen3_omni_prefix_cache.py index bb5e48fd6a..b9e0ff0755 100644 --- a/test/modular/test_qwen3_omni_prefix_cache.py +++ b/test/modular/test_qwen3_omni_prefix_cache.py @@ -7,15 +7,6 @@ import pytest import torch from torch import nn -from torch._dynamo.utils import counters -from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( - Qwen3OmniMoeAudioEncoderConfig, - Qwen3OmniMoeVisionEncoderConfig, -) -from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( - Qwen3OmniMoeAudioEncoder, - Qwen3OmniMoeVisionEncoder, -) from mstar.engine.engine import Engine from mstar.engine.resources import PositionStep, Resource, StepContext, StepRunner @@ -32,9 +23,6 @@ NODE = "Thinker" HIDDEN = 8 -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="flash-attn runs on a GPU") - - class _StubThinker(nn.Module): def __init__(self): super().__init__() @@ -256,36 +244,6 @@ def test_the_layout_lists_each_walk_the_schedule_runs_with_the_slots_it_writes(m ) -def _vision(attn_implementation: str = "sdpa") -> VisionEncoderSubmodule: - # The tensors `prepare_inputs` builds depend only on the grid, the merge size and - # `num_grid_per_side`, so a small encoder is enough. - config = Qwen3OmniMoeVisionEncoderConfig( - depth=2, hidden_size=32, intermediate_size=64, num_heads=2, out_hidden_size=16, - num_position_embeddings=64, deepstack_visual_indexes=[0, 1], - ) - encoder = Qwen3OmniMoeVisionEncoder._from_config(config, attn_implementation=attn_implementation).eval() - return VisionEncoderSubmodule(encoder, config=None) - - -def _image(vision): - grid, pixels = torch.tensor([1, 4, 6]), torch.randn(24, 3 * 2 * 16 * 16) - inputs = {"pixel_values": [pixels], "image_grid_thw": [grid]} - return grid, pixels, vision.prepare_inputs("prefill_vision", None, inputs) - - -def test_the_vision_encoder_embeds_an_image_from_its_prepared_inputs_as_it_would_alone(): - vision = _vision() - grid, pixels, prepared = _image(vision) - encoder = vision.vision_encoder - with torch.no_grad(): - ours = vision.forward("prefill_vision", None, **prepared.tensor_inputs) - alone = encoder(pixels, grid_thw=grid[None]) - - assert torch.equal(ours["vision_embeds"][0], alone.pooler_output) and all( - torch.equal(layer, own) for layer, own in zip(ours["deepstack"], alone.deepstack_features, strict=True) - ), "the inputs built before the compiled call embed the image otherwise than the encoder does" - - def test_a_prompt_with_a_video_keys_only_its_leading_text(model): out = _prompt( model, ["text", "video", "text"], {"video_inputs": [torch.zeros(4, 3, 64, 96)]}, @@ -297,96 +255,21 @@ def test_a_prompt_with_a_video_keys_only_its_leading_text(model): "a prompt with a video keyed other than its leading text" ) -def _audio(attn_implementation: str = "sdpa") -> AudioEncoderSubmodule: - # The tensors `prepare_inputs` builds depend only on the clip's length and the encoder's - # windows, so a small encoder is enough. - config = Qwen3OmniMoeAudioEncoderConfig( - encoder_layers=1, d_model=32, encoder_attention_heads=2, encoder_ffn_dim=64, output_dim=16, - downsample_hidden_size=8, - ) - encoder = Qwen3OmniMoeAudioEncoder._from_config(config, attn_implementation=attn_implementation).eval() - return AudioEncoderSubmodule(encoder, config=None) - - -def _clip(frames: int, device: str = "cpu") -> dict: - return {"audio_features": [torch.randn(128, frames, device=device)], - "audio_seqlens": [torch.tensor([frames], device=device)]} - - -# shorter than one chunk, shorter than one attention window, and three windows -@pytest.mark.parametrize("frames", [60, 600, 1773]) -def test_the_audio_encoder_embeds_a_clip_from_its_prepared_inputs_as_it_would_alone(frames): - audio, clip = _audio(), _clip(frames) - prepared = audio.prepare_inputs("prefill_audio", None, clip) - # sdpa ignores `max_seqlen`, so comparing the outputs can't test it - assert audio._max_seqlen >= prepared.tensor_inputs["cu_seqlens"].diff().max(), ( - "flash-attn would be told a bound shorter than the clip's longest attention window" - ) - with torch.no_grad(): - ours = audio.forward("prefill_audio", None, **prepared.tensor_inputs) - alone = audio.audio_encoder(clip["audio_features"][0], feature_lens=clip["audio_seqlens"][0]) - - assert torch.equal(ours["audio_embeds"][0], alone.last_hidden_state), ( - "the inputs built before the compiled call embed the clip otherwise than the encoder does" - ) - - -def test_the_audio_encoder_compiles_with_the_default_while_its_output_is_cached(monkeypatch): +@pytest.mark.parametrize("node", ["audio_encoder", "vision_encoder"]) +def test_the_encoders_are_not_compiled_while_their_output_is_cached(monkeypatch, node): + """The engine compiles an item encoder's forward with dynamic=True, and a + graph compiled for a new size rounds otherwise than the one whose output + the prefix cache holds. Both encoders capture their block stacks and opt + out of the compile, so no size compiles a graph of its own.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch, "compile", lambda fn, dynamic, **kwargs: dynamic) + compiled = [] + monkeypatch.setattr(torch, "compile", lambda fn, **kwargs: compiled.append(fn) or fn) + cls = {"audio_encoder": AudioEncoderSubmodule, "vision_encoder": VisionEncoderSubmodule}[node] engine = Engine.__new__(Engine) engine._device = torch.device("cuda") - engine._item_encoders = {"audio_encoder"} - engine._submodules = {"audio_encoder": SimpleNamespace(submodule=_audio(), forward="eager")} + engine._item_encoders = {node} + engine._submodules = {node: SimpleNamespace(submodule=object.__new__(cls), forward="eager")} engine._compile_submodules() - assert engine._submodules["audio_encoder"].forward is None, ( - "the audio encoder compiled with dynamic=True or stayed eager while its output is cached" - ) - -@requires_cuda -@pytest.mark.parametrize(("node", "seen", "unseen"), [ - ("vision", [(1, 4, 6)], [(1, 8, 10)]), - # The clips of 2, 60 and 50100 frames have a dimension of size 1 inside the encoder, so they compile - # graphs of their own. - ("audio", [1773, 60, 50100, 2], [600, 35]), -]) -def test_an_encoder_compiles_no_graph_for_a_size_inside_the_ranges_it_has_compiled(node, seen, unseen): - pytest.importorskip("flash_attn") - torch._dynamo.reset() - if node == "vision": - submodule, walk = _vision("flash_attention_2"), "prefill_vision" - - def inputs(grid): - return {"pixel_values": [torch.randn(grid[0] * grid[1] * grid[2], 3 * 2 * 16 * 16, device="cuda")], - "image_grid_thw": [torch.tensor(grid, device="cuda")]} - else: - submodule, walk = _audio("flash_attention_2"), "prefill_audio" - - def inputs(frames): - return _clip(frames, "cuda") - - submodule = submodule.to(device="cuda", dtype=torch.bfloat16) - engine = Engine.__new__(Engine) - engine._device = torch.device("cuda") - engine._item_encoders = {node} - engine._submodules = {node: SimpleNamespace(submodule=submodule)} - engine._compile_submodules() - forward = engine._submodules[node].forward - - def encode(size): - prepared = submodule.prepare_inputs(walk, None, inputs(size)) - with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): - forward(walk, None, **prepared.tensor_inputs) - - for size in seen: - encode(size) - frames = counters["frames"]["ok"] - with torch._dynamo.config.patch(error_on_recompile=True): - for size in unseen: - encode(size) - - assert counters["frames"]["ok"] == frames, ( - f"the {node} encoder compiled a graph for a size inside a range it had compiled" - ) + assert not compiled and engine._submodules[node].forward == "eager" diff --git a/test/modular/test_trust_remote_code.py b/test/modular/test_trust_remote_code.py index 6217b7fd04..b78edd7c47 100644 --- a/test/modular/test_trust_remote_code.py +++ b/test/modular/test_trust_remote_code.py @@ -62,21 +62,6 @@ def test_qwen3_omni_tokenizer_and_processor_without_remote_code(qwen3_omni, monk assert [c["trust_remote_code"] for c in tok_calls + proc_calls] == [False, False] -@pytest.mark.parametrize("loader", ["_create_audio_encoder_submodule", "_create_vision_encoder_submodule"]) -def test_qwen3_omni_encoder_config_without_remote_code(qwen3_omni, monkeypatch, tmp_path, loader): - from transformers import AutoConfig - - calls = [] - _patch_from_pretrained(monkeypatch, AutoConfig, calls, result=_StopLoadError) - model = qwen3_omni.Qwen3OmniModel.__new__(qwen3_omni.Qwen3OmniModel) - model.local_dir = str(tmp_path) - - with pytest.raises(_StopLoadError): - getattr(model, loader)("cpu") - - assert [c["trust_remote_code"] for c in calls] == [False] - - def test_omnivoice_checkpoint_config_without_remote_code(monkeypatch): from transformers import AutoConfig From de85d7b17bd6bbb7c253b0df8373b85ad548c6a9 Mon Sep 17 00:00:00 2001 From: merceod Date: Sat, 31 Oct 2026 04:22:53 -0700 Subject: [PATCH 3/3] docs: the flash-attn section heading matches its references --- docs/installation.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 32d7a7dab8..4de40f9e1f 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -302,8 +302,8 @@ mstar reads only the speech-tokenizer source files and runs them under transform lists them (``transformers==4.57.3``, ``accelerate==1.12.0``, ``gradio``, ``sox`` and others) as unmet for ``qwen-tts``; that is expected. -flash-attn (Qwen3-Omni, LingBot-World) --------------------------------------- +flash-attn (LingBot-World) +-------------------------- ``flash-attn`` is only needed for **LingBot-World** with ``reference_numerics: true``, and it is **not** pulled in by ``.[lingbot_world]`` or ``.[all]`` — you install it as a separate step. The reason: flash-attn