class_docstr: Inference model implementation for Mixtral models.
class_docstr_tok: ["Inference","model","implementation","for","Mixtral","models","."]
docstr: Builds the attention layer for the model. This sets the `self.attn` attribute.
docstr_tok: ["Builds","the","attention","layer","for","the","model",".","This","sets","the","`","self.attn","`","attribute","."]
code: def make_attn_layer(self) -> None:
      """
      Builds the attention layer for the model. This sets the `self.attn` attribute.
      """
      softmax_scale = 1.0 / (self.head_size**0.5)
  
      rotary_config = RotateHalfConfig(
          theta_base=self._config.rope_theta
      )
  
      attn_config = DSSelfAttentionConfig(
          max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
          n_heads_q=self.n_heads_q_local,
          n_heads_kv=self.n_heads_kv_local,
          head_size=self.head_size,
          max_sequences=self._engine_config.state_manager.max_ragged_sequence_count,
          scale_factor=softmax_scale,
          input_dtype=self.activation_dtype,
          output_dtype=self.activation_dtype,
          positional_embedding_type=self.positional_embedding_type,
          positional_embedding_config=rotary_config,
      )
  
      self.attn = heuristics.instantiate_attention(
          attn_config, self._engine_config
      )
  
imports: from typing import Iterable, Optional, Tuple
  import torch
  import deepspeed.comm
  from ...allocator import empty_from
  from ...config_v2 import RaggedInferenceEngineConfig
  from ...inference_utils import ActivationType, DtypeEnum
  from ...model_implementations import *
  from ...modules.configs import *
  from ...modules.interfaces import *
  from ...modules import heuristics
  from ...ragged import RaggedBatchWrapper
  from ..inference_model_base import DSModelImplementationConfig, MPType
  from .container import MixtralNonTransformerContainer, MixtralTransformerContainer
label_desc: Use image node_id 17 for calling the MixtralInferenceModel obj's underlying member method code with example usage: obj.make_attn_layer() without return types