class_docstr: Built on the paper and library Megablocks as described in
  https://arxiv.org/abs/2211.15841. This implementation is
  strictly equivalent to standard MoE with full capacity (no
  dropped tokens). It's faster since it formulates MoE operations
  in terms of block-sparse operations to accomodate imbalanced
  assignments of tokens to experts, whereas standard MoE either
  (1) drop tokens at the cost of reduced performance or (2) set
  capacity factor to number of experts and thus waste computation
  and memory on padding.
class_docstr_tok: ["Built","on","the","paper","and","library","Megablocks","as","described","in","https",":","\/\/arxiv.org\/abs\/2211.15841",".","This","implementation","is","strictly","equivalent","to","standard","MoE","with","full","capacity","(","no","dropped","tokens",")",".","It","'s","faster","since","it","formulates","MoE","operations","in","terms","of","block-sparse","operations","to","accomodate","imbalanced","assignments","of","tokens","to","experts",",","whereas","standard","MoE","either","(","1",")","drop","tokens","at","the","cost","of","reduced","performance","or","(","2",")","set","capacity","factor","to","number","of","experts","and","thus","waste","computation","and","memory","on","padding","."]
code: def topology(self, x: torch.Tensor, padded_bins: torch.Tensor):
      padded_tokens, _ = x.size()
      assert padded_tokens % self.blocking == 0
      assert self.ffn_dim % self.blocking == 0
  
      # Offsets for the sparse matrix. All rows have the
      # same number of nonzero blocks dictated by the
      # dimensionality of a single expert.
      block_rows = padded_tokens // self.blocking
      blocks_per_row = self.ffn_dim // self.blocking
      if self.offsets is None or block_rows > self.offsets_block_rows:
          self.offsets = torch.arange(
              0,
              block_rows * blocks_per_row + 1,
              blocks_per_row,
              dtype=torch.int32,
              device=x.device,
          )
          self.offsets_block_rows = block_rows
          offsets = self.offsets
      else:
          offsets = self.offsets[: block_rows + 1]
  
      # Indices for the sparse matrix. The indices for
      # the intermediate matrix are dynamic depending
      # on the mapping of tokens to experts.
      column_indices = ops.topology(
          padded_bins, self.blocking, block_rows, blocks_per_row
      )
  
      # For now, use meta init to save the device memory.
      data = torch.empty(
          column_indices.numel(),
          self.blocking,
          self.blocking,
          dtype=x.dtype,
          device="meta",
      )
      shape = (padded_tokens, self.ffn_dim * self.num_experts)
      row_indices = stk.ops.row_indices(
          shape, data, offsets, column_indices
      )
      return stk.Matrix(
          shape,
          data,
          row_indices,
          column_indices,
          offsets,
          False,
          False,
          False,
      )
  
imports: import torch
  import torch.distributed
  import numpy
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from loguru import logger
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import FastLinear, FastRMSNorm, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear
label_desc: Use image node_id 2 for calling the BlockSparseMoE obj's underlying member method code with example usage: obj.topology(x, padded_bins) and returns: stk