Use updated Triton stateful backend behavior for streaming states

[alecgunny]

Triton's documentation now contains a description of a state management feature which should help to improve the efficiency of the stateful models in hermes.quiver by removing the need for updating the snapshot weight in the model itself. This in turn makes it unnecessary to implement the stateful models in TensorFlow. A simple implementation of this for the snapshotter might look like:

import torch
from typing import Sequence, Tuple

class Snapshotter(torch.nn.Module):
    def __init__(self, snapshot_size: int, channels: Sequence) -> None:
        super().__init__()
        self.snapshot_size = snapshot_size
        self.channels = channels

    def forward(self, update: torch.Tensor, snapshot: torch.Tensor) -> Tuple[torch.tensor, ...]:
        snapshot = torch.cat([update, snapshot[len(update):]], axis=2)
        snapshots = torch.split(snapshot, list(self.channels))
        return tuple(snapshots) + (snapshot,)

And the model config would include a section that looks something like

sequence_batching {
  state [
    {
      input_name: "old_snapshot"
      output_name: "new_snapshot"
      data_type: TYPE_INT32
      dims: [ -1 ]
      initial_state: {
       data_type: TYPE_INT32
       dims: [ 1 ]
       zero_data: true
       name: "initial state"
      }
    }
  ]
}

With the actual state naming mechanism needing to be worked out.

This will require adding functionality to the Exporter classes to be able to mark certain inputs as states.