Code

Add inference cache (cleaner version)#4

Code

Open

borzunov wants to merge 19 commits intoweight-sharinglearning-at-home/dalle-pytorch:weight-sharingfrom inference-cache-v2learning-at-home/dalle-pytorch:inference-cache-v2Copy head branch name to clipboard

+149-18

Commits

Commits on Dec 20, 2021

• Add FFN caching

  borzunov

  committed
  1cd8e20

• Cache to_qkv and to_out in sparse attn, add debug prints

  borzunov

  committed
  4d431ac

• Cache full Attention

  borzunov

  committed
  2603776

• Remove debug outputs

  borzunov

  committed
  112ea05

• Cache pre-logits MLP

  borzunov

  committed
  6ba4cb6

• Further optimize attention caching

  borzunov

  committed
  c333ea7

Commits on Dec 21, 2021

• Fix mask in attention

  borzunov

  committed
  1fd45ca

• Don't cache MLPs since we can just pass only last item

  borzunov

  committed
  2b77018

• Revert excess changes in attentions

  borzunov

  committed
  8e8dea8

• Rename FixCacheKey -> CachedAs

  borzunov

  committed
  df89951

• Save the current offset in cache

  borzunov

  committed
  059fe1b

• Use static masks to simulate axial attn

  borzunov

  committed
  b76b78e

• Add option to disable caching

  borzunov

  committed
  4c833a2

• Make the cached version work

  borzunov

  committed
  1ff47c6

• Speed up PreShiftToken

  borzunov

  committed
  94fda36

Commits on Jan 10, 2022

• Remove excess changes

  borzunov

  committed
  adfce34

• Add NonCached wrapper

  borzunov

  committed
  59cfc49

• Rename use_static_masks -> optimize_for_inference

  borzunov

  committed
  4f496a4

• Improve names and comments

  borzunov

  committed
  732226d