Method forward

tensorrt_llm/layers/ssm.py:172–228

Parameters: hidden_states: [B, L, D] or [T, D] conv_state: [B, W, D] or [1] of type int64 for paged state ssm_state: [B, N, D] or [1] of type int64 for paged state host_request_types: [B] last_token_ids: [B] host_contex

(self,
                hidden_states: Tensor,
                conv_state: Tensor,
                ssm_state: Tensor,
                host_request_types: Tensor,
                last_token_ids: Tensor,
                host_context_lengths: Optional[Tensor] = None,
                slot_mapping: Optional[Tensor] = None,
                conv_indices: Optional[Tensor] = None)

                               gather_output=False)

    def forward(self,
                hidden_states: Tensor,
                conv_state: Tensor,
                ssm_state: Tensor,
                host_request_types: Tensor,
                last_token_ids: Tensor,
                host_context_lengths: Optional[Tensor] = None,
                slot_mapping: Optional[Tensor] = None,
                conv_indices: Optional[Tensor] = None):
        '''
        Parameters:
            hidden_states: [B, L, D] or [T, D]
            conv_state: [B, W, D] or [1] of type int64 for paged state
            ssm_state: [B, N, D] or [1] of type int64 for paged state
            host_request_types: [B]
            last_token_ids: [B]
            host_context_lengths: [B]
            slot_mapping: [B]
            conv_indices: [B]
        '''
        # in_proj
        x = self.in_proj_x(hidden_states)
        z = self.in_proj_z(hidden_states)

        x_conv, conv_state = self.conv1d(x, conv_state, host_request_types,
                                         last_token_ids, host_context_lengths,
                                         slot_mapping, conv_indices)

        # Get dt, B and C
        x_dbl = self.x_proj(x_conv)
        if default_net().plugin_config.gemm_plugin:
            dt = self.dt_proj(x_dbl)
        else:
            dt, _ = split(x_dbl, [self.dt_rank, self.d_state * 2], dim=-1)
            dt = self.dt_proj(dt)

        # selective scan
        y, ssm_state = selective_scan(x_conv,
                                      ssm_state,
                                      dt,
                                      self.dt_bias.value,
                                      self.A.value,
                                      x_dbl,
                                      self.D.value,
                                      host_request_types,
                                      last_token_ids,
                                      self.d_inner,
                                      self.d_state,
                                      self.dt_rank,
                                      delta_softplus=True,
                                      dtype=self.dtype,
                                      z=z,
                                      host_context_lengths=host_context_lengths,
                                      slot_mapping=slot_mapping)
        # out_proj
        out = self.out_proj(y)
        return out, conv_state, ssm_state