From e1c383d315bb3d7778bef3b2aaa40e60cff7371d Mon Sep 17 00:00:00 2001 From: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 7 Jul 2026 01:01:05 -0700 Subject: [PATCH 1/3] perf(multimodal): optimize Qwen vision CUDA graph replay Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> --- .../tokenspeed/runtime/models/qwen3_vision.py | 233 ++++++++++++----- .../runtime/multimodal/encoder_cudagraph.py | 238 +++++++++++++++--- .../test_encoder_cudagraph_grid_rows.py | 223 ++++++++++++++++ 3 files changed, 592 insertions(+), 102 deletions(-) create mode 100644 test/runtime/test_encoder_cudagraph_grid_rows.py diff --git a/python/tokenspeed/runtime/models/qwen3_vision.py b/python/tokenspeed/runtime/models/qwen3_vision.py index 0b79f1a33..945eab0ad 100644 --- a/python/tokenspeed/runtime/models/qwen3_vision.py +++ b/python/tokenspeed/runtime/models/qwen3_vision.py @@ -48,6 +48,20 @@ from tokenspeed.runtime.utils import add_prefix +def _same_grid_row_repeat( + grid_thw: list[list[int]], +) -> tuple[tuple[int, int, int], int] | None: + if not grid_thw: + return None + first = tuple(int(value) for value in grid_thw[0]) + if len(first) != 3: + return None + for row in grid_thw[1:]: + if tuple(int(value) for value in row) != first: + return None + return first, len(grid_thw) + + @lru_cache(maxsize=1024) def _rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor: if isinstance(h, torch.Tensor): @@ -388,20 +402,39 @@ def device(self) -> torch.device: def rot_pos_emb( self, grid_thw: list[list[int]] ) -> tuple[torch.Tensor, torch.Tensor]: - pos_ids = [] + device = self.device + same_row = _same_grid_row_repeat(grid_thw) + if same_row is not None: + (t, h, w), repeats = same_row + cos, sin = self._rot_pos_emb_one(t, h, w, device) + if repeats == 1: + return cos, sin + return cos.repeat(repeats, 1), sin.repeat(repeats, 1) + + cos_outputs = [] + sin_outputs = [] for t, h, w in grid_thw: - base = _rot_pos_ids(h, w, self.spatial_merge_size) - pos_ids.append(base if t == 1 else base.repeat(t, 1)) + cos, sin = self._rot_pos_emb_one(int(t), int(h), int(w), device) + cos_outputs.append(cos) + sin_outputs.append(sin) - pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True) - max_grid_size = max(max(h, w) for _, h, w in grid_thw) + if len(cos_outputs) == 1: + return cos_outputs[0], sin_outputs[0] + return torch.cat(cos_outputs, dim=0), torch.cat(sin_outputs, dim=0) - cos, sin = self._get_rotary_cos_sin(max_grid_size) - - cos_combined = cos[pos_ids].flatten(1) - sin_combined = sin[pos_ids].flatten(1) + def _rot_pos_emb_one( + self, + t: int, + h: int, + w: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + base = _rot_pos_ids(h, w, self.spatial_merge_size) + pos_ids = base if t == 1 else base.repeat(t, 1) + pos_ids = pos_ids.to(device, non_blocking=True) - return cos_combined, sin_combined + cos, sin = self._get_rotary_cos_sin(max(h, w)) + return cos[pos_ids].flatten(1), sin[pos_ids].flatten(1) def _get_rotary_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]: cos_sin = self.rotary_pos_emb.cos_sin_cache[:seqlen].to(self.device) @@ -411,62 +444,109 @@ def fast_pos_embed_interpolate_from_list(self, grid_thw): num_grid_per_side = self.num_grid_per_side m_size = self.spatial_merge_size hidden_dim = self.pos_embed.embedding_dim + device = self.device + dtype = self.dtype + + same_row = _same_grid_row_repeat(grid_thw) + if same_row is not None: + (t, h, w), repeats = same_row + interpolated = self._pos_embed_interpolate_one( + t, + h, + w, + num_grid_per_side, + m_size, + hidden_dim, + device, + dtype, + ) + if repeats == 1: + return interpolated + return interpolated.repeat(repeats, 1) outputs = [] for t, h, w in grid_thw: - h_idxs = torch.linspace( - 0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device - ) - w_idxs = torch.linspace( - 0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device - ) - - h_floor = h_idxs.to(torch.long) - w_floor = w_idxs.to(torch.long) - h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) - w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) - - dh = h_idxs - h_floor - dw = w_idxs - w_floor - - # Create meshgrid view for all h, w vars - dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") - h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") - h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") - - # original computation of weights - # w00 = (1 - dh_grid) * (1 - dw_grid) - # w01 = (1 - dh_grid) * dw_grid - # w10 = dh_grid * (1 - dw_grid) - # w11 = dh_grid * dw_grid - # we reuse w11 here to avoid duplicate - # dh_grid * dw_grid computation - w11 = dh_grid * dw_grid - w10 = dh_grid - w11 - w01 = dw_grid - w11 - w00 = 1 - dh_grid - w01 - - h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid]) - w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid]) - h_grid_idx = h_grid * num_grid_per_side - - indices = (h_grid_idx + w_grid).reshape(4, -1) - weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) - weights = weights.to(dtype=self.dtype) - - embeds = self.pos_embed(indices) - embeds *= weights - combined = embeds.sum(dim=0) - - combined = combined.reshape( - h // m_size, m_size, w // m_size, m_size, hidden_dim + t = int(t) + h = int(h) + w = int(w) + outputs.append( + self._pos_embed_interpolate_one( + t, + h, + w, + num_grid_per_side, + m_size, + hidden_dim, + device, + dtype, + ) ) - combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim) - repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim) - outputs.append(repeated) + if len(outputs) == 1: + return outputs[0] return torch.cat(outputs, dim=0) + def _pos_embed_interpolate_one( + self, + t: int, + h: int, + w: int, + num_grid_per_side: int, + m_size: int, + hidden_dim: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + h_idxs = torch.linspace( + 0, num_grid_per_side - 1, h, dtype=torch.float32, device=device + ) + w_idxs = torch.linspace( + 0, num_grid_per_side - 1, w, dtype=torch.float32, device=device + ) + + h_floor = h_idxs.to(torch.long) + w_floor = w_idxs.to(torch.long) + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # Create meshgrid view for all h, w vars + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") + + # original computation of weights + # w00 = (1 - dh_grid) * (1 - dw_grid) + # w01 = (1 - dh_grid) * dw_grid + # w10 = dh_grid * (1 - dw_grid) + # w11 = dh_grid * dw_grid + # we reuse w11 here to avoid duplicate + # dh_grid * dw_grid computation + w11 = dh_grid * dw_grid + w10 = dh_grid - w11 + w01 = dw_grid - w11 + w00 = 1 - dh_grid - w01 + + h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid]) + w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid]) + h_grid_idx = h_grid * num_grid_per_side + + indices = (h_grid_idx + w_grid).reshape(4, -1) + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) + weights = weights.to(dtype=dtype) + + embeds = self.pos_embed(indices) + embeds *= weights + combined = embeds.sum(dim=0) + + combined = combined.reshape( + h // m_size, m_size, w // m_size, m_size, hidden_dim + ) + combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim) + return combined.expand(t, -1, -1).reshape(-1, hidden_dim) + def compute_cudnn_batch_offsets_packed( self, token_cu_seqlens: np.ndarray, @@ -560,8 +640,8 @@ def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict: grid_thw_list = grid_thw grid_thw_np = np.array(grid_thw, dtype=np.int32) else: - grid_thw_list = grid_thw.tolist() grid_thw_np = grid_thw.cpu().numpy() + grid_thw_list = grid_thw_np.tolist() rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list) @@ -615,12 +695,17 @@ def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict: "sequence_lengths": sequence_lengths, } - def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: - """Capture-safe encoder body: block loop + deepstack mergers + merger. + def _forward_blocks_parts( + self, + x: torch.Tensor, + metadata: dict, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Encoder body returning main and deepstack tensors separately. No host syncs and no data-dependent control flow, so this region is - safe to record into a CUDA graph. ``metadata`` comes from - :meth:`prepare_metadata`; ``x`` from :meth:`prepare_patch_embed`. + safe to record into a CUDA graph when wrapped by ``forward_blocks``. + ``metadata`` comes from :meth:`prepare_metadata`; ``x`` from + :meth:`prepare_patch_embed`. """ cu_seqlens = metadata["cu_seqlens"] rotary_pos_emb_cos = metadata["rotary_pos_emb_cos"] @@ -646,5 +731,25 @@ def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: deepstack_feature_lists.append(deepstack_feature) num_deepstack_captured += 1 x = self.merger(x) + return x, deepstack_feature_lists + + def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: + """Capture-safe encoder body: block loop + deepstack mergers + merger.""" + x, deepstack_feature_lists = self._forward_blocks_parts(x, metadata) # [seq_len, out_hidden_size * (1 + depth_of_deepstack)] + if not deepstack_feature_lists: + return x return torch.cat([x] + deepstack_feature_lists, dim=1) + + def forward_blocks_split( + self, + x: torch.Tensor, + metadata: dict, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Eager encoder body avoiding the main+deepstack concat roundtrip.""" + x, deepstack_feature_lists = self._forward_blocks_parts(x, metadata) + if not deepstack_feature_lists: + return x, None + if len(deepstack_feature_lists) == 1: + return x, deepstack_feature_lists[0] + return x, torch.cat(deepstack_feature_lists, dim=1) diff --git a/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py index 32449e0aa..a5bf7de07 100644 --- a/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py +++ b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py @@ -104,6 +104,8 @@ def prepare_metadata( batch: EncoderCudaGraphBatch, encoder_output_token_budget: int | None, metadata_sequence_budget: int, + *, + pad_cu_seqlens: bool = True, ) -> dict[str, Any]: ... def forward( @@ -131,6 +133,7 @@ class VisionEncoderBatch: tokens: torch.Tensor grid: torch.Tensor out_div: int + grid_rows: list[list[int]] | None = None @property def input_tensors(self) -> dict[str, torch.Tensor]: @@ -138,6 +141,8 @@ def input_tensors(self) -> dict[str, torch.Tensor]: @cached_property def _grid_rows(self) -> list[list[int]]: + if self.grid_rows is not None: + return self.grid_rows return self.grid.tolist() def num_items(self) -> int: @@ -162,15 +167,55 @@ def select(self, indices: list[int]) -> "VisionEncoderBatch": """Sub-batch at ``indices``, preserving order.""" cu = self.cu_input if indices: - rows = torch.cat( - [ - torch.arange(cu[i], cu[i + 1], device=self.tokens.device) - for i in indices - ] - ) + if len(indices) == 1: + item_idx = indices[0] + tokens = self.tokens[cu[item_idx] : cu[item_idx + 1]] + grid = self.grid[item_idx : item_idx + 1] + elif _is_contiguous(indices): + start_idx = indices[0] + end_idx = indices[-1] + 1 + tokens = self.tokens[cu[start_idx] : cu[end_idx]] + grid = self.grid[start_idx:end_idx] + else: + total_rows = sum(cu[i + 1] - cu[i] for i in indices) + if total_rows <= 8192: + rows = _small_row_indices_from_cu(cu, indices, self.tokens.device) + else: + rows = torch.cat( + [ + torch.arange(cu[i], cu[i + 1], device=self.tokens.device) + for i in indices + ] + ) + tokens = self.tokens.index_select(0, rows) + grid_indices = torch.as_tensor( + indices, + dtype=torch.long, + device=self.grid.device, + ) + grid = self.grid.index_select(0, grid_indices) else: - rows = torch.zeros(0, dtype=torch.long, device=self.tokens.device) - return VisionEncoderBatch(self.tokens[rows], self.grid[indices], self.out_div) + tokens = self.tokens[:0] + grid = self.grid[:0] + grid_rows = None + if self.grid_rows is not None: + grid_rows = [self.grid_rows[i] for i in indices] + return VisionEncoderBatch(tokens, grid, self.out_div, grid_rows) + + +def _is_contiguous(indices: list[int]) -> bool: + return all(curr == prev + 1 for prev, curr in zip(indices, indices[1:])) + + +def _small_row_indices_from_cu( + cu: list[int], + indices: list[int], + device: torch.device, +) -> torch.Tensor: + rows: list[int] = [] + for i in indices: + rows.extend(range(cu[i], cu[i + 1])) + return torch.tensor(rows, dtype=torch.long, device=device) @dataclass @@ -178,7 +223,11 @@ class VisionEncoderCudaGraphAdapter: """Adapter for Qwen/Kimi-style ``grid_thw`` vision encoders.""" tower: Any - pre_encode: Callable[[list[Any]], tuple[torch.Tensor, torch.Tensor]] + pre_encode: Callable[ + [list[Any]], + tuple[torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, list[list[int]]], + ] post_encode: Callable[[list[torch.Tensor], torch.Tensor], torch.Tensor] out_div: int merge: int @@ -209,10 +258,9 @@ def synthetic_grid(self, encoder_output_token_budget: int) -> list[list[int]]: b = units // a return [[1, a * self.merge, b * self.merge]] - def pad_cu_seqlens( - self, metadata: dict[str, Any], metadata_sequence_budget: int - ) -> None: - cu = metadata["cu_seqlens"] + def _checked_cu_seqlens_target( + self, cu: torch.Tensor, metadata_sequence_budget: int + ) -> int: target = metadata_sequence_budget + 1 if cu.shape[0] > target: raise RuntimeError( @@ -220,13 +268,25 @@ def pad_cu_seqlens( f"metadata sequences, but the configured limit is " f"{metadata_sequence_budget}" ) + return target + + def pad_cu_seqlens( + self, metadata: dict[str, Any], metadata_sequence_budget: int + ) -> None: + cu = metadata["cu_seqlens"] + target = self._checked_cu_seqlens_target(cu, metadata_sequence_budget) pad = target - cu.shape[0] if pad > 0: metadata["cu_seqlens"] = torch.cat([cu, cu[-1:].expand(pad)]) def batch_from_items(self, items: list[Any]) -> VisionEncoderBatch: - tokens, grid = self.pre_encode(items) - return VisionEncoderBatch(tokens, grid, self.out_div) + pre_encoded = self.pre_encode(items) + if len(pre_encoded) == 3: + tokens, grid, grid_rows = pre_encoded + else: + tokens, grid = pre_encoded + grid_rows = None + return VisionEncoderBatch(tokens, grid, self.out_div, grid_rows) def capture_batch_for_budget( self, @@ -253,15 +313,23 @@ def prepare_metadata( batch: EncoderCudaGraphBatch, encoder_output_token_budget: int | None, metadata_sequence_budget: int, + *, + pad_cu_seqlens: bool = True, ) -> dict[str, Any]: if not isinstance(batch, VisionEncoderBatch): raise TypeError( f"{self.modality_name} encoder cudagraph expected " f"VisionEncoderBatch, got {type(batch).__name__}" ) - metadata = dict(self.tower.prepare_metadata(batch.grid)) + metadata_grid = batch._grid_rows if batch.grid_rows is not None else batch.grid + metadata = dict(self.tower.prepare_metadata(metadata_grid)) if encoder_output_token_budget is not None: - self.pad_cu_seqlens(metadata, metadata_sequence_budget) + if pad_cu_seqlens: + self.pad_cu_seqlens(metadata, metadata_sequence_budget) + else: + self._checked_cu_seqlens_target( + metadata["cu_seqlens"], metadata_sequence_budget + ) # Non-tensor scalar gets baked at capture. Use the per-budget worst # case so replay never exceeds the captured attention max seqlen. metadata["max_seqlen"] = encoder_output_token_budget * self.out_div @@ -474,6 +542,80 @@ def _scatter_output_slices( dest[idx] = sliced.clone() if clone else sliced offset += n_tokens + @staticmethod + def _scatter_cloned_output_slices( + output: torch.Tensor, + indices: list[int], + per_item_encoder_output_tokens: list[int], + dest: dict[int, torch.Tensor], + ) -> None: + """Clone one compact sub-batch output, then scatter views by item. + + Captured graph outputs live in a reusable output buffer. Cloning each + item slice separately creates many small GPU copies under high + concurrency; one contiguous clone per replay keeps item tensors stable + without changing values. + """ + owned = EncoderCudaGraphWrapper._clone_compact_output( + output, + indices, + per_item_encoder_output_tokens, + ) + EncoderCudaGraphWrapper._scatter_output_slices( + owned, + indices, + per_item_encoder_output_tokens, + dest, + ) + + @staticmethod + def _clone_compact_output( + output: torch.Tensor, + indices: list[int], + per_item_encoder_output_tokens: list[int], + ) -> torch.Tensor: + actual_tokens = sum(per_item_encoder_output_tokens[idx] for idx in indices) + return output[:actual_tokens].clone() + + @staticmethod + def _sorted_indices_by_encoder_output_tokens( + per_item_encoder_output_tokens: list[int], + ) -> list[int]: + if len(per_item_encoder_output_tokens) <= 1: + return list(range(len(per_item_encoder_output_tokens))) + if all( + prev <= curr + for prev, curr in zip( + per_item_encoder_output_tokens, + per_item_encoder_output_tokens[1:], + ) + ): + return list(range(len(per_item_encoder_output_tokens))) + return sorted( + range(len(per_item_encoder_output_tokens)), + key=lambda i: per_item_encoder_output_tokens[i], + ) + + @staticmethod + def _copy_prefix_zero_tail(buf: torch.Tensor, src: torch.Tensor) -> None: + n = src.shape[0] + if n == buf.shape[0]: + buf.copy_(src) + return + buf[:n].copy_(src) + buf[n:].zero_() + + @staticmethod + def _copy_prefix_repeat_last_tail(buf: torch.Tensor, src: torch.Tensor) -> None: + n = src.shape[0] + if n == 0: + raise RuntimeError("cannot repeat the last row from an empty tensor") + if n == buf.shape[0]: + buf.copy_(src) + return + buf[:n].copy_(src) + buf[n:].copy_(src[-1:].expand_as(buf[n:])) + def _run_budget_graph( self, batch: EncoderCudaGraphBatch, @@ -510,12 +652,14 @@ def _run_budget_graph( f"shape changed after dim0: capture={tuple(buf.shape)}, " f"replay={tuple(src.shape)}" ) - buf.zero_() - buf[:n].copy_(src) + self._copy_prefix_zero_tail(buf, src) metadata = dict( self.adapter.prepare_metadata( - batch, encoder_output_token_budget, metadata_sequence_budget + batch, + encoder_output_token_budget, + metadata_sequence_budget, + pad_cu_seqlens=False, ) ) replay_buffers = { @@ -546,8 +690,10 @@ def _run_budget_graph( f"has {new.shape[0]} rows, but the captured buffer only " f"has {buf.shape[0]} rows" ) - buf.zero_() - buf[: new.shape[0]].copy_(new) + if key == "cu_seqlens": + self._copy_prefix_repeat_last_tail(buf, new) + else: + self._copy_prefix_zero_tail(buf, new) graph_meta.graph.replay() return graph_meta.output_buffer @@ -570,9 +716,10 @@ def _dispatch(self, batch: EncoderCudaGraphBatch) -> list[torch.Tensor]: per_item_encoder_output_tokens = batch.encoder_output_tokens per_item_metadata_sequences = batch.metadata_sequences - sorted_indices = sorted( - range(num_items), key=lambda i: per_item_encoder_output_tokens[i] + sorted_indices = self._sorted_indices_by_encoder_output_tokens( + per_item_encoder_output_tokens ) + preserve_batch_chunks = sorted_indices == list(range(num_items)) batches: list[tuple[list[int], int | None]] = [] current_batch: list[int] = [] @@ -617,27 +764,42 @@ def _dispatch(self, batch: EncoderCudaGraphBatch) -> list[torch.Tensor]: ) # Packing reorders; restore original order before return. + ordered_outputs: list[torch.Tensor] = [] outputs_by_orig_idx: dict[int, torch.Tensor] = {} for batch_orig_indices, encoder_output_token_budget in batches: sub_batch = batch.select(batch_orig_indices) if encoder_output_token_budget is None: with torch.inference_mode(): raw = self._run_eager(sub_batch) - self._scatter_output_slices( - raw, - batch_orig_indices, - per_item_encoder_output_tokens, - outputs_by_orig_idx, - ) + if preserve_batch_chunks: + ordered_outputs.append(raw) + else: + self._scatter_output_slices( + raw, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + ) else: output = self._run_budget_graph(sub_batch, encoder_output_token_budget) - # clone: output is the shared, reused output_buffer. - self._scatter_output_slices( - output, - batch_orig_indices, - per_item_encoder_output_tokens, - outputs_by_orig_idx, - clone=True, - ) + # output is the shared, reused output_buffer; clone once per + # replay and scatter views out of the owned compact tensor. + if preserve_batch_chunks: + ordered_outputs.append( + self._clone_compact_output( + output, + batch_orig_indices, + per_item_encoder_output_tokens, + ) + ) + else: + self._scatter_cloned_output_slices( + output, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + ) + if preserve_batch_chunks: + return ordered_outputs return [outputs_by_orig_idx[i] for i in range(num_items)] diff --git a/test/runtime/test_encoder_cudagraph_grid_rows.py b/test/runtime/test_encoder_cudagraph_grid_rows.py new file mode 100644 index 000000000..c7f993b27 --- /dev/null +++ b/test/runtime/test_encoder_cudagraph_grid_rows.py @@ -0,0 +1,223 @@ +import torch + +from tokenspeed.runtime.models.qwen3_vision import ( + Qwen3VLMoeVisionModel, +) +from tokenspeed.runtime.multimodal.encoder_cudagraph import ( + EncoderCudaGraphWrapper, + VisionEncoderBatch, + VisionEncoderCudaGraphAdapter, +) + + +def test_vision_encoder_batch_selects_single_and_contiguous_items(): + tokens = torch.arange(5, dtype=torch.float32).reshape(5, 1) + grid = torch.tensor([[1, 1, 1], [1, 2, 1], [1, 1, 2]], dtype=torch.int32) + grid_rows = [[1, 1, 1], [1, 2, 1], [1, 1, 2]] + batch = VisionEncoderBatch(tokens, grid, out_div=1, grid_rows=grid_rows) + + assert batch.encoder_output_tokens == [1, 2, 2] + contiguous = batch.select([0, 1]) + assert torch.equal(contiguous.tokens, tokens[:3]) + assert torch.equal(contiguous.grid, grid[:2]) + assert contiguous.grid_rows == [[1, 1, 1], [1, 2, 1]] + + single = batch.select([2]) + assert torch.equal(single.tokens, tokens[3:5]) + assert torch.equal(single.grid, grid[2:3]) + assert single.grid_rows == [[1, 1, 2]] + + +def test_vision_encoder_batch_selects_noncontiguous_items(): + tokens = torch.arange(10, dtype=torch.float32).reshape(10, 1) + grid = torch.tensor( + [[1, 1, 1], [1, 2, 1], [1, 3, 1], [1, 4, 1]], + dtype=torch.int32, + ) + grid_rows = grid.tolist() + batch = VisionEncoderBatch(tokens, grid, out_div=1, grid_rows=grid_rows) + + selected = batch.select([2, 0]) + + assert torch.equal(selected.tokens, torch.cat([tokens[3:6], tokens[:1]], dim=0)) + assert torch.equal(selected.grid, grid[[2, 0]]) + assert selected.grid_rows == [grid_rows[2], grid_rows[0]] + + +def test_qwen3_vision_rot_pos_emb_repeats_uniform_rows_exactly(): + class Rotary: + cos_sin_cache = torch.empty((1, 2), dtype=torch.float32) + + class DummyTower: + device = torch.device("cpu") + rotary_pos_emb = Rotary() + + def _rot_pos_emb_one(self, *_args): + return ( + torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + torch.tensor([[5.0, 6.0], [7.0, 8.0]]), + ) + + cos, sin = Qwen3VLMoeVisionModel.rot_pos_emb( + DummyTower(), + [[1, 2, 2], [1, 2, 2]], + ) + + assert torch.equal( + cos, + torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]]), + ) + assert torch.equal( + sin, + torch.tensor([[5.0, 6.0], [7.0, 8.0], [5.0, 6.0], [7.0, 8.0]]), + ) + + +def test_qwen3_vision_pos_embed_repeats_uniform_rows_exactly(): + class PosEmbed: + embedding_dim = 2 + + class DummyTower: + num_grid_per_side = 8 + spatial_merge_size = 2 + pos_embed = PosEmbed() + device = torch.device("cpu") + dtype = torch.float32 + + def _pos_embed_interpolate_one(self, *_args): + return torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + tower = DummyTower() + + out = Qwen3VLMoeVisionModel.fast_pos_embed_interpolate_from_list( + tower, + [[1, 2, 2], [1, 2, 2]], + ) + + assert torch.equal( + out, + torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]]), + ) + + +def test_vision_encoder_adapter_passes_precomputed_grid_rows_to_metadata(): + tokens = torch.zeros((2, 1), dtype=torch.float32) + grid = torch.tensor([[1, 2, 1]], dtype=torch.int32) + grid_rows = [[1, 2, 1]] + + class DummyTower: + def __init__(self): + self.seen_grid = None + + def prepare_metadata(self, grid_arg): + self.seen_grid = grid_arg + return {} + + tower = DummyTower() + + def pre_encode(_items): + return tokens, grid, grid_rows + + adapter = VisionEncoderCudaGraphAdapter( + tower=tower, + pre_encode=pre_encode, + post_encode=lambda encoder_outs, _grid: torch.cat(encoder_outs, dim=0), + out_div=1, + merge=1, + input_feature_shape=(1,), + ) + + batch = adapter.batch_from_items([object()]) + adapter.prepare_metadata( + batch, encoder_output_token_budget=None, metadata_sequence_budget=1 + ) + + assert tower.seen_grid is grid_rows + + +def test_vision_encoder_adapter_can_skip_replay_cu_seqlens_padding(): + tokens = torch.zeros((2, 1), dtype=torch.float32) + grid = torch.tensor([[1, 2, 1]], dtype=torch.int32) + + class DummyTower: + def prepare_metadata(self, _grid_arg): + return {"cu_seqlens": torch.tensor([0, 2], dtype=torch.int32)} + + adapter = VisionEncoderCudaGraphAdapter( + tower=DummyTower(), + pre_encode=lambda _items: (tokens, grid), + post_encode=lambda encoder_outs, _grid: torch.cat(encoder_outs, dim=0), + out_div=1, + merge=1, + input_feature_shape=(1,), + ) + batch = VisionEncoderBatch(tokens, grid, out_div=1) + + capture_metadata = adapter.prepare_metadata( + batch, encoder_output_token_budget=4, metadata_sequence_budget=4 + ) + replay_metadata = adapter.prepare_metadata( + batch, + encoder_output_token_budget=4, + metadata_sequence_budget=4, + pad_cu_seqlens=False, + ) + + assert torch.equal(capture_metadata["cu_seqlens"], torch.tensor([0, 2, 2, 2, 2])) + assert torch.equal(replay_metadata["cu_seqlens"], torch.tensor([0, 2])) + + +def test_encoder_cudagraph_copy_prefix_zero_tail_only_clears_unused_rows(): + buf = torch.full((4, 2), 9.0) + src = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + EncoderCudaGraphWrapper._copy_prefix_zero_tail(buf, src) + + assert torch.equal( + buf, + torch.tensor( + [ + [1.0, 2.0], + [3.0, 4.0], + [0.0, 0.0], + [0.0, 0.0], + ] + ), + ) + exact = torch.full_like(src, 9.0) + EncoderCudaGraphWrapper._copy_prefix_zero_tail(exact, src) + assert torch.equal(exact, src) + + +def test_encoder_cudagraph_copy_prefix_repeat_last_tail_matches_padding(): + buf = torch.full((5,), -1, dtype=torch.int32) + src = torch.tensor([0, 2], dtype=torch.int32) + + EncoderCudaGraphWrapper._copy_prefix_repeat_last_tail(buf, src) + + assert torch.equal(buf, torch.tensor([0, 2, 2, 2, 2], dtype=torch.int32)) + exact = torch.full_like(src, -1) + EncoderCudaGraphWrapper._copy_prefix_repeat_last_tail(exact, src) + assert torch.equal(exact, src) + + +def test_encoder_cudagraph_scatter_clones_once_then_scatters_views(): + output = torch.arange(12, dtype=torch.float32).reshape(6, 2) + dest = {} + + EncoderCudaGraphWrapper._scatter_cloned_output_slices( + output, + [0, 1, 2], + [1, 2, 3], + dest, + ) + + assert torch.equal(dest[0], output[0:1]) + assert torch.equal(dest[1], output[1:3]) + assert torch.equal(dest[2], output[3:6]) + assert dest[0].untyped_storage().data_ptr() != output.untyped_storage().data_ptr() + assert ( + dest[0].untyped_storage().data_ptr() + == dest[1].untyped_storage().data_ptr() + == dest[2].untyped_storage().data_ptr() + ) From e8852837059fdde344119dd3fe9ad9dddae77df1 Mon Sep 17 00:00:00 2001 From: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 7 Jul 2026 01:01:16 -0700 Subject: [PATCH 2/3] perf(multimodal): optimize Qwen vision embedding assembly Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> --- python/tokenspeed/runtime/models/qwen3_5.py | 128 +++++--- .../tokenspeed/runtime/models/qwen3_vision.py | 13 - .../tokenspeed/runtime/multimodal/embedder.py | 285 +++++++++++++++--- .../test_qwen_vision_runtime_optimizations.py | 159 ++++++++++ 4 files changed, 500 insertions(+), 85 deletions(-) create mode 100644 test/runtime/test_qwen_vision_runtime_optimizations.py diff --git a/python/tokenspeed/runtime/models/qwen3_5.py b/python/tokenspeed/runtime/models/qwen3_5.py index 10ff57ff3..e8e667151 100644 --- a/python/tokenspeed/runtime/models/qwen3_5.py +++ b/python/tokenspeed/runtime/models/qwen3_5.py @@ -115,6 +115,41 @@ logger = logging.getLogger(__name__) +def _cat_or_storage_view(tensors: list[torch.Tensor]) -> torch.Tensor: + """Return a dim-0 view when Qwen item tensors are adjacent storage slices.""" + if len(tensors) == 1: + return tensors[0] + first = tensors[0] + if first.dim() == 0 or not first.is_contiguous(): + return torch.cat(tensors, dim=0) + + tail_shape = tuple(first.shape[1:]) + stride = first.stride() + storage_ptr = first.untyped_storage().data_ptr() + expected_offset = first.storage_offset() + total_rows = 0 + for tensor in tensors: + if ( + tensor.dtype != first.dtype + or tensor.device != first.device + or tuple(tensor.shape[1:]) != tail_shape + or tensor.stride() != stride + or not tensor.is_contiguous() + or tensor.untyped_storage().data_ptr() != storage_ptr + or tensor.storage_offset() != expected_offset + ): + return torch.cat(tensors, dim=0) + expected_offset += tensor.numel() + total_rows += int(tensor.shape[0]) + + return torch.as_strided( + first, + (total_rows, *tail_shape), + stride, + first.storage_offset(), + ) + + class Qwen3_5GatedDeltaNet(nn.Module): def __init__( self, @@ -1209,18 +1244,16 @@ def pad_input_ids(self, input_ids: list[int], mm_inputs: MultimodalInputs): return pad_input_tokens(input_ids, mm_inputs) def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: - """Eager image encode via the ``pre_encode`` / ``forward_blocks`` / - ``post_encode`` decomposition the cudagraph wrapper uses, so eager - and captured paths share a single source of truth.""" - tokens, grid = self.pre_encode(items) - metadata = self.visual.prepare_metadata(grid) - encoded = self.visual.forward_blocks(tokens, metadata) - return self.post_encode([encoded], grid) + """Encode images using the existing tensor return contract.""" + return self._get_feature(items) def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: - """Eager video encode; the cudagraph path uses the same pre/post hooks.""" - tokens, grid = self.pre_encode(items) - metadata = self.visual.prepare_metadata(grid) + """Encode videos using the existing tensor return contract.""" + return self._get_feature(items) + + def _get_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + tokens, grid, grid_rows = self._pre_encode_with_metadata(items) + metadata = self.visual.prepare_metadata(grid_rows) encoded = self.visual.forward_blocks(tokens, metadata) return self.post_encode([encoded], grid) @@ -1228,40 +1261,61 @@ def pre_encode( self, items: list[MultimodalDataItem], ) -> tuple[torch.Tensor, torch.Tensor]: - """Eager patch-embed before the captured region; returns ``(tokens, grid)``. + """Preserve the existing eager pre-encode return contract.""" + tokens, grid, _ = self._pre_encode_with_metadata(items) + return tokens, grid + + def _pre_encode_with_metadata( + self, + items: list[MultimodalDataItem], + ) -> tuple[torch.Tensor, torch.Tensor, list[list[int]]]: + """Eager patch-embed before the captured region. The grid field is selected per item by modality (``video_grid_thw`` for video, ``image_grid_thw`` otherwise) so a single shared encoder cudagraph - wrapper can serve both image and video batches. + wrapper can serve both image and video batches. The grid rows are + materialized once for patch embedding, metadata, and graph packing. """ - pixel_values = torch.cat([item.feature for item in items], dim=0).type( - self.visual.dtype - ) - grid = torch.concat( - [ - getattr( - item, - ( - "video_grid_thw" - if item.modality == Modality.VIDEO - else "image_grid_thw" - ), - ) - for item in items - ], - dim=0, - ) + if len(items) == 1: + item = items[0] + pixel_values = item.feature + grid = getattr( + item, + ( + "video_grid_thw" + if item.modality == Modality.VIDEO + else "image_grid_thw" + ), + ) + else: + pixel_values = _cat_or_storage_view([item.feature for item in items]) + grid = _cat_or_storage_view( + [ + getattr( + item, + ( + "video_grid_thw" + if item.modality == Modality.VIDEO + else "image_grid_thw" + ), + ) + for item in items + ] + ) if pixel_values.dim() != 2: raise ValueError(f"pixel_values must be 2D, got {pixel_values.dim()}D.") if grid.dim() != 2: raise ValueError(f"grid must be 2D, got {grid.dim()}D.") - x = self.visual.prepare_patch_embed(pixel_values, grid) - return x, grid + grid_rows = grid.tolist() + x = self.visual.prepare_patch_embed(pixel_values, grid_rows) + return x, grid, grid_rows def post_encode( self, encoder_outs: list[torch.Tensor], grid: torch.Tensor ) -> torch.Tensor: """Eager step after the captured region; returns features.""" + if len(encoder_outs) == 1: + return encoder_outs[0] return torch.cat(encoder_outs, dim=0) def _build_encoder_cudagraph_wrapper( @@ -1278,7 +1332,7 @@ def _build_encoder_cudagraph_wrapper( # ``spatial_merge_size ** 2 * budget`` patches. adapter = VisionEncoderCudaGraphAdapter( tower=self.visual, - pre_encode=self.pre_encode, + pre_encode=self._pre_encode_with_metadata, post_encode=self.post_encode, out_div=self.visual.spatial_merge_size**2, merge=self.visual.spatial_merge_size, @@ -1350,8 +1404,14 @@ def forward( text_embedding=self.model.get_input_embeddings(), ctx=multimodal_context, encoders={ - Modality.IMAGE: EncoderSpec(self.image_encoder, deepstack=True), - Modality.VIDEO: EncoderSpec(self.video_encoder, deepstack=True), + Modality.IMAGE: EncoderSpec( + self.image_encoder, + deepstack=self.num_deepstack_embeddings > 0, + ), + Modality.VIDEO: EncoderSpec( + self.video_encoder, + deepstack=self.num_deepstack_embeddings > 0, + ), }, multimodal_model=self, is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(), diff --git a/python/tokenspeed/runtime/models/qwen3_vision.py b/python/tokenspeed/runtime/models/qwen3_vision.py index 945eab0ad..a1816ee3e 100644 --- a/python/tokenspeed/runtime/models/qwen3_vision.py +++ b/python/tokenspeed/runtime/models/qwen3_vision.py @@ -740,16 +740,3 @@ def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: if not deepstack_feature_lists: return x return torch.cat([x] + deepstack_feature_lists, dim=1) - - def forward_blocks_split( - self, - x: torch.Tensor, - metadata: dict, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Eager encoder body avoiding the main+deepstack concat roundtrip.""" - x, deepstack_feature_lists = self._forward_blocks_parts(x, metadata) - if not deepstack_feature_lists: - return x, None - if len(deepstack_feature_lists) == 1: - return x, deepstack_feature_lists[0] - return x, torch.cat(deepstack_feature_lists, dim=1) diff --git a/python/tokenspeed/runtime/multimodal/embedder.py b/python/tokenspeed/runtime/multimodal/embedder.py index b2ae773ce..19e9977ef 100644 --- a/python/tokenspeed/runtime/multimodal/embedder.py +++ b/python/tokenspeed/runtime/multimodal/embedder.py @@ -145,6 +145,80 @@ class ScatterRange: item_src_end: int +def _merged_visual_ranges( + scatter_ranges: list[ScatterRange], + total_rows: int, +) -> list[tuple[int, int]]: + """Return half-open, merged destination ranges filled by vision tokens.""" + if total_rows <= 0: + return [] + ranges = sorted( + (max(0, int(r.flat_dst_start)), min(total_rows, int(r.flat_dst_end) + 1)) + for r in scatter_ranges + ) + ranges = [(start, end) for start, end in ranges if start < end] + if not ranges: + return [] + + merged = [ranges[0]] + for start, end in ranges[1:]: + last_start, last_end = merged[-1] + if start <= last_end: + merged[-1] = (last_start, max(last_end, end)) + else: + merged.append((start, end)) + return merged + + +def _text_token_spans( + total_rows: int, + visual_ranges: list[tuple[int, int]], +) -> list[tuple[int, int]]: + """Return half-open ranges that still need text embedding lookup.""" + spans: list[tuple[int, int]] = [] + cursor = 0 + for start, end in visual_ranges: + if cursor < start: + spans.append((cursor, start)) + cursor = max(cursor, end) + if cursor < total_rows: + spans.append((cursor, total_rows)) + return spans + + +def _coalesced_scatter_ranges( + scatter_ranges: list[ScatterRange], +) -> list[ScatterRange]: + """Merge adjacent scatter work for the same item and contiguous source rows.""" + if len(scatter_ranges) <= 1: + return scatter_ranges + + merged: list[ScatterRange] = [] + scatter_iter = iter(scatter_ranges) + current = next(scatter_iter) + changed = False + for scatter_range in scatter_iter: + if ( + current.item is scatter_range.item + and current.flat_dst_end + 1 == scatter_range.flat_dst_start + and current.item_src_end + 1 == scatter_range.item_src_start + ): + current = ScatterRange( + current.flat_dst_start, + scatter_range.flat_dst_end, + current.item, + current.item_src_start, + scatter_range.item_src_end, + ) + changed = True + continue + + merged.append(current) + current = scatter_range + merged.append(current) + return merged if changed else scatter_ranges + + @dataclass class EncodePlan: """Work to do this prefill iteration. @@ -174,6 +248,94 @@ def _item_token_count(item: MultimodalDataItem) -> int: return sum(end - start + 1 for start, end in item.offsets) +def _ensure_tensor_layout( + tensor: torch.Tensor, *, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + if tensor.dtype == dtype and tensor.device == device: + return tensor + return tensor.to(dtype=dtype, device=device) + + +def _allocate_deepstack_buffer( + input_embeds: torch.Tensor, + num_deepstack: int, + scatter_ranges: list[ScatterRange], + visual_ranges: list[tuple[int, int]] | None = None, + text_spans: list[tuple[int, int]] | None = None, +) -> torch.Tensor: + """Allocate deepstack embeddings for multimodal scatter. + + Qwen3.5 adds the full ``input_deepstack_embeds`` tensor into the first + decoder layers. Text rows must therefore be zero while vision rows are + populated from encoder output. + """ + shape = input_embeds.shape[:-1] + (input_embeds.shape[-1] * num_deepstack,) + if input_embeds.dim() != 2: + return torch.zeros(shape, dtype=input_embeds.dtype, device=input_embeds.device) + + total_rows = int(input_embeds.shape[0]) + if visual_ranges is None: + visual_ranges = _merged_visual_ranges(scatter_ranges, total_rows) + if not visual_ranges: + return torch.zeros(shape, dtype=input_embeds.dtype, device=input_embeds.device) + + deepstack_buffer = torch.empty( + shape, dtype=input_embeds.dtype, device=input_embeds.device + ) + if text_spans is None: + text_spans = _text_token_spans(total_rows, visual_ranges) + for start, end in text_spans: + deepstack_buffer[start:end].zero_() + return deepstack_buffer + + +def _copy_deepstack_rows( + dst: torch.Tensor, + src: torch.Tensor, + *, + dtype: torch.dtype, + device: torch.device, + src_start: int = 0, + src_end: int | None = None, +) -> None: + if src_end is not None: + src = src[src_start:src_end] + dst.copy_(_ensure_tensor_layout(src, dtype=dtype, device=device)) + + +def _dense_deepstack_view_for_full_visual_chunk( + scatter_ranges: list[ScatterRange], + total_rows: int, + *, + expected_width: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + """Return a direct deepstack view when one scatter covers every row.""" + if total_rows <= 0 or len(scatter_ranges) != 1: + return None + + scatter_range = scatter_ranges[0] + if ( + scatter_range.flat_dst_start != 0 + or scatter_range.flat_dst_end != total_rows - 1 + ): + return None + + src = scatter_range.item.encoded_deepstack + if src is None: + return None + + src = src[scatter_range.item_src_start : scatter_range.item_src_end + 1] + if ( + src.dim() != 2 + or int(src.shape[0]) != total_rows + or int(src.shape[-1]) != expected_width + ): + return None + return _ensure_tensor_layout(src, dtype=dtype, device=device) + + # --------------------------------------------------------------------------- # VisionEmbedder # --------------------------------------------------------------------------- @@ -292,17 +454,15 @@ def _plan(self, ctx: MultimodalForwardContext) -> EncodePlan: # the caller hands us. Requests without mm input contribute # nothing but still advance ``base``. base = 0 - for req_idx, mm_inputs in enumerate(ctx.mm_inputs): - if req_idx >= len(ctx.extend_seq_lens) or req_idx >= len( - ctx.extend_prefix_lens - ): - break - seq = ctx.extend_seq_lens[req_idx] + for mm_inputs, seq, prefix in zip( + ctx.mm_inputs, + ctx.extend_seq_lens, + ctx.extend_prefix_lens, + ): if mm_inputs is None or seq <= 0: base += max(seq, 0) continue - prefix = ctx.extend_prefix_lens[req_idx] chunk_start = prefix chunk_end_inc = prefix + seq - 1 @@ -322,7 +482,6 @@ def _plan(self, ctx: MultimodalForwardContext) -> EncodePlan: if canonical is not item: plan.aliases_by_canonical[canonical].append(item) - # src_cursor: start of current subgrid inside item.encoded. src_cursor = 0 for offset_start, offset_end in item.offsets: span = offset_end - offset_start + 1 @@ -362,6 +521,7 @@ def _encode( for modality, items in plan.misses_by_modality.items(): if not items: continue + spec = encoders.get(modality) if spec is None: raise RuntimeError( @@ -386,17 +546,32 @@ def _encode( ) output = output.reshape(-1, output.shape[-1]) - per_item_lens = [_item_token_count(it) for it in items] - per_item_embs = torch.split(output, per_item_lens, dim=0) - - if spec.deepstack: - for item, emb in zip(items, per_item_embs): - main, deep = multimodal_model.separate_deepstack_embeds(emb) - item.encoded = main - item.encoded_deepstack = deep - else: - for item, emb in zip(items, per_item_embs): - item.encoded = emb + output_rows = int(output.shape[0]) + cursor = 0 + per_item_lens = [] if LOG_MM_TIMING else None + for item in items: + n_tokens = _item_token_count(item) + if per_item_lens is not None: + per_item_lens.append(n_tokens) + next_cursor = cursor + n_tokens + if next_cursor > output_rows: + raise RuntimeError( + "VisionEmbedder encoder output is shorter than " + f"planned tokens for {modality}: needed " + f"{next_cursor}, got {output_rows}" + ) + item.encoded = output[cursor:next_cursor] + if spec.deepstack: + item.encoded, item.encoded_deepstack = ( + multimodal_model.separate_deepstack_embeds(item.encoded) + ) + cursor = next_cursor + if cursor != output_rows: + raise RuntimeError( + "VisionEmbedder encoder output token count does not match " + f"planned tokens for {modality}: planned {cursor}, " + f"got {output_rows}" + ) if LOG_MM_TIMING: logger.info( "mm_timing encoder_ms modality=%s items=%d " @@ -432,24 +607,47 @@ def _assemble( encoders: dict[Modality, EncoderSpec], multimodal_model: nn.Module, ) -> tuple[torch.Tensor, dict[str, Any]]: - # Placeholder positions hold large content-derived IDs that exceed - # vocab_size; the lookup we run here is overwritten for those rows - # by the scatter below, but the lookup still needs valid indices. - vocab_size = text_embedding.num_embeddings - safe_ids = input_ids.clamp(min=0, max=vocab_size - 1) - input_embeds = text_embedding(safe_ids) - + # Placeholder positions may hold content-derived IDs beyond vocab_size. + # Clamp them for the lookup; vision rows are overwritten below. + scatter_ranges = _coalesced_scatter_ranges(plan.scatter_ranges) + vocab_size = getattr(text_embedding, "num_embeddings") + input_embeds = text_embedding(input_ids.clamp(min=0, max=int(vocab_size) - 1)) + + has_deepstack = any(spec.deepstack for spec in encoders.values()) + visual_ranges = None + text_spans = None + if has_deepstack and input_ids.dim() == 1: + input_rows = int(input_ids.shape[0]) + visual_ranges = _merged_visual_ranges( + scatter_ranges, + input_rows, + ) + if visual_ranges: + text_spans = _text_token_spans(input_rows, visual_ranges) kwargs: dict[str, Any] = {} deepstack_buffer: torch.Tensor | None = None - if any(spec.deepstack for spec in encoders.values()): + deepstack_buffer_is_direct = False + if has_deepstack: num_deepstack = len(multimodal_model.deepstack_visual_indexes) - shape = input_embeds.shape[:-1] + (input_embeds.shape[-1] * num_deepstack,) - deepstack_buffer = torch.zeros( - shape, dtype=input_embeds.dtype, device=input_embeds.device + deepstack_buffer = _dense_deepstack_view_for_full_visual_chunk( + scatter_ranges, + int(input_embeds.shape[0]) if input_embeds.dim() == 2 else 0, + expected_width=int(input_embeds.shape[-1]) * num_deepstack, + dtype=input_embeds.dtype, + device=input_embeds.device, ) + deepstack_buffer_is_direct = deepstack_buffer is not None + if deepstack_buffer is None: + deepstack_buffer = _allocate_deepstack_buffer( + input_embeds, + num_deepstack, + scatter_ranges, + visual_ranges, + text_spans, + ) kwargs["input_deepstack_embeds"] = deepstack_buffer - for r in plan.scatter_ranges: + for r in scatter_ranges: main = r.item.encoded if main is None: raise RuntimeError( @@ -457,17 +655,28 @@ def _assemble( "encoded tensor after _encode; this is a bug" ) src = main[r.item_src_start : r.item_src_end + 1] - input_embeds[r.flat_dst_start : r.flat_dst_end + 1] = src.to( - dtype=input_embeds.dtype, device=input_embeds.device + input_embeds[r.flat_dst_start : r.flat_dst_end + 1] = _ensure_tensor_layout( + src, + dtype=input_embeds.dtype, + device=input_embeds.device, ) + if deepstack_buffer_is_direct: + continue + if deepstack_buffer is not None and r.item.encoded_deepstack is not None: - deep_src = r.item.encoded_deepstack[ - r.item_src_start : r.item_src_end + 1 - ] - deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1] = deep_src.to( - dtype=input_embeds.dtype, device=input_embeds.device + deep_src = r.item.encoded_deepstack + deep_dst = deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1] + _copy_deepstack_rows( + deep_dst, + deep_src, + dtype=input_embeds.dtype, + device=input_embeds.device, + src_start=r.item_src_start, + src_end=r.item_src_end + 1, ) + elif deepstack_buffer is not None: + deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1].zero_() return input_embeds, kwargs diff --git a/test/runtime/test_qwen_vision_runtime_optimizations.py b/test/runtime/test_qwen_vision_runtime_optimizations.py new file mode 100644 index 000000000..4aee3e90d --- /dev/null +++ b/test/runtime/test_qwen_vision_runtime_optimizations.py @@ -0,0 +1,159 @@ +import pytest +import torch + +from tokenspeed.runtime.multimodal.embedder import ( + EncodePlan, + EncoderSpec, + ScatterRange, + VisionEmbedder, + _coalesced_scatter_ranges, + _merged_visual_ranges, + _text_token_spans, +) +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalForwardContext, + MultimodalInputs, +) + + +def test_visual_ranges_are_coalesced_for_embedding_assembly(): + item = MultimodalDataItem(modality=Modality.VIDEO) + other = MultimodalDataItem(modality=Modality.VIDEO) + scatter_ranges = [ + ScatterRange(1, 2, item, 0, 1), + ScatterRange(3, 5, item, 2, 4), + ScatterRange(7, 8, other, 0, 1), + ] + + coalesced = _coalesced_scatter_ranges(scatter_ranges) + visual_ranges = _merged_visual_ranges(coalesced, total_rows=10) + + assert coalesced == [ + ScatterRange(1, 5, item, 0, 4), + ScatterRange(7, 8, other, 0, 1), + ] + assert visual_ranges == [(1, 6), (7, 9)] + assert _text_token_spans(10, visual_ranges) == [(0, 1), (6, 7), (9, 10)] + + +def test_embedding_assembly_clamps_placeholder_ids_before_visual_overwrite(): + embedding = torch.nn.Embedding(10, 3) + with torch.no_grad(): + embedding.weight.copy_(torch.arange(30, dtype=torch.float32).reshape(10, 3)) + item = MultimodalDataItem( + modality=Modality.IMAGE, + encoded=torch.tensor( + [[100.0, 101.0, 102.0], [110.0, 111.0, 112.0]], + dtype=torch.float32, + ), + ) + plan = EncodePlan( + scatter_ranges=[ScatterRange(2, 3, item, 0, 1)], + ) + input_ids = torch.tensor([1, 2, 999, 999, 3], dtype=torch.long) + + input_embeds, kwargs = VisionEmbedder()._assemble( + input_ids=input_ids, + text_embedding=embedding, + plan=plan, + encoders={Modality.IMAGE: EncoderSpec(lambda _: torch.empty(0))}, + multimodal_model=None, + ) + + expected = embedding(input_ids.clamp(max=embedding.num_embeddings - 1)) + expected[2:4] = item.encoded + assert kwargs == {} + assert torch.equal(input_embeds, expected) + + +@pytest.mark.parametrize("modality", [Modality.IMAGE, Modality.VIDEO]) +def test_deepstack_combined_contract_for_chunked_prefill(modality): + embedding = torch.nn.Embedding(10, 3, dtype=torch.bfloat16) + with torch.no_grad(): + embedding.weight.copy_(torch.arange(30, dtype=torch.bfloat16).reshape(10, 3)) + + main = torch.arange(12, dtype=torch.bfloat16).reshape(4, 3) + deepstack = torch.arange(24, dtype=torch.bfloat16).reshape(4, 6) + combined = torch.cat([main, deepstack], dim=1) + input_ids = torch.tensor([999, 1, 2, 999, 999], dtype=torch.long) + + class Model: + deepstack_visual_indexes = [1, 2] + + @staticmethod + def separate_deepstack_embeds(output): + return output[:, :3], output[:, 3:] + + def run(): + item = MultimodalDataItem( + modality=modality, + offsets=[(2, 3), (6, 7)], + feature=torch.empty(0), + ) + context = MultimodalForwardContext( + mm_inputs=[MultimodalInputs(mm_items=[item])], + extend_prefix_lens=[3], + extend_seq_lens=[5], + ) + return VisionEmbedder().apply( + input_ids=input_ids, + text_embedding=embedding, + ctx=context, + encoders={modality: EncoderSpec(lambda _items: combined, deepstack=True)}, + multimodal_model=Model(), + ) + + input_embeds, kwargs = run() + + expected = embedding(input_ids.clamp(max=embedding.num_embeddings - 1)) + expected[0] = main[1] + expected[3:] = main[2:] + assert torch.equal(input_embeds, expected) + assert torch.equal(kwargs["input_deepstack_embeds"][0], deepstack[1]) + assert torch.equal(kwargs["input_deepstack_embeds"][3:], deepstack[2:]) + + +def test_dense_deepstack_reuses_full_visual_tensor(): + item = MultimodalDataItem( + modality=Modality.VIDEO, + encoded=torch.arange(6, dtype=torch.float32).reshape(2, 3), + encoded_deepstack=torch.arange(6, 12, dtype=torch.float32).reshape(2, 3), + ) + plan = EncodePlan(scatter_ranges=[ScatterRange(0, 1, item, 0, 1)]) + + class Model: + deepstack_visual_indexes = [8] + + _, kwargs = VisionEmbedder()._assemble( + input_ids=torch.tensor([999, 999], dtype=torch.long), + text_embedding=torch.nn.Embedding(10, 3), + plan=plan, + encoders={ + Modality.VIDEO: EncoderSpec(lambda _: torch.empty(0), deepstack=True) + }, + multimodal_model=Model(), + ) + + deepstack = kwargs["input_deepstack_embeds"] + assert torch.equal(deepstack, item.encoded_deepstack) + assert ( + deepstack.untyped_storage().data_ptr() + == item.encoded_deepstack.untyped_storage().data_ptr() + ) + + +def test_encoder_output_must_match_planned_tokens(): + item = MultimodalDataItem(modality=Modality.IMAGE, offsets=[(0, 0)]) + plan = EncodePlan() + plan.misses_by_modality[Modality.IMAGE] = [item] + output = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + with pytest.raises(RuntimeError, match="does not match planned tokens"): + VisionEmbedder()._encode( + plan=plan, + encoders={Modality.IMAGE: EncoderSpec(lambda _: output)}, + multimodal_model=None, + device=torch.device("cpu"), + ) From 68e23cb62824d9ac551c09555b5b612d3f2bd735 Mon Sep 17 00:00:00 2001 From: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:43:05 -0700 Subject: [PATCH 3/3] refactor(multimodal): narrow Qwen vision runtime scope Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com> --- python/tokenspeed/runtime/models/qwen3_5.py | 101 ++----- .../tokenspeed/runtime/models/qwen3_vision.py | 220 ++++---------- .../tokenspeed/runtime/multimodal/embedder.py | 285 +++--------------- .../runtime/multimodal/encoder_cudagraph.py | 238 +++------------ .../test_encoder_cudagraph_grid_rows.py | 223 -------------- test/runtime/test_qwen35_vision_assembly.py | 69 +++++ .../test_qwen_vision_runtime_optimizations.py | 159 ---------- 7 files changed, 233 insertions(+), 1062 deletions(-) delete mode 100644 test/runtime/test_encoder_cudagraph_grid_rows.py create mode 100644 test/runtime/test_qwen35_vision_assembly.py delete mode 100644 test/runtime/test_qwen_vision_runtime_optimizations.py diff --git a/python/tokenspeed/runtime/models/qwen3_5.py b/python/tokenspeed/runtime/models/qwen3_5.py index e8e667151..78274b7cd 100644 --- a/python/tokenspeed/runtime/models/qwen3_5.py +++ b/python/tokenspeed/runtime/models/qwen3_5.py @@ -115,41 +115,6 @@ logger = logging.getLogger(__name__) -def _cat_or_storage_view(tensors: list[torch.Tensor]) -> torch.Tensor: - """Return a dim-0 view when Qwen item tensors are adjacent storage slices.""" - if len(tensors) == 1: - return tensors[0] - first = tensors[0] - if first.dim() == 0 or not first.is_contiguous(): - return torch.cat(tensors, dim=0) - - tail_shape = tuple(first.shape[1:]) - stride = first.stride() - storage_ptr = first.untyped_storage().data_ptr() - expected_offset = first.storage_offset() - total_rows = 0 - for tensor in tensors: - if ( - tensor.dtype != first.dtype - or tensor.device != first.device - or tuple(tensor.shape[1:]) != tail_shape - or tensor.stride() != stride - or not tensor.is_contiguous() - or tensor.untyped_storage().data_ptr() != storage_ptr - or tensor.storage_offset() != expected_offset - ): - return torch.cat(tensors, dim=0) - expected_offset += tensor.numel() - total_rows += int(tensor.shape[0]) - - return torch.as_strided( - first, - (total_rows, *tail_shape), - stride, - first.storage_offset(), - ) - - class Qwen3_5GatedDeltaNet(nn.Module): def __init__( self, @@ -1244,16 +1209,18 @@ def pad_input_ids(self, input_ids: list[int], mm_inputs: MultimodalInputs): return pad_input_tokens(input_ids, mm_inputs) def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: - """Encode images using the existing tensor return contract.""" - return self._get_feature(items) + """Eager image encode via the ``pre_encode`` / ``forward_blocks`` / + ``post_encode`` decomposition the cudagraph wrapper uses, so eager + and captured paths share a single source of truth.""" + tokens, grid = self.pre_encode(items) + metadata = self.visual.prepare_metadata(grid) + encoded = self.visual.forward_blocks(tokens, metadata) + return self.post_encode([encoded], grid) def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: - """Encode videos using the existing tensor return contract.""" - return self._get_feature(items) - - def _get_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: - tokens, grid, grid_rows = self._pre_encode_with_metadata(items) - metadata = self.visual.prepare_metadata(grid_rows) + """Eager video encode; the cudagraph path uses the same pre/post hooks.""" + tokens, grid = self.pre_encode(items) + metadata = self.visual.prepare_metadata(grid) encoded = self.visual.forward_blocks(tokens, metadata) return self.post_encode([encoded], grid) @@ -1261,25 +1228,15 @@ def pre_encode( self, items: list[MultimodalDataItem], ) -> tuple[torch.Tensor, torch.Tensor]: - """Preserve the existing eager pre-encode return contract.""" - tokens, grid, _ = self._pre_encode_with_metadata(items) - return tokens, grid - - def _pre_encode_with_metadata( - self, - items: list[MultimodalDataItem], - ) -> tuple[torch.Tensor, torch.Tensor, list[list[int]]]: - """Eager patch-embed before the captured region. + """Eager patch-embed before the captured region; returns ``(tokens, grid)``. The grid field is selected per item by modality (``video_grid_thw`` for video, ``image_grid_thw`` otherwise) so a single shared encoder cudagraph - wrapper can serve both image and video batches. The grid rows are - materialized once for patch embedding, metadata, and graph packing. + wrapper can serve both image and video batches. """ - if len(items) == 1: - item = items[0] - pixel_values = item.feature - grid = getattr( + features = [item.feature for item in items] + grids = [ + getattr( item, ( "video_grid_thw" @@ -1287,28 +1244,18 @@ def _pre_encode_with_metadata( else "image_grid_thw" ), ) - else: - pixel_values = _cat_or_storage_view([item.feature for item in items]) - grid = _cat_or_storage_view( - [ - getattr( - item, - ( - "video_grid_thw" - if item.modality == Modality.VIDEO - else "image_grid_thw" - ), - ) - for item in items - ] - ) + for item in items + ] + pixel_values = ( + features[0] if len(features) == 1 else torch.cat(features, dim=0) + ).type(self.visual.dtype) + grid = grids[0] if len(grids) == 1 else torch.concat(grids, dim=0) if pixel_values.dim() != 2: raise ValueError(f"pixel_values must be 2D, got {pixel_values.dim()}D.") if grid.dim() != 2: raise ValueError(f"grid must be 2D, got {grid.dim()}D.") - grid_rows = grid.tolist() - x = self.visual.prepare_patch_embed(pixel_values, grid_rows) - return x, grid, grid_rows + x = self.visual.prepare_patch_embed(pixel_values, grid) + return x, grid def post_encode( self, encoder_outs: list[torch.Tensor], grid: torch.Tensor @@ -1332,7 +1279,7 @@ def _build_encoder_cudagraph_wrapper( # ``spatial_merge_size ** 2 * budget`` patches. adapter = VisionEncoderCudaGraphAdapter( tower=self.visual, - pre_encode=self._pre_encode_with_metadata, + pre_encode=self.pre_encode, post_encode=self.post_encode, out_div=self.visual.spatial_merge_size**2, merge=self.visual.spatial_merge_size, diff --git a/python/tokenspeed/runtime/models/qwen3_vision.py b/python/tokenspeed/runtime/models/qwen3_vision.py index a1816ee3e..0b79f1a33 100644 --- a/python/tokenspeed/runtime/models/qwen3_vision.py +++ b/python/tokenspeed/runtime/models/qwen3_vision.py @@ -48,20 +48,6 @@ from tokenspeed.runtime.utils import add_prefix -def _same_grid_row_repeat( - grid_thw: list[list[int]], -) -> tuple[tuple[int, int, int], int] | None: - if not grid_thw: - return None - first = tuple(int(value) for value in grid_thw[0]) - if len(first) != 3: - return None - for row in grid_thw[1:]: - if tuple(int(value) for value in row) != first: - return None - return first, len(grid_thw) - - @lru_cache(maxsize=1024) def _rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor: if isinstance(h, torch.Tensor): @@ -402,39 +388,20 @@ def device(self) -> torch.device: def rot_pos_emb( self, grid_thw: list[list[int]] ) -> tuple[torch.Tensor, torch.Tensor]: - device = self.device - same_row = _same_grid_row_repeat(grid_thw) - if same_row is not None: - (t, h, w), repeats = same_row - cos, sin = self._rot_pos_emb_one(t, h, w, device) - if repeats == 1: - return cos, sin - return cos.repeat(repeats, 1), sin.repeat(repeats, 1) - - cos_outputs = [] - sin_outputs = [] + pos_ids = [] for t, h, w in grid_thw: - cos, sin = self._rot_pos_emb_one(int(t), int(h), int(w), device) - cos_outputs.append(cos) - sin_outputs.append(sin) + base = _rot_pos_ids(h, w, self.spatial_merge_size) + pos_ids.append(base if t == 1 else base.repeat(t, 1)) - if len(cos_outputs) == 1: - return cos_outputs[0], sin_outputs[0] - return torch.cat(cos_outputs, dim=0), torch.cat(sin_outputs, dim=0) + pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True) + max_grid_size = max(max(h, w) for _, h, w in grid_thw) - def _rot_pos_emb_one( - self, - t: int, - h: int, - w: int, - device: torch.device, - ) -> tuple[torch.Tensor, torch.Tensor]: - base = _rot_pos_ids(h, w, self.spatial_merge_size) - pos_ids = base if t == 1 else base.repeat(t, 1) - pos_ids = pos_ids.to(device, non_blocking=True) + cos, sin = self._get_rotary_cos_sin(max_grid_size) + + cos_combined = cos[pos_ids].flatten(1) + sin_combined = sin[pos_ids].flatten(1) - cos, sin = self._get_rotary_cos_sin(max(h, w)) - return cos[pos_ids].flatten(1), sin[pos_ids].flatten(1) + return cos_combined, sin_combined def _get_rotary_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]: cos_sin = self.rotary_pos_emb.cos_sin_cache[:seqlen].to(self.device) @@ -444,108 +411,61 @@ def fast_pos_embed_interpolate_from_list(self, grid_thw): num_grid_per_side = self.num_grid_per_side m_size = self.spatial_merge_size hidden_dim = self.pos_embed.embedding_dim - device = self.device - dtype = self.dtype - - same_row = _same_grid_row_repeat(grid_thw) - if same_row is not None: - (t, h, w), repeats = same_row - interpolated = self._pos_embed_interpolate_one( - t, - h, - w, - num_grid_per_side, - m_size, - hidden_dim, - device, - dtype, - ) - if repeats == 1: - return interpolated - return interpolated.repeat(repeats, 1) outputs = [] for t, h, w in grid_thw: - t = int(t) - h = int(h) - w = int(w) - outputs.append( - self._pos_embed_interpolate_one( - t, - h, - w, - num_grid_per_side, - m_size, - hidden_dim, - device, - dtype, - ) + h_idxs = torch.linspace( + 0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device + ) + w_idxs = torch.linspace( + 0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device ) - if len(outputs) == 1: - return outputs[0] - return torch.cat(outputs, dim=0) - - def _pos_embed_interpolate_one( - self, - t: int, - h: int, - w: int, - num_grid_per_side: int, - m_size: int, - hidden_dim: int, - device: torch.device, - dtype: torch.dtype, - ) -> torch.Tensor: - h_idxs = torch.linspace( - 0, num_grid_per_side - 1, h, dtype=torch.float32, device=device - ) - w_idxs = torch.linspace( - 0, num_grid_per_side - 1, w, dtype=torch.float32, device=device - ) + h_floor = h_idxs.to(torch.long) + w_floor = w_idxs.to(torch.long) + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # Create meshgrid view for all h, w vars + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") + + # original computation of weights + # w00 = (1 - dh_grid) * (1 - dw_grid) + # w01 = (1 - dh_grid) * dw_grid + # w10 = dh_grid * (1 - dw_grid) + # w11 = dh_grid * dw_grid + # we reuse w11 here to avoid duplicate + # dh_grid * dw_grid computation + w11 = dh_grid * dw_grid + w10 = dh_grid - w11 + w01 = dw_grid - w11 + w00 = 1 - dh_grid - w01 + + h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid]) + w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid]) + h_grid_idx = h_grid * num_grid_per_side + + indices = (h_grid_idx + w_grid).reshape(4, -1) + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) + weights = weights.to(dtype=self.dtype) + + embeds = self.pos_embed(indices) + embeds *= weights + combined = embeds.sum(dim=0) + + combined = combined.reshape( + h // m_size, m_size, w // m_size, m_size, hidden_dim + ) + combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim) + repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim) + outputs.append(repeated) - h_floor = h_idxs.to(torch.long) - w_floor = w_idxs.to(torch.long) - h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) - w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) - - dh = h_idxs - h_floor - dw = w_idxs - w_floor - - # Create meshgrid view for all h, w vars - dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") - h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") - h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") - - # original computation of weights - # w00 = (1 - dh_grid) * (1 - dw_grid) - # w01 = (1 - dh_grid) * dw_grid - # w10 = dh_grid * (1 - dw_grid) - # w11 = dh_grid * dw_grid - # we reuse w11 here to avoid duplicate - # dh_grid * dw_grid computation - w11 = dh_grid * dw_grid - w10 = dh_grid - w11 - w01 = dw_grid - w11 - w00 = 1 - dh_grid - w01 - - h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid]) - w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid]) - h_grid_idx = h_grid * num_grid_per_side - - indices = (h_grid_idx + w_grid).reshape(4, -1) - weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) - weights = weights.to(dtype=dtype) - - embeds = self.pos_embed(indices) - embeds *= weights - combined = embeds.sum(dim=0) - - combined = combined.reshape( - h // m_size, m_size, w // m_size, m_size, hidden_dim - ) - combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim) - return combined.expand(t, -1, -1).reshape(-1, hidden_dim) + return torch.cat(outputs, dim=0) def compute_cudnn_batch_offsets_packed( self, @@ -640,8 +560,8 @@ def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict: grid_thw_list = grid_thw grid_thw_np = np.array(grid_thw, dtype=np.int32) else: + grid_thw_list = grid_thw.tolist() grid_thw_np = grid_thw.cpu().numpy() - grid_thw_list = grid_thw_np.tolist() rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list) @@ -695,17 +615,12 @@ def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict: "sequence_lengths": sequence_lengths, } - def _forward_blocks_parts( - self, - x: torch.Tensor, - metadata: dict, - ) -> tuple[torch.Tensor, list[torch.Tensor]]: - """Encoder body returning main and deepstack tensors separately. + def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: + """Capture-safe encoder body: block loop + deepstack mergers + merger. No host syncs and no data-dependent control flow, so this region is - safe to record into a CUDA graph when wrapped by ``forward_blocks``. - ``metadata`` comes from :meth:`prepare_metadata`; ``x`` from - :meth:`prepare_patch_embed`. + safe to record into a CUDA graph. ``metadata`` comes from + :meth:`prepare_metadata`; ``x`` from :meth:`prepare_patch_embed`. """ cu_seqlens = metadata["cu_seqlens"] rotary_pos_emb_cos = metadata["rotary_pos_emb_cos"] @@ -731,12 +646,5 @@ def _forward_blocks_parts( deepstack_feature_lists.append(deepstack_feature) num_deepstack_captured += 1 x = self.merger(x) - return x, deepstack_feature_lists - - def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: - """Capture-safe encoder body: block loop + deepstack mergers + merger.""" - x, deepstack_feature_lists = self._forward_blocks_parts(x, metadata) # [seq_len, out_hidden_size * (1 + depth_of_deepstack)] - if not deepstack_feature_lists: - return x return torch.cat([x] + deepstack_feature_lists, dim=1) diff --git a/python/tokenspeed/runtime/multimodal/embedder.py b/python/tokenspeed/runtime/multimodal/embedder.py index 19e9977ef..b2ae773ce 100644 --- a/python/tokenspeed/runtime/multimodal/embedder.py +++ b/python/tokenspeed/runtime/multimodal/embedder.py @@ -145,80 +145,6 @@ class ScatterRange: item_src_end: int -def _merged_visual_ranges( - scatter_ranges: list[ScatterRange], - total_rows: int, -) -> list[tuple[int, int]]: - """Return half-open, merged destination ranges filled by vision tokens.""" - if total_rows <= 0: - return [] - ranges = sorted( - (max(0, int(r.flat_dst_start)), min(total_rows, int(r.flat_dst_end) + 1)) - for r in scatter_ranges - ) - ranges = [(start, end) for start, end in ranges if start < end] - if not ranges: - return [] - - merged = [ranges[0]] - for start, end in ranges[1:]: - last_start, last_end = merged[-1] - if start <= last_end: - merged[-1] = (last_start, max(last_end, end)) - else: - merged.append((start, end)) - return merged - - -def _text_token_spans( - total_rows: int, - visual_ranges: list[tuple[int, int]], -) -> list[tuple[int, int]]: - """Return half-open ranges that still need text embedding lookup.""" - spans: list[tuple[int, int]] = [] - cursor = 0 - for start, end in visual_ranges: - if cursor < start: - spans.append((cursor, start)) - cursor = max(cursor, end) - if cursor < total_rows: - spans.append((cursor, total_rows)) - return spans - - -def _coalesced_scatter_ranges( - scatter_ranges: list[ScatterRange], -) -> list[ScatterRange]: - """Merge adjacent scatter work for the same item and contiguous source rows.""" - if len(scatter_ranges) <= 1: - return scatter_ranges - - merged: list[ScatterRange] = [] - scatter_iter = iter(scatter_ranges) - current = next(scatter_iter) - changed = False - for scatter_range in scatter_iter: - if ( - current.item is scatter_range.item - and current.flat_dst_end + 1 == scatter_range.flat_dst_start - and current.item_src_end + 1 == scatter_range.item_src_start - ): - current = ScatterRange( - current.flat_dst_start, - scatter_range.flat_dst_end, - current.item, - current.item_src_start, - scatter_range.item_src_end, - ) - changed = True - continue - - merged.append(current) - current = scatter_range - merged.append(current) - return merged if changed else scatter_ranges - - @dataclass class EncodePlan: """Work to do this prefill iteration. @@ -248,94 +174,6 @@ def _item_token_count(item: MultimodalDataItem) -> int: return sum(end - start + 1 for start, end in item.offsets) -def _ensure_tensor_layout( - tensor: torch.Tensor, *, dtype: torch.dtype, device: torch.device -) -> torch.Tensor: - if tensor.dtype == dtype and tensor.device == device: - return tensor - return tensor.to(dtype=dtype, device=device) - - -def _allocate_deepstack_buffer( - input_embeds: torch.Tensor, - num_deepstack: int, - scatter_ranges: list[ScatterRange], - visual_ranges: list[tuple[int, int]] | None = None, - text_spans: list[tuple[int, int]] | None = None, -) -> torch.Tensor: - """Allocate deepstack embeddings for multimodal scatter. - - Qwen3.5 adds the full ``input_deepstack_embeds`` tensor into the first - decoder layers. Text rows must therefore be zero while vision rows are - populated from encoder output. - """ - shape = input_embeds.shape[:-1] + (input_embeds.shape[-1] * num_deepstack,) - if input_embeds.dim() != 2: - return torch.zeros(shape, dtype=input_embeds.dtype, device=input_embeds.device) - - total_rows = int(input_embeds.shape[0]) - if visual_ranges is None: - visual_ranges = _merged_visual_ranges(scatter_ranges, total_rows) - if not visual_ranges: - return torch.zeros(shape, dtype=input_embeds.dtype, device=input_embeds.device) - - deepstack_buffer = torch.empty( - shape, dtype=input_embeds.dtype, device=input_embeds.device - ) - if text_spans is None: - text_spans = _text_token_spans(total_rows, visual_ranges) - for start, end in text_spans: - deepstack_buffer[start:end].zero_() - return deepstack_buffer - - -def _copy_deepstack_rows( - dst: torch.Tensor, - src: torch.Tensor, - *, - dtype: torch.dtype, - device: torch.device, - src_start: int = 0, - src_end: int | None = None, -) -> None: - if src_end is not None: - src = src[src_start:src_end] - dst.copy_(_ensure_tensor_layout(src, dtype=dtype, device=device)) - - -def _dense_deepstack_view_for_full_visual_chunk( - scatter_ranges: list[ScatterRange], - total_rows: int, - *, - expected_width: int, - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor | None: - """Return a direct deepstack view when one scatter covers every row.""" - if total_rows <= 0 or len(scatter_ranges) != 1: - return None - - scatter_range = scatter_ranges[0] - if ( - scatter_range.flat_dst_start != 0 - or scatter_range.flat_dst_end != total_rows - 1 - ): - return None - - src = scatter_range.item.encoded_deepstack - if src is None: - return None - - src = src[scatter_range.item_src_start : scatter_range.item_src_end + 1] - if ( - src.dim() != 2 - or int(src.shape[0]) != total_rows - or int(src.shape[-1]) != expected_width - ): - return None - return _ensure_tensor_layout(src, dtype=dtype, device=device) - - # --------------------------------------------------------------------------- # VisionEmbedder # --------------------------------------------------------------------------- @@ -454,15 +292,17 @@ def _plan(self, ctx: MultimodalForwardContext) -> EncodePlan: # the caller hands us. Requests without mm input contribute # nothing but still advance ``base``. base = 0 - for mm_inputs, seq, prefix in zip( - ctx.mm_inputs, - ctx.extend_seq_lens, - ctx.extend_prefix_lens, - ): + for req_idx, mm_inputs in enumerate(ctx.mm_inputs): + if req_idx >= len(ctx.extend_seq_lens) or req_idx >= len( + ctx.extend_prefix_lens + ): + break + seq = ctx.extend_seq_lens[req_idx] if mm_inputs is None or seq <= 0: base += max(seq, 0) continue + prefix = ctx.extend_prefix_lens[req_idx] chunk_start = prefix chunk_end_inc = prefix + seq - 1 @@ -482,6 +322,7 @@ def _plan(self, ctx: MultimodalForwardContext) -> EncodePlan: if canonical is not item: plan.aliases_by_canonical[canonical].append(item) + # src_cursor: start of current subgrid inside item.encoded. src_cursor = 0 for offset_start, offset_end in item.offsets: span = offset_end - offset_start + 1 @@ -521,7 +362,6 @@ def _encode( for modality, items in plan.misses_by_modality.items(): if not items: continue - spec = encoders.get(modality) if spec is None: raise RuntimeError( @@ -546,32 +386,17 @@ def _encode( ) output = output.reshape(-1, output.shape[-1]) - output_rows = int(output.shape[0]) - cursor = 0 - per_item_lens = [] if LOG_MM_TIMING else None - for item in items: - n_tokens = _item_token_count(item) - if per_item_lens is not None: - per_item_lens.append(n_tokens) - next_cursor = cursor + n_tokens - if next_cursor > output_rows: - raise RuntimeError( - "VisionEmbedder encoder output is shorter than " - f"planned tokens for {modality}: needed " - f"{next_cursor}, got {output_rows}" - ) - item.encoded = output[cursor:next_cursor] - if spec.deepstack: - item.encoded, item.encoded_deepstack = ( - multimodal_model.separate_deepstack_embeds(item.encoded) - ) - cursor = next_cursor - if cursor != output_rows: - raise RuntimeError( - "VisionEmbedder encoder output token count does not match " - f"planned tokens for {modality}: planned {cursor}, " - f"got {output_rows}" - ) + per_item_lens = [_item_token_count(it) for it in items] + per_item_embs = torch.split(output, per_item_lens, dim=0) + + if spec.deepstack: + for item, emb in zip(items, per_item_embs): + main, deep = multimodal_model.separate_deepstack_embeds(emb) + item.encoded = main + item.encoded_deepstack = deep + else: + for item, emb in zip(items, per_item_embs): + item.encoded = emb if LOG_MM_TIMING: logger.info( "mm_timing encoder_ms modality=%s items=%d " @@ -607,47 +432,24 @@ def _assemble( encoders: dict[Modality, EncoderSpec], multimodal_model: nn.Module, ) -> tuple[torch.Tensor, dict[str, Any]]: - # Placeholder positions may hold content-derived IDs beyond vocab_size. - # Clamp them for the lookup; vision rows are overwritten below. - scatter_ranges = _coalesced_scatter_ranges(plan.scatter_ranges) - vocab_size = getattr(text_embedding, "num_embeddings") - input_embeds = text_embedding(input_ids.clamp(min=0, max=int(vocab_size) - 1)) - - has_deepstack = any(spec.deepstack for spec in encoders.values()) - visual_ranges = None - text_spans = None - if has_deepstack and input_ids.dim() == 1: - input_rows = int(input_ids.shape[0]) - visual_ranges = _merged_visual_ranges( - scatter_ranges, - input_rows, - ) - if visual_ranges: - text_spans = _text_token_spans(input_rows, visual_ranges) + # Placeholder positions hold large content-derived IDs that exceed + # vocab_size; the lookup we run here is overwritten for those rows + # by the scatter below, but the lookup still needs valid indices. + vocab_size = text_embedding.num_embeddings + safe_ids = input_ids.clamp(min=0, max=vocab_size - 1) + input_embeds = text_embedding(safe_ids) + kwargs: dict[str, Any] = {} deepstack_buffer: torch.Tensor | None = None - deepstack_buffer_is_direct = False - if has_deepstack: + if any(spec.deepstack for spec in encoders.values()): num_deepstack = len(multimodal_model.deepstack_visual_indexes) - deepstack_buffer = _dense_deepstack_view_for_full_visual_chunk( - scatter_ranges, - int(input_embeds.shape[0]) if input_embeds.dim() == 2 else 0, - expected_width=int(input_embeds.shape[-1]) * num_deepstack, - dtype=input_embeds.dtype, - device=input_embeds.device, + shape = input_embeds.shape[:-1] + (input_embeds.shape[-1] * num_deepstack,) + deepstack_buffer = torch.zeros( + shape, dtype=input_embeds.dtype, device=input_embeds.device ) - deepstack_buffer_is_direct = deepstack_buffer is not None - if deepstack_buffer is None: - deepstack_buffer = _allocate_deepstack_buffer( - input_embeds, - num_deepstack, - scatter_ranges, - visual_ranges, - text_spans, - ) kwargs["input_deepstack_embeds"] = deepstack_buffer - for r in scatter_ranges: + for r in plan.scatter_ranges: main = r.item.encoded if main is None: raise RuntimeError( @@ -655,28 +457,17 @@ def _assemble( "encoded tensor after _encode; this is a bug" ) src = main[r.item_src_start : r.item_src_end + 1] - input_embeds[r.flat_dst_start : r.flat_dst_end + 1] = _ensure_tensor_layout( - src, - dtype=input_embeds.dtype, - device=input_embeds.device, + input_embeds[r.flat_dst_start : r.flat_dst_end + 1] = src.to( + dtype=input_embeds.dtype, device=input_embeds.device ) - if deepstack_buffer_is_direct: - continue - if deepstack_buffer is not None and r.item.encoded_deepstack is not None: - deep_src = r.item.encoded_deepstack - deep_dst = deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1] - _copy_deepstack_rows( - deep_dst, - deep_src, - dtype=input_embeds.dtype, - device=input_embeds.device, - src_start=r.item_src_start, - src_end=r.item_src_end + 1, + deep_src = r.item.encoded_deepstack[ + r.item_src_start : r.item_src_end + 1 + ] + deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1] = deep_src.to( + dtype=input_embeds.dtype, device=input_embeds.device ) - elif deepstack_buffer is not None: - deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1].zero_() return input_embeds, kwargs diff --git a/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py index a5bf7de07..32449e0aa 100644 --- a/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py +++ b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py @@ -104,8 +104,6 @@ def prepare_metadata( batch: EncoderCudaGraphBatch, encoder_output_token_budget: int | None, metadata_sequence_budget: int, - *, - pad_cu_seqlens: bool = True, ) -> dict[str, Any]: ... def forward( @@ -133,7 +131,6 @@ class VisionEncoderBatch: tokens: torch.Tensor grid: torch.Tensor out_div: int - grid_rows: list[list[int]] | None = None @property def input_tensors(self) -> dict[str, torch.Tensor]: @@ -141,8 +138,6 @@ def input_tensors(self) -> dict[str, torch.Tensor]: @cached_property def _grid_rows(self) -> list[list[int]]: - if self.grid_rows is not None: - return self.grid_rows return self.grid.tolist() def num_items(self) -> int: @@ -167,55 +162,15 @@ def select(self, indices: list[int]) -> "VisionEncoderBatch": """Sub-batch at ``indices``, preserving order.""" cu = self.cu_input if indices: - if len(indices) == 1: - item_idx = indices[0] - tokens = self.tokens[cu[item_idx] : cu[item_idx + 1]] - grid = self.grid[item_idx : item_idx + 1] - elif _is_contiguous(indices): - start_idx = indices[0] - end_idx = indices[-1] + 1 - tokens = self.tokens[cu[start_idx] : cu[end_idx]] - grid = self.grid[start_idx:end_idx] - else: - total_rows = sum(cu[i + 1] - cu[i] for i in indices) - if total_rows <= 8192: - rows = _small_row_indices_from_cu(cu, indices, self.tokens.device) - else: - rows = torch.cat( - [ - torch.arange(cu[i], cu[i + 1], device=self.tokens.device) - for i in indices - ] - ) - tokens = self.tokens.index_select(0, rows) - grid_indices = torch.as_tensor( - indices, - dtype=torch.long, - device=self.grid.device, - ) - grid = self.grid.index_select(0, grid_indices) + rows = torch.cat( + [ + torch.arange(cu[i], cu[i + 1], device=self.tokens.device) + for i in indices + ] + ) else: - tokens = self.tokens[:0] - grid = self.grid[:0] - grid_rows = None - if self.grid_rows is not None: - grid_rows = [self.grid_rows[i] for i in indices] - return VisionEncoderBatch(tokens, grid, self.out_div, grid_rows) - - -def _is_contiguous(indices: list[int]) -> bool: - return all(curr == prev + 1 for prev, curr in zip(indices, indices[1:])) - - -def _small_row_indices_from_cu( - cu: list[int], - indices: list[int], - device: torch.device, -) -> torch.Tensor: - rows: list[int] = [] - for i in indices: - rows.extend(range(cu[i], cu[i + 1])) - return torch.tensor(rows, dtype=torch.long, device=device) + rows = torch.zeros(0, dtype=torch.long, device=self.tokens.device) + return VisionEncoderBatch(self.tokens[rows], self.grid[indices], self.out_div) @dataclass @@ -223,11 +178,7 @@ class VisionEncoderCudaGraphAdapter: """Adapter for Qwen/Kimi-style ``grid_thw`` vision encoders.""" tower: Any - pre_encode: Callable[ - [list[Any]], - tuple[torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, list[list[int]]], - ] + pre_encode: Callable[[list[Any]], tuple[torch.Tensor, torch.Tensor]] post_encode: Callable[[list[torch.Tensor], torch.Tensor], torch.Tensor] out_div: int merge: int @@ -258,9 +209,10 @@ def synthetic_grid(self, encoder_output_token_budget: int) -> list[list[int]]: b = units // a return [[1, a * self.merge, b * self.merge]] - def _checked_cu_seqlens_target( - self, cu: torch.Tensor, metadata_sequence_budget: int - ) -> int: + def pad_cu_seqlens( + self, metadata: dict[str, Any], metadata_sequence_budget: int + ) -> None: + cu = metadata["cu_seqlens"] target = metadata_sequence_budget + 1 if cu.shape[0] > target: raise RuntimeError( @@ -268,25 +220,13 @@ def _checked_cu_seqlens_target( f"metadata sequences, but the configured limit is " f"{metadata_sequence_budget}" ) - return target - - def pad_cu_seqlens( - self, metadata: dict[str, Any], metadata_sequence_budget: int - ) -> None: - cu = metadata["cu_seqlens"] - target = self._checked_cu_seqlens_target(cu, metadata_sequence_budget) pad = target - cu.shape[0] if pad > 0: metadata["cu_seqlens"] = torch.cat([cu, cu[-1:].expand(pad)]) def batch_from_items(self, items: list[Any]) -> VisionEncoderBatch: - pre_encoded = self.pre_encode(items) - if len(pre_encoded) == 3: - tokens, grid, grid_rows = pre_encoded - else: - tokens, grid = pre_encoded - grid_rows = None - return VisionEncoderBatch(tokens, grid, self.out_div, grid_rows) + tokens, grid = self.pre_encode(items) + return VisionEncoderBatch(tokens, grid, self.out_div) def capture_batch_for_budget( self, @@ -313,23 +253,15 @@ def prepare_metadata( batch: EncoderCudaGraphBatch, encoder_output_token_budget: int | None, metadata_sequence_budget: int, - *, - pad_cu_seqlens: bool = True, ) -> dict[str, Any]: if not isinstance(batch, VisionEncoderBatch): raise TypeError( f"{self.modality_name} encoder cudagraph expected " f"VisionEncoderBatch, got {type(batch).__name__}" ) - metadata_grid = batch._grid_rows if batch.grid_rows is not None else batch.grid - metadata = dict(self.tower.prepare_metadata(metadata_grid)) + metadata = dict(self.tower.prepare_metadata(batch.grid)) if encoder_output_token_budget is not None: - if pad_cu_seqlens: - self.pad_cu_seqlens(metadata, metadata_sequence_budget) - else: - self._checked_cu_seqlens_target( - metadata["cu_seqlens"], metadata_sequence_budget - ) + self.pad_cu_seqlens(metadata, metadata_sequence_budget) # Non-tensor scalar gets baked at capture. Use the per-budget worst # case so replay never exceeds the captured attention max seqlen. metadata["max_seqlen"] = encoder_output_token_budget * self.out_div @@ -542,80 +474,6 @@ def _scatter_output_slices( dest[idx] = sliced.clone() if clone else sliced offset += n_tokens - @staticmethod - def _scatter_cloned_output_slices( - output: torch.Tensor, - indices: list[int], - per_item_encoder_output_tokens: list[int], - dest: dict[int, torch.Tensor], - ) -> None: - """Clone one compact sub-batch output, then scatter views by item. - - Captured graph outputs live in a reusable output buffer. Cloning each - item slice separately creates many small GPU copies under high - concurrency; one contiguous clone per replay keeps item tensors stable - without changing values. - """ - owned = EncoderCudaGraphWrapper._clone_compact_output( - output, - indices, - per_item_encoder_output_tokens, - ) - EncoderCudaGraphWrapper._scatter_output_slices( - owned, - indices, - per_item_encoder_output_tokens, - dest, - ) - - @staticmethod - def _clone_compact_output( - output: torch.Tensor, - indices: list[int], - per_item_encoder_output_tokens: list[int], - ) -> torch.Tensor: - actual_tokens = sum(per_item_encoder_output_tokens[idx] for idx in indices) - return output[:actual_tokens].clone() - - @staticmethod - def _sorted_indices_by_encoder_output_tokens( - per_item_encoder_output_tokens: list[int], - ) -> list[int]: - if len(per_item_encoder_output_tokens) <= 1: - return list(range(len(per_item_encoder_output_tokens))) - if all( - prev <= curr - for prev, curr in zip( - per_item_encoder_output_tokens, - per_item_encoder_output_tokens[1:], - ) - ): - return list(range(len(per_item_encoder_output_tokens))) - return sorted( - range(len(per_item_encoder_output_tokens)), - key=lambda i: per_item_encoder_output_tokens[i], - ) - - @staticmethod - def _copy_prefix_zero_tail(buf: torch.Tensor, src: torch.Tensor) -> None: - n = src.shape[0] - if n == buf.shape[0]: - buf.copy_(src) - return - buf[:n].copy_(src) - buf[n:].zero_() - - @staticmethod - def _copy_prefix_repeat_last_tail(buf: torch.Tensor, src: torch.Tensor) -> None: - n = src.shape[0] - if n == 0: - raise RuntimeError("cannot repeat the last row from an empty tensor") - if n == buf.shape[0]: - buf.copy_(src) - return - buf[:n].copy_(src) - buf[n:].copy_(src[-1:].expand_as(buf[n:])) - def _run_budget_graph( self, batch: EncoderCudaGraphBatch, @@ -652,14 +510,12 @@ def _run_budget_graph( f"shape changed after dim0: capture={tuple(buf.shape)}, " f"replay={tuple(src.shape)}" ) - self._copy_prefix_zero_tail(buf, src) + buf.zero_() + buf[:n].copy_(src) metadata = dict( self.adapter.prepare_metadata( - batch, - encoder_output_token_budget, - metadata_sequence_budget, - pad_cu_seqlens=False, + batch, encoder_output_token_budget, metadata_sequence_budget ) ) replay_buffers = { @@ -690,10 +546,8 @@ def _run_budget_graph( f"has {new.shape[0]} rows, but the captured buffer only " f"has {buf.shape[0]} rows" ) - if key == "cu_seqlens": - self._copy_prefix_repeat_last_tail(buf, new) - else: - self._copy_prefix_zero_tail(buf, new) + buf.zero_() + buf[: new.shape[0]].copy_(new) graph_meta.graph.replay() return graph_meta.output_buffer @@ -716,10 +570,9 @@ def _dispatch(self, batch: EncoderCudaGraphBatch) -> list[torch.Tensor]: per_item_encoder_output_tokens = batch.encoder_output_tokens per_item_metadata_sequences = batch.metadata_sequences - sorted_indices = self._sorted_indices_by_encoder_output_tokens( - per_item_encoder_output_tokens + sorted_indices = sorted( + range(num_items), key=lambda i: per_item_encoder_output_tokens[i] ) - preserve_batch_chunks = sorted_indices == list(range(num_items)) batches: list[tuple[list[int], int | None]] = [] current_batch: list[int] = [] @@ -764,42 +617,27 @@ def _dispatch(self, batch: EncoderCudaGraphBatch) -> list[torch.Tensor]: ) # Packing reorders; restore original order before return. - ordered_outputs: list[torch.Tensor] = [] outputs_by_orig_idx: dict[int, torch.Tensor] = {} for batch_orig_indices, encoder_output_token_budget in batches: sub_batch = batch.select(batch_orig_indices) if encoder_output_token_budget is None: with torch.inference_mode(): raw = self._run_eager(sub_batch) - if preserve_batch_chunks: - ordered_outputs.append(raw) - else: - self._scatter_output_slices( - raw, - batch_orig_indices, - per_item_encoder_output_tokens, - outputs_by_orig_idx, - ) + self._scatter_output_slices( + raw, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + ) else: output = self._run_budget_graph(sub_batch, encoder_output_token_budget) - # output is the shared, reused output_buffer; clone once per - # replay and scatter views out of the owned compact tensor. - if preserve_batch_chunks: - ordered_outputs.append( - self._clone_compact_output( - output, - batch_orig_indices, - per_item_encoder_output_tokens, - ) - ) - else: - self._scatter_cloned_output_slices( - output, - batch_orig_indices, - per_item_encoder_output_tokens, - outputs_by_orig_idx, - ) + # clone: output is the shared, reused output_buffer. + self._scatter_output_slices( + output, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + clone=True, + ) - if preserve_batch_chunks: - return ordered_outputs return [outputs_by_orig_idx[i] for i in range(num_items)] diff --git a/test/runtime/test_encoder_cudagraph_grid_rows.py b/test/runtime/test_encoder_cudagraph_grid_rows.py deleted file mode 100644 index c7f993b27..000000000 --- a/test/runtime/test_encoder_cudagraph_grid_rows.py +++ /dev/null @@ -1,223 +0,0 @@ -import torch - -from tokenspeed.runtime.models.qwen3_vision import ( - Qwen3VLMoeVisionModel, -) -from tokenspeed.runtime.multimodal.encoder_cudagraph import ( - EncoderCudaGraphWrapper, - VisionEncoderBatch, - VisionEncoderCudaGraphAdapter, -) - - -def test_vision_encoder_batch_selects_single_and_contiguous_items(): - tokens = torch.arange(5, dtype=torch.float32).reshape(5, 1) - grid = torch.tensor([[1, 1, 1], [1, 2, 1], [1, 1, 2]], dtype=torch.int32) - grid_rows = [[1, 1, 1], [1, 2, 1], [1, 1, 2]] - batch = VisionEncoderBatch(tokens, grid, out_div=1, grid_rows=grid_rows) - - assert batch.encoder_output_tokens == [1, 2, 2] - contiguous = batch.select([0, 1]) - assert torch.equal(contiguous.tokens, tokens[:3]) - assert torch.equal(contiguous.grid, grid[:2]) - assert contiguous.grid_rows == [[1, 1, 1], [1, 2, 1]] - - single = batch.select([2]) - assert torch.equal(single.tokens, tokens[3:5]) - assert torch.equal(single.grid, grid[2:3]) - assert single.grid_rows == [[1, 1, 2]] - - -def test_vision_encoder_batch_selects_noncontiguous_items(): - tokens = torch.arange(10, dtype=torch.float32).reshape(10, 1) - grid = torch.tensor( - [[1, 1, 1], [1, 2, 1], [1, 3, 1], [1, 4, 1]], - dtype=torch.int32, - ) - grid_rows = grid.tolist() - batch = VisionEncoderBatch(tokens, grid, out_div=1, grid_rows=grid_rows) - - selected = batch.select([2, 0]) - - assert torch.equal(selected.tokens, torch.cat([tokens[3:6], tokens[:1]], dim=0)) - assert torch.equal(selected.grid, grid[[2, 0]]) - assert selected.grid_rows == [grid_rows[2], grid_rows[0]] - - -def test_qwen3_vision_rot_pos_emb_repeats_uniform_rows_exactly(): - class Rotary: - cos_sin_cache = torch.empty((1, 2), dtype=torch.float32) - - class DummyTower: - device = torch.device("cpu") - rotary_pos_emb = Rotary() - - def _rot_pos_emb_one(self, *_args): - return ( - torch.tensor([[1.0, 2.0], [3.0, 4.0]]), - torch.tensor([[5.0, 6.0], [7.0, 8.0]]), - ) - - cos, sin = Qwen3VLMoeVisionModel.rot_pos_emb( - DummyTower(), - [[1, 2, 2], [1, 2, 2]], - ) - - assert torch.equal( - cos, - torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]]), - ) - assert torch.equal( - sin, - torch.tensor([[5.0, 6.0], [7.0, 8.0], [5.0, 6.0], [7.0, 8.0]]), - ) - - -def test_qwen3_vision_pos_embed_repeats_uniform_rows_exactly(): - class PosEmbed: - embedding_dim = 2 - - class DummyTower: - num_grid_per_side = 8 - spatial_merge_size = 2 - pos_embed = PosEmbed() - device = torch.device("cpu") - dtype = torch.float32 - - def _pos_embed_interpolate_one(self, *_args): - return torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - - tower = DummyTower() - - out = Qwen3VLMoeVisionModel.fast_pos_embed_interpolate_from_list( - tower, - [[1, 2, 2], [1, 2, 2]], - ) - - assert torch.equal( - out, - torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]]), - ) - - -def test_vision_encoder_adapter_passes_precomputed_grid_rows_to_metadata(): - tokens = torch.zeros((2, 1), dtype=torch.float32) - grid = torch.tensor([[1, 2, 1]], dtype=torch.int32) - grid_rows = [[1, 2, 1]] - - class DummyTower: - def __init__(self): - self.seen_grid = None - - def prepare_metadata(self, grid_arg): - self.seen_grid = grid_arg - return {} - - tower = DummyTower() - - def pre_encode(_items): - return tokens, grid, grid_rows - - adapter = VisionEncoderCudaGraphAdapter( - tower=tower, - pre_encode=pre_encode, - post_encode=lambda encoder_outs, _grid: torch.cat(encoder_outs, dim=0), - out_div=1, - merge=1, - input_feature_shape=(1,), - ) - - batch = adapter.batch_from_items([object()]) - adapter.prepare_metadata( - batch, encoder_output_token_budget=None, metadata_sequence_budget=1 - ) - - assert tower.seen_grid is grid_rows - - -def test_vision_encoder_adapter_can_skip_replay_cu_seqlens_padding(): - tokens = torch.zeros((2, 1), dtype=torch.float32) - grid = torch.tensor([[1, 2, 1]], dtype=torch.int32) - - class DummyTower: - def prepare_metadata(self, _grid_arg): - return {"cu_seqlens": torch.tensor([0, 2], dtype=torch.int32)} - - adapter = VisionEncoderCudaGraphAdapter( - tower=DummyTower(), - pre_encode=lambda _items: (tokens, grid), - post_encode=lambda encoder_outs, _grid: torch.cat(encoder_outs, dim=0), - out_div=1, - merge=1, - input_feature_shape=(1,), - ) - batch = VisionEncoderBatch(tokens, grid, out_div=1) - - capture_metadata = adapter.prepare_metadata( - batch, encoder_output_token_budget=4, metadata_sequence_budget=4 - ) - replay_metadata = adapter.prepare_metadata( - batch, - encoder_output_token_budget=4, - metadata_sequence_budget=4, - pad_cu_seqlens=False, - ) - - assert torch.equal(capture_metadata["cu_seqlens"], torch.tensor([0, 2, 2, 2, 2])) - assert torch.equal(replay_metadata["cu_seqlens"], torch.tensor([0, 2])) - - -def test_encoder_cudagraph_copy_prefix_zero_tail_only_clears_unused_rows(): - buf = torch.full((4, 2), 9.0) - src = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - - EncoderCudaGraphWrapper._copy_prefix_zero_tail(buf, src) - - assert torch.equal( - buf, - torch.tensor( - [ - [1.0, 2.0], - [3.0, 4.0], - [0.0, 0.0], - [0.0, 0.0], - ] - ), - ) - exact = torch.full_like(src, 9.0) - EncoderCudaGraphWrapper._copy_prefix_zero_tail(exact, src) - assert torch.equal(exact, src) - - -def test_encoder_cudagraph_copy_prefix_repeat_last_tail_matches_padding(): - buf = torch.full((5,), -1, dtype=torch.int32) - src = torch.tensor([0, 2], dtype=torch.int32) - - EncoderCudaGraphWrapper._copy_prefix_repeat_last_tail(buf, src) - - assert torch.equal(buf, torch.tensor([0, 2, 2, 2, 2], dtype=torch.int32)) - exact = torch.full_like(src, -1) - EncoderCudaGraphWrapper._copy_prefix_repeat_last_tail(exact, src) - assert torch.equal(exact, src) - - -def test_encoder_cudagraph_scatter_clones_once_then_scatters_views(): - output = torch.arange(12, dtype=torch.float32).reshape(6, 2) - dest = {} - - EncoderCudaGraphWrapper._scatter_cloned_output_slices( - output, - [0, 1, 2], - [1, 2, 3], - dest, - ) - - assert torch.equal(dest[0], output[0:1]) - assert torch.equal(dest[1], output[1:3]) - assert torch.equal(dest[2], output[3:6]) - assert dest[0].untyped_storage().data_ptr() != output.untyped_storage().data_ptr() - assert ( - dest[0].untyped_storage().data_ptr() - == dest[1].untyped_storage().data_ptr() - == dest[2].untyped_storage().data_ptr() - ) diff --git a/test/runtime/test_qwen35_vision_assembly.py b/test/runtime/test_qwen35_vision_assembly.py new file mode 100644 index 000000000..2df7f0eba --- /dev/null +++ b/test/runtime/test_qwen35_vision_assembly.py @@ -0,0 +1,69 @@ +from types import SimpleNamespace + +import pytest +import torch + +from tokenspeed.runtime.models.qwen3_5 import Qwen3_5ForConditionalGeneration +from tokenspeed.runtime.multimodal.inputs import Modality + + +class _Visual: + dtype = torch.float32 + + def prepare_patch_embed(self, pixel_values, grid): + self.pixel_values = pixel_values + self.grid = grid + return pixel_values + + +@pytest.mark.parametrize( + ("modality", "grid_field"), + [ + (Modality.IMAGE, "image_grid_thw"), + (Modality.VIDEO, "video_grid_thw"), + ], +) +def test_single_item_pre_encode_reuses_feature_and_grid(modality, grid_field): + visual = _Visual() + model = SimpleNamespace(visual=visual) + feature = torch.arange(12, dtype=torch.float32).reshape(3, 4) + grid = torch.tensor([[1, 2, 2]], dtype=torch.int64) + item = SimpleNamespace(feature=feature, modality=modality, **{grid_field: grid}) + + tokens, result_grid = Qwen3_5ForConditionalGeneration.pre_encode(model, [item]) + + assert tokens.data_ptr() == feature.data_ptr() + assert visual.pixel_values.data_ptr() == feature.data_ptr() + assert result_grid.data_ptr() == grid.data_ptr() + assert visual.grid.data_ptr() == grid.data_ptr() + + +def test_multi_item_pre_encode_preserves_concatenation(): + visual = _Visual() + model = SimpleNamespace(visual=visual) + first = SimpleNamespace( + feature=torch.ones((2, 4), dtype=torch.float32), + modality=Modality.IMAGE, + image_grid_thw=torch.tensor([[1, 2, 2]], dtype=torch.int64), + ) + second = SimpleNamespace( + feature=torch.full((3, 4), 2.0, dtype=torch.float32), + modality=Modality.IMAGE, + image_grid_thw=torch.tensor([[1, 3, 2]], dtype=torch.int64), + ) + + tokens, grid = Qwen3_5ForConditionalGeneration.pre_encode(model, [first, second]) + + assert torch.equal(tokens, torch.cat([first.feature, second.feature], dim=0)) + assert torch.equal( + grid, + torch.cat([first.image_grid_thw, second.image_grid_thw], dim=0), + ) + + +def test_single_encoder_output_is_returned_without_concatenation(): + output = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + result = Qwen3_5ForConditionalGeneration.post_encode(None, [output], None) + + assert result is output diff --git a/test/runtime/test_qwen_vision_runtime_optimizations.py b/test/runtime/test_qwen_vision_runtime_optimizations.py deleted file mode 100644 index 4aee3e90d..000000000 --- a/test/runtime/test_qwen_vision_runtime_optimizations.py +++ /dev/null @@ -1,159 +0,0 @@ -import pytest -import torch - -from tokenspeed.runtime.multimodal.embedder import ( - EncodePlan, - EncoderSpec, - ScatterRange, - VisionEmbedder, - _coalesced_scatter_ranges, - _merged_visual_ranges, - _text_token_spans, -) -from tokenspeed.runtime.multimodal.inputs import ( - Modality, - MultimodalDataItem, - MultimodalForwardContext, - MultimodalInputs, -) - - -def test_visual_ranges_are_coalesced_for_embedding_assembly(): - item = MultimodalDataItem(modality=Modality.VIDEO) - other = MultimodalDataItem(modality=Modality.VIDEO) - scatter_ranges = [ - ScatterRange(1, 2, item, 0, 1), - ScatterRange(3, 5, item, 2, 4), - ScatterRange(7, 8, other, 0, 1), - ] - - coalesced = _coalesced_scatter_ranges(scatter_ranges) - visual_ranges = _merged_visual_ranges(coalesced, total_rows=10) - - assert coalesced == [ - ScatterRange(1, 5, item, 0, 4), - ScatterRange(7, 8, other, 0, 1), - ] - assert visual_ranges == [(1, 6), (7, 9)] - assert _text_token_spans(10, visual_ranges) == [(0, 1), (6, 7), (9, 10)] - - -def test_embedding_assembly_clamps_placeholder_ids_before_visual_overwrite(): - embedding = torch.nn.Embedding(10, 3) - with torch.no_grad(): - embedding.weight.copy_(torch.arange(30, dtype=torch.float32).reshape(10, 3)) - item = MultimodalDataItem( - modality=Modality.IMAGE, - encoded=torch.tensor( - [[100.0, 101.0, 102.0], [110.0, 111.0, 112.0]], - dtype=torch.float32, - ), - ) - plan = EncodePlan( - scatter_ranges=[ScatterRange(2, 3, item, 0, 1)], - ) - input_ids = torch.tensor([1, 2, 999, 999, 3], dtype=torch.long) - - input_embeds, kwargs = VisionEmbedder()._assemble( - input_ids=input_ids, - text_embedding=embedding, - plan=plan, - encoders={Modality.IMAGE: EncoderSpec(lambda _: torch.empty(0))}, - multimodal_model=None, - ) - - expected = embedding(input_ids.clamp(max=embedding.num_embeddings - 1)) - expected[2:4] = item.encoded - assert kwargs == {} - assert torch.equal(input_embeds, expected) - - -@pytest.mark.parametrize("modality", [Modality.IMAGE, Modality.VIDEO]) -def test_deepstack_combined_contract_for_chunked_prefill(modality): - embedding = torch.nn.Embedding(10, 3, dtype=torch.bfloat16) - with torch.no_grad(): - embedding.weight.copy_(torch.arange(30, dtype=torch.bfloat16).reshape(10, 3)) - - main = torch.arange(12, dtype=torch.bfloat16).reshape(4, 3) - deepstack = torch.arange(24, dtype=torch.bfloat16).reshape(4, 6) - combined = torch.cat([main, deepstack], dim=1) - input_ids = torch.tensor([999, 1, 2, 999, 999], dtype=torch.long) - - class Model: - deepstack_visual_indexes = [1, 2] - - @staticmethod - def separate_deepstack_embeds(output): - return output[:, :3], output[:, 3:] - - def run(): - item = MultimodalDataItem( - modality=modality, - offsets=[(2, 3), (6, 7)], - feature=torch.empty(0), - ) - context = MultimodalForwardContext( - mm_inputs=[MultimodalInputs(mm_items=[item])], - extend_prefix_lens=[3], - extend_seq_lens=[5], - ) - return VisionEmbedder().apply( - input_ids=input_ids, - text_embedding=embedding, - ctx=context, - encoders={modality: EncoderSpec(lambda _items: combined, deepstack=True)}, - multimodal_model=Model(), - ) - - input_embeds, kwargs = run() - - expected = embedding(input_ids.clamp(max=embedding.num_embeddings - 1)) - expected[0] = main[1] - expected[3:] = main[2:] - assert torch.equal(input_embeds, expected) - assert torch.equal(kwargs["input_deepstack_embeds"][0], deepstack[1]) - assert torch.equal(kwargs["input_deepstack_embeds"][3:], deepstack[2:]) - - -def test_dense_deepstack_reuses_full_visual_tensor(): - item = MultimodalDataItem( - modality=Modality.VIDEO, - encoded=torch.arange(6, dtype=torch.float32).reshape(2, 3), - encoded_deepstack=torch.arange(6, 12, dtype=torch.float32).reshape(2, 3), - ) - plan = EncodePlan(scatter_ranges=[ScatterRange(0, 1, item, 0, 1)]) - - class Model: - deepstack_visual_indexes = [8] - - _, kwargs = VisionEmbedder()._assemble( - input_ids=torch.tensor([999, 999], dtype=torch.long), - text_embedding=torch.nn.Embedding(10, 3), - plan=plan, - encoders={ - Modality.VIDEO: EncoderSpec(lambda _: torch.empty(0), deepstack=True) - }, - multimodal_model=Model(), - ) - - deepstack = kwargs["input_deepstack_embeds"] - assert torch.equal(deepstack, item.encoded_deepstack) - assert ( - deepstack.untyped_storage().data_ptr() - == item.encoded_deepstack.untyped_storage().data_ptr() - ) - - -def test_encoder_output_must_match_planned_tokens(): - item = MultimodalDataItem(modality=Modality.IMAGE, offsets=[(0, 0)]) - plan = EncodePlan() - plan.misses_by_modality[Modality.IMAGE] = [item] - output = torch.arange(6, dtype=torch.float32).reshape(2, 3) - - with pytest.raises(RuntimeError, match="does not match planned tokens"): - VisionEmbedder()._encode( - plan=plan, - encoders={Modality.IMAGE: EncoderSpec(lambda _: output)}, - multimodal_model=None, - device=torch.device("cpu"), - )