Context, tokens, KV stores & storage, Solidigm presents at #AIFD8

Solidigm presented (video here) at AIFD8 this month and as part of their presentation they spent time disecting what happens to a prompt, how token growth happens, and where storage can help speed up prompt processing.

The token count explosion

It all starts at a simple prompt something as simple as “run a benchmark against a drive” maybe a 12 token prompt but when it actually gets processed can balloon into something that’s much larger. As an LLM processes the prompt it goes through a number of steps building context, calling tools, obtaining and interpreting results, persisting knowledge and finally, responding to the prompt.

Digging a level deeper, here’s what the token counts look like during prompt processing. First step is to understand the environment of the prompt, rules, safety requirements, methodology at it’s disposal, then there’s retrieval activity that gathers information needed to actually process and perform the prompt, then identifying tools and their APIs needed to process the prompt, and at some point when the LLM has all that it plans out the steps needed to actually perform the prompt, tool results are generated, interpreted and fed back to LLM processing to determine next step. All of which at some point, prompt precessing completes and the prompt reply is sent back to the issuer.

As one can see in the above, the prompt itself was minuscule in token counts in the vast scheme of activity needed to process the prompt. And this is just how one (albeit complex), ~12 token prompt can grow into a 42K token context.

Inferencing and Time To First Token

Inferencing consists of two phases:

  • PreFill phase – which is the processing that goes on to take the context token stream and convert it into a KV (Key:Value) store which the LLM can use for subsequent processing so it doesn’t have to go back to the token context. PreFill ends up with a fully populated KV store representing all the tokens in the current context, and generates the first token in the LLM response to the prompt
  • Decode – which is all subsequent processing needed to generate the rest of the prompt response, uses the that KV store to underpin it’s processing to generate any more tokens needed to answer the prompt.

Solidigm went on to describe how these activities impact the Time To First Token (TTFT), or how long it takes from the time the prompt is issued until the LLM responds with the first word (token) of the prompt response.

(Although in the Solidigm’s chart they show Decode in the TTFT path. I believe this incorrect as PreFill generates the first token. Nonetheless, there is a portion of PreFill that “decodes” the prompt response first token and I assume that’s what they are showing here. Of course I could be mistaken.)

Storage can impact both the time it takes to assemble context tokens and to perform PreFill.

While storage can matter a lot during context assemble (lots of potential IO activity reading files, RAGs and other documents), storage’s impact on PreFill is less widely known. That is until you understand how prompt processing can be held up for KV store recalculation (going back to context tokens and rebuilding some or all of the current KV store for the prompt).

Increasing context, leads to more tokens, leads to larger KV stores, all of which impacts TTFT

Although, it’s only conjecture on my part, but the biggest portion of the Tprefill above seems to be calculating and converting context/memory tokens into KV elements stored in the prompts KV store. KV stores are used during prompt downstream processing because they can be easily accessed and each KV item represents intepreted token information in an easily used (by LLM) fashion.

And what’s not evident in the above TTFT decomposition chart is that tool use, generates even more tokens, as tool result (tokens), all of which need to be processed into more KV store elements in order to determine what to do next.

What happens to large KV stores during prompt processing

If there is a single GPU running a single prompt it’s possible, depending on model and HBM size, that it will run out of GPU HBM memory and offload or move some portion of its KV (store) cache to CPU memory. But if that GPU is processing 100s to 1000s of prompts concurrently, even CPU memory may not be large enough to hold every KV cache segment that no longer fits in GPU HBM. And of course most enterprise AI servers hold anywhere from 4 to 10 GPUs, each running 100s to 1000s of prompts concurrently.

KV cache offload is where fast storage can significantly speed up prompt processing

There’s an obvious tradeoff here with respect to KV stores. One can always go back to the Prefill phase, reread all the tokens in current context and recompute the KV store or one can offload KV store segments to memory, local storage or network storage and later retrieve the already computed KV store from wherever it ended up.

The tradeoff is how long it takes to recompute vs do the data transfers to offload and retrieve the KV cache segments. Larger contexts, increase KV store size, which lead to more need to offload or jettison KV store segments when running out of GPU HBM space. Both KV caching to memory-storage vs jetisoning KV store segments and reconstituting them, add time to TTFT. The question is which is faster.

One can see how this would be made ever more of an issue as prompts token counts (& KV elements) skyrocket. Also when more prompts are running concurrently on the same GPU(s) in a single server.

Obviously local, large SSDs with very fast random read would be ideal for KV cache offload activity which has the KV cache segment written out once (and extended as prompt processing adds context) but read back multiple times. Which s is great application for Large capacity, fast read NVMe SSDs which, I must say, are Solidigm’s forte.

NVIDIA and others have started to add KV cache offloading to their inferencing stacks. As they do, large fast NVMe SSDs activity during AI prompt processing will become one of the critical factors in TTFT.

In the meantime, if anyone has any large, fast NVMe SSDs they don’t need anymore, please let me know. 🙂

Data and code versioning For MLops

Read an interesting article (Ex-Apple engineers raise … data storage startup) and research paper (Git is for data) about a of group of ML engineers from Apple forming a new “data storage” startup targeted at MLOps teams just like Apple. It turns out that MLops has some very unique data requirements that go way beyond just data storage.

The paper discusses some of the unusual data requirements for MLOps such as:

  • Infrequent updates – yes there are some MLOps datasets where updates are streamed in but the vast majority of MLOps datasets are updated on a slower cadence. The authors think monthly works for most MLOps teams
  • Small changes/lots of copies – The changes to MLOps data are relatively small compared to the overall dataset size and usually consist of data additions, record deletions, label updates, etc. But uncommon to most data, MLOps data are often subsetted or extracted into smaller datasets used for testing, experimentation and other “off-label” activities.
  • Variety of file types – depending on the application domain, MLOps file types range all over the place. But there’s often a lot of CSV files in combination with text, images, audio, and semi-structured data (DICOM, FASTQ, sensor streams, etc.). However within a single domain, MLOps file types are pretty much all the same.
  • Variety of file directory trees – this is very MLOps team and model dependent. Usually there are train/validate/test splits to every MLOps dataset but what’s underneath each of these can vary a lot and needs to be user customizable.
  • Data often requires pre-processing to be cleansed and made into something appropriate and more useable by ML models
  • Code and data must co-evolve together, over time – as data changes, the code that uses them change. Adding more data may not cause changes to code but models are constantly under scrutiny to improve performance, accuracy or remove biases. Bias elimination often requires data changes but code changes may also be needed.

It’s that last requirement, MLOps data and code must co-evolve and thus, need to be versioned together that’s most unusual. Data-code co-evolution is needed for reproducibility, rollback and QA but also for many other reasons as well.

In the paper they show a typical MLOps data pipeline.

Versioning can also provide data (and code) provenance, identifying the origin of data (and code). MLOps teams undergoing continuous integration need to know where data and code came from and who changed them. And as most MLOps teams collaborate in the development, they also need a way to identify data and code conflicts when multiple changes occur to the same artifact.

Source version control

Code has had this versioning problem forever and the solution became revision control systems (RCS) or source version control (SVC) systems. The most popular solutions for code RCS are Git (software) and GitHub (SaaS). Both provide repositories and source code version control (clone, checkout, diff, add/merge, commit, etc.) as well as a number of other features that enable teams of developers to collaborate on code development.

The only thing holding Git/GitHub back from being the answer to MLOps data and code version control is that they don’t handle large (>1MB) files very well.

The solution seems to be adding better data handling capabilities to Git or GitHub. And that’s what XetHub has created for Git.

XetHub’s “Git is Data” paper (see link above) explains what they do in much detail, as to how they provide a better data layer to Git, but it boils down to using Git for code versioning and as a metadata database for their deduplicating data store. They are using a Merkle trees to track the chunks of data in a deduped dataset.

How XetHub works

XetHub support (dedupe) variable chunking capabilities for their data store. This allows them to use relatively small files checked into Git to provide the metadata to point to the current (and all) previous versions of data files checked into the system.

Their mean chunk size is ~4KB. Data chunks are stored in their data store. But the manifest for dataset versions is effectively stored in the Git repository.

The paper shows how using a deduplicated data store can support data versioning.

XetHub uses a content addressable store (CAS) to store the file data chunk(s) as objects or BLOBs. The key to getting good IO performance out of such a system is to have small chunks but large objects.

They map data chunks to files using a CDMT (content defined merkle tree[s]). Each chunk of data resides in at least two different CDMTs, one associated with the file version and the other associated with the data storage elements.

XetHub’s variable chunking approach is done using a statistical approach and multiple checksums but they also offer one specialized file type chunking for CSV files. As it is, even with their general purpose variable chunking method, they can offer ~9X dedupe ratio for text data (embeddings).

They end up using Git commands for code and data but provide hooks (Git filters) to support data cloning, add/checkin, commits, etc.). So they can take advantage of all the capabilities of Git that have grown up over the years to support code collaborative development but use these for data as well as code.

In addition to normal Git services for code and data, XetHub also offers a read-only, NFSv3 file system interface to XetHub datases. Doing this eliminates having to reconstitute and copy TB of data from their code-data repo to user workstations. With NFSv3 front end access to XetHub data, users can easily incorporate data access for experimentation, testing and other uses.

Results from using XetHub

XetHub showed some benchmarks comparing their solution to GIT LFS, another Git based large data storage solution. For their benchmark, they used the CORD-19 (and ArXiv paper, and Kaggle CORD-I9 dataset) which is a corpus of all COVID-19 papers since COVID started. The corpus is updated daily, released periodically and they used the last 50 versions (up to June 2022) of the research corpus for their benchmark.

Each version of the CORD-19 corpus consists of JSON files (research reports, up to 700K each) and 2 large CSV files one with paper information and the other paper (word?) embeddings (a more useable version of the paper text/tables used for ML modeling).

For CORD-19, XetHub are able to store all the 2.45TB of research reports and CSV files in only 287GB of Git (metadata) and datastore data, or with a dedupe factor 8.7X. With XetHub’s specialized CSV chunking (Xet w/ CSV chunking above), the CORD-19 50 versions can be stored in 87GB or with a 28.8X dedupe ratio. And of that 87GB, only 82GB is data and the rest ~5GB is metadata (of which 1.7GB is the merle tree).

In the paper, they also showed the cost of branching this data by extracting and adding one version which consisted of a 75-25% (random) split of a version. This split was accomplished by changing only the two (paper metadata and paper word embeddings) CSV files. Adding this single split version to their code-data repository/datastore only took an additional 11GB of space An aligned split (only partitioning on a CSV record boundary, unclear but presumably with CSV chunking), only added 185KB.

XETHUB Potential Enhancements

XetHub envisions many enhancements to their solution, including adding other specific file type chunking strategies, adding a “time series” view to their NFS frontend to view code/data versions over time, finer granularity data provenance (at the record level rather than at the change level), and RW NFS access to data. Further, XetHub’s dedupe metadata (on the Git repo) only grows over time, supporting updates and deletes to dedupe metadata would help reduce data requirements.

Read the paper to find out more.

Picture/Graphic credit(s):