lightseekorg--tokenspeed/python/tokenspeed/runtime/models/qwen3_vision.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Qwen3 visual tower base reused by the Qwen3.5 target models."""

from __future__ import annotations

from collections.abc import Callable
from functools import lru_cache, partial

import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
from transformers.activations import ACT2FN

from tokenspeed.runtime.configs.qwen3_vision_config import Qwen3VLVisionConfig
from tokenspeed.runtime.distributed.mapping import Mapping
from tokenspeed.runtime.layers.attention.mm_encoder_attention import (
    VIT_CUDNN_BATCH_BUCKETS,
    VIT_CUDNN_SEQLEN_BUCKETS,
    VIT_CUDNN_WORKSPACE_BYTES,
    VisionAttention,
    round_up_to_bucket,
)
from tokenspeed.runtime.layers.conv import Conv3dLayer
from tokenspeed.runtime.layers.linear import ColumnParallelLinear, RowParallelLinear
from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.rotary_embedding import get_rope
from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding
from tokenspeed.runtime.utils import add_prefix


@lru_cache(maxsize=1024)
def _rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
    if isinstance(h, torch.Tensor):
        h = int(h.item())
    if isinstance(w, torch.Tensor):
        w = int(w.item())
    if isinstance(spatial_merge_size, torch.Tensor):
        spatial_merge_size = int(spatial_merge_size.item())
    hpos_ids = np.broadcast_to(np.arange(h).reshape(h, 1), (h, w))
    h_div = h // spatial_merge_size
    w_div = w // spatial_merge_size
    hpos_ids = hpos_ids.reshape(h_div, spatial_merge_size, w_div, spatial_merge_size)
    hpos_ids = hpos_ids.transpose(0, 2, 1, 3).flatten()

    wpos_ids = np.broadcast_to(np.arange(w).reshape(1, w), (h, w))
    wpos_ids = wpos_ids.reshape(h_div, spatial_merge_size, w_div, spatial_merge_size)
    wpos_ids = wpos_ids.transpose(0, 2, 1, 3).flatten()
    return torch.from_numpy(np.stack([hpos_ids, wpos_ids], axis=-1))


class Qwen3VLVisionMLP(nn.Module):

    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        mapping: Mapping,
        bias: bool = True,
        hidden_act="silu",
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        vision = mapping.vision
        self.linear_fc1 = ColumnParallelLinear(
            in_features,
            hidden_features,
            bias=bias,
            quant_config=quant_config,
            prefix=add_prefix("linear_fc1", prefix),
            tp_size=vision.tp_size,
            tp_rank=vision.tp_rank,
            tp_group=vision.tp_group,
        )
        self.linear_fc2 = RowParallelLinear(
            hidden_features,
            in_features,
            bias=bias,
            quant_config=quant_config,
            prefix=add_prefix("linear_fc2", prefix),
            tp_size=vision.tp_size,
            tp_rank=vision.tp_rank,
            tp_group=vision.tp_group,
            reduce_results=True,
        )
        self.act = ACT2FN[hidden_act]

    def forward(self, x: torch.Tensor):
        x_fc1, _ = self.linear_fc1(x)
        x_act = self.act(x_fc1)
        mlp_output, _ = self.linear_fc2(x_act)
        return mlp_output


class Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = Conv3dLayer(
            self.in_channels,
            self.embed_dim,
            kernel_size=kernel_size,
            stride=kernel_size,
            bias=True,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1,
            self.in_channels,
            self.temporal_patch_size,
            self.patch_size,
            self.patch_size,
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(
            -1, self.embed_dim
        )
        return hidden_states


class Qwen3VLVisionBlock(nn.Module):

    def __init__(
        self,
        dim: int,
        num_heads: int,
        intermediate_dim: int,
        mapping: Mapping,
        head_size: int | None = None,
        hidden_act="silu",
        norm_layer: Callable[[int], nn.Module] | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        workspace_buffer: torch.Tensor | None = None,
        mm_attention_backend: str | None = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = partial(nn.LayerNorm, eps=1e-6)
        self.norm1 = norm_layer(dim)
        self.norm2 = norm_layer(dim)

        self.attn = VisionAttention(
            embed_dim=dim,
            num_heads=num_heads,
            head_size=head_size,
            proj_bias=True,
            quant_config=quant_config,
            prefix=add_prefix("attn", prefix),
            workspace_buffer=workspace_buffer,
            mapping=mapping,
            mm_attention_backend=mm_attention_backend,
        )
        self.mlp = Qwen3VLVisionMLP(
            dim,
            intermediate_dim,
            hidden_act=hidden_act,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
            mapping=mapping,
        )

    def forward(
        self,
        x: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb_cos: torch.Tensor,
        rotary_pos_emb_sin: torch.Tensor,
        max_seqlen: int | None = None,
        sequence_lengths: torch.Tensor | None = None,
    ) -> torch.Tensor:
        hidden_states = self.norm1(x)
        hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
        attn = self.attn(
            hidden_states,
            cu_seqlens=cu_seqlens,
            rotary_pos_emb_cos=rotary_pos_emb_cos,
            rotary_pos_emb_sin=rotary_pos_emb_sin,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )
        attn = rearrange(attn, "b s ... -> s b ...")
        x += attn
        norm2 = self.norm2(x)
        mlp = self.mlp(norm2)
        x += mlp
        return x


class Qwen3VLMoeVisionPatchMerger(nn.Module):

    def __init__(
        self,
        dim: int,
        context_dim: int,
        padded_context_dim: int,
        mapping: Mapping,
        norm_layer: Callable[[int], nn.Module] | None = None,
        spatial_merge_size: int = 2,
        use_postshuffle_norm: bool = False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = context_dim * (spatial_merge_size**2)
        self.padded_context_dim = padded_context_dim * (spatial_merge_size**2)

        self.use_postshuffle_norm = use_postshuffle_norm

        if norm_layer is None:
            norm_layer = partial(nn.LayerNorm, eps=1e-6)
        self.norm = norm_layer(
            self.hidden_size if use_postshuffle_norm else context_dim
        )
        vision = mapping.vision
        self.linear_fc1 = ColumnParallelLinear(
            self.hidden_size,
            self.padded_context_dim,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("linear_fc1", prefix),
            tp_size=vision.tp_size,
            tp_rank=vision.tp_rank,
            tp_group=vision.tp_group,
        )
        self.act_fn = nn.GELU()
        self.linear_fc2 = RowParallelLinear(
            self.padded_context_dim,
            dim,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("linear_fc2", prefix),
            tp_size=vision.tp_size,
            tp_rank=vision.tp_rank,
            tp_group=vision.tp_group,
            reduce_results=True,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.use_postshuffle_norm:
            x = self.norm(x.view(-1, self.hidden_size))
        else:
            x = self.norm(x).view(-1, self.hidden_size)

        x_parallel, _ = self.linear_fc1(x)
        x_parallel = self.act_fn(x_parallel)
        out, _ = self.linear_fc2(x_parallel)
        return out


class Qwen3VLMoeVisionModel(nn.Module):

    def __init__(
        self,
        vision_config: Qwen3VLVisionConfig,
        mapping: Mapping,
        norm_eps: float = 1e-6,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        mm_attention_backend: str | None = None,
    ) -> None:
        super().__init__()
        vision = mapping.vision
        self.mm_attention_backend = mm_attention_backend
        self.hidden_size = vision_config.hidden_size
        self.num_heads = vision_config.num_heads
        self.num_position_embeddings = vision_config.num_position_embeddings
        self.num_grid_per_side = int(self.num_position_embeddings**0.5)
        self.spatial_merge_size = vision_config.spatial_merge_size
        # layer indices whose outputs feed the deepstack mergers
        self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes
        self.patch_embed = Qwen3VLVisionPatchEmbed(config=vision_config)
        self.pos_embed = VocabParallelEmbedding(
            self.num_position_embeddings,
            self.hidden_size,
            quant_config=quant_config,
            tp_rank=vision.tp_rank,
            tp_size=vision.tp_size,
            tp_group=vision.tp_group,
            prefix=add_prefix("pos_embed", prefix),
        )

        norm_layer = partial(nn.LayerNorm, eps=norm_eps)
        head_dim = self.hidden_size // self.num_heads
        self.rotary_pos_emb = get_rope(
            head_size=head_dim,
            rotary_dim=head_dim // 2,
            max_position=8192,
            base=10000.0,
            is_neox_style=True,
        )

        workspace_buffer = None
        if self.mm_attention_backend == "flashinfer_cudnn":
            if torch.cuda.is_available():
                ws_device = torch.device("cuda", torch.cuda.current_device())
            else:
                ws_device = self.device
            workspace_buffer = torch.empty(
                VIT_CUDNN_WORKSPACE_BYTES,
                dtype=torch.uint8,
                device=ws_device,
            )

        self.blocks = nn.ModuleList(
            [
                Qwen3VLVisionBlock(
                    dim=self.hidden_size,
                    num_heads=self.num_heads,
                    intermediate_dim=vision_config.intermediate_size,
                    head_size=head_dim,
                    hidden_act=vision_config.hidden_act,
                    norm_layer=norm_layer,
                    quant_config=quant_config,
                    prefix=add_prefix(f"blocks.{layer_idx}", prefix),
                    workspace_buffer=workspace_buffer,
                    mapping=mapping,
                    mm_attention_backend=mm_attention_backend,
                )
                for layer_idx in range(vision_config.depth)
            ]
        )
        self.merger = Qwen3VLMoeVisionPatchMerger(
            dim=vision_config.out_hidden_size,
            context_dim=self.hidden_size,
            padded_context_dim=self.num_heads * head_dim,
            norm_layer=norm_layer,
            spatial_merge_size=self.spatial_merge_size,
            quant_config=quant_config,
            prefix=add_prefix("merger", prefix),
            mapping=mapping,
        )

        self.deepstack_merger_list = nn.ModuleList(
            [
                Qwen3VLMoeVisionPatchMerger(
                    dim=vision_config.out_hidden_size,
                    context_dim=self.hidden_size,
                    padded_context_dim=self.num_heads * head_dim,
                    spatial_merge_size=self.spatial_merge_size,
                    use_postshuffle_norm=True,
                    norm_layer=norm_layer,
                    quant_config=quant_config,
                    prefix=add_prefix(f"deepstack_merger_list.{layer_idx}", prefix),
                    mapping=mapping,
                )
                for layer_idx in range(len(self.deepstack_visual_indexes))
            ]
        )

        self.tp_size = vision.tp_size

    @property
    def dtype(self) -> torch.dtype:
        return self.patch_embed.proj.weight.dtype

    @property
    def device(self) -> torch.device:
        return self.patch_embed.proj.weight.device

    def rot_pos_emb(
        self, grid_thw: list[list[int]]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        pos_ids = []
        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))

        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)

        cos, sin = self._get_rotary_cos_sin(max_grid_size)

        cos_combined = cos[pos_ids].flatten(1)
        sin_combined = 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)
        return cos_sin.chunk(2, dim=-1)

    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

        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
            )
            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)

        return torch.cat(outputs, dim=0)

    def compute_cudnn_batch_offsets_packed(
        self,
        token_cu_seqlens: np.ndarray,
        *,
        elem_per_token: int,
    ) -> np.ndarray:
        """
        Build packed *element* indptrs for cuDNN prefill.

        Input:
        token_cu_seqlens: (B+1,) token indptr
        elem_per_token: per-token element width on THIS TP rank
                        (usually hidden_size / attn_tp_size)

        Output:
        packed_offsets: (3 * (B_padded + 1),) int32
            [qk_indptr, v_indptr, o_indptr] concatenated,
            each indptr is (B_padded + 1,) in element units.
        """
        if token_cu_seqlens.ndim != 1 or token_cu_seqlens.size < 2:
            raise ValueError(
                "token_cu_seqlens must be a 1D array with at least 2 entries."
            )
        B = int(token_cu_seqlens.size - 1)
        B_padded = round_up_to_bucket(B, VIT_CUDNN_BATCH_BUCKETS)

        # token indptr -> pad to (B_padded+1,) by appending total_tokens for extra empty sequences
        token_indptr = token_cu_seqlens.astype(np.int64, copy=False)  # (B+1,)
        if B_padded != B:
            pad = np.full((B_padded - B,), token_indptr[-1], dtype=token_indptr.dtype)
            token_indptr = np.concatenate([token_indptr, pad], axis=0)  # (B_padded+1,)

        # convert token indptr -> element indptr
        elem_indptr = (token_indptr * int(elem_per_token)).astype(
            np.int32
        )  # (B_padded+1,)

        # q/k/v/o in this vision encoder path share the same indptr
        return np.concatenate([elem_indptr, elem_indptr, elem_indptr], axis=0)

    def compute_cudnn_sequence_lengths_padded(
        self,
        token_cu_seqlens: np.ndarray,
    ) -> np.ndarray:
        """
        token_cu_seqlens: (B+1,) token indptr
        return: (B_padded,) token lengths (padded with 0)
        """
        if token_cu_seqlens.ndim != 1 or token_cu_seqlens.size < 2:
            raise ValueError(
                "token_cu_seqlens must be a 1D array with at least 2 entries."
            )
        B = int(token_cu_seqlens.size - 1)

        seq_lens = (token_cu_seqlens[1:] - token_cu_seqlens[:-1]).astype(
            np.int32
        )  # (B,)

        B_padded = round_up_to_bucket(B, VIT_CUDNN_BATCH_BUCKETS)
        if B_padded != B:
            pad = np.zeros((B_padded - B,), dtype=np.int32)
            seq_lens = np.concatenate([seq_lens, pad], axis=0)  # (B_padded,)
        return seq_lens

    def prepare_patch_embed(
        self, x: torch.Tensor, grid_thw: torch.Tensor | list
    ) -> torch.Tensor:
        """Eager patch-embed (runs before the captured region): Conv patch embed
        + interpolated position embedding + the ``[s, 1, h]`` reshape the block
        loop expects.

        Kept eager (outside the capture-safe region) -- the interpolation does
        host/numpy work.
        """
        x = x.to(device=self.device, dtype=self.dtype)
        x = self.patch_embed(x)
        grid_thw_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()
        x = x + self.fast_pos_embed_interpolate_from_list(grid_thw_list)
        return x.unsqueeze(1)

    def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict:
        """Eager metadata pass: rotary embeddings, cu_seqlens, sequence lengths,
        and ``max_seqlen`` as a Python int.

        Everything here involves a host sync or a data-dependent shape, so it
        lives outside the capture-safe block loop. ``max_seqlen`` is
        materialized as a plain int (CPU/numpy, no GPU sync) so the captured
        block loop never hits the attention backend's ``.item()`` fallback.
        """
        if isinstance(grid_thw, list):
            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()

        rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)

        # ---- build token indptr (B+1,) ----
        token_cu_seqlens = np.concatenate(
            [
                np.zeros(1, dtype=np.int32),
                np.repeat(
                    grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]
                ).cumsum(axis=0, dtype=np.int32),
            ]
        )
        real_seq_lens = token_cu_seqlens[1:] - token_cu_seqlens[:-1]
        real_max_seqlen = int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0

        if self.mm_attention_backend == "flashinfer_cudnn":
            # (B_padded,) token lengths
            seq_lens_padded = self.compute_cudnn_sequence_lengths_padded(
                token_cu_seqlens
            )
            # element-per-token width on this vision TP rank
            elem_per_token = (
                self.hidden_size // self.tp_size
            )  # == heads_per_rank * head_dim
            # (3*(B_padded+1),) packed element indptrs
            offsets_packed = self.compute_cudnn_batch_offsets_packed(
                token_cu_seqlens,
                elem_per_token=elem_per_token,
            )
            sequence_lengths = (
                torch.from_numpy(seq_lens_padded)
                .to(device=self.device, dtype=torch.int32, non_blocking=True)
                .view(-1, 1, 1, 1)
            )  # match cuDNN test style
            cu_seqlens = torch.from_numpy(offsets_packed).to(
                device=self.device, dtype=torch.int32, non_blocking=True
            )
            max_seqlen = round_up_to_bucket(real_max_seqlen, VIT_CUDNN_SEQLEN_BUCKETS)
        else:
            sequence_lengths = None
            cu_seqlens = torch.from_numpy(token_cu_seqlens).to(
                device=self.device, dtype=torch.int32, non_blocking=True
            )
            max_seqlen = real_max_seqlen

        return {
            "cu_seqlens": cu_seqlens,
            "rotary_pos_emb_cos": rotary_pos_emb_cos,
            "rotary_pos_emb_sin": rotary_pos_emb_sin,
            "max_seqlen": max_seqlen,
            "sequence_lengths": sequence_lengths,
        }

    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. ``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"]
        rotary_pos_emb_sin = metadata["rotary_pos_emb_sin"]
        max_seqlen = metadata["max_seqlen"]
        sequence_lengths = metadata["sequence_lengths"]

        deepstack_feature_lists = []
        num_deepstack_captured = 0
        for layer_num, blk in enumerate(self.blocks):
            x = blk(
                x,
                cu_seqlens=cu_seqlens,
                rotary_pos_emb_cos=rotary_pos_emb_cos,
                rotary_pos_emb_sin=rotary_pos_emb_sin,
                max_seqlen=max_seqlen,
                sequence_lengths=sequence_lengths,
            )
            if layer_num in self.deepstack_visual_indexes:
                deepstack_feature = self.deepstack_merger_list[num_deepstack_captured](
                    x
                )
                deepstack_feature_lists.append(deepstack_feature)
                num_deepstack_captured += 1
        x = self.merger(x)
        # [seq_len, out_hidden_size * (1 + depth_of_deepstack)]
        return torch.cat([x] + deepstack_feature_lists, dim=1)