PagedAttention tracking issue

[EricLBuehler]

Tracking issue for PagedAttention

General overview

• Use updated, refactored PagedAttention impl from vLLM
• Implement attention biases
• Implement _memory_efficient_attention or equiv.
  • Implement xtransformers memory efficient transformers: https://arxiv.org/abs/2112.05682
• Implement and use block cache (Scheduler, CacheEngine) in inference.
• Finalize FFI linking to compiled CUDA kernels
• Test PagedAttention in llama , mistral models

Sampler

• Implement according to this

ModulePipeline refactoring

Need to redesign ModulePipeline design to integrate with the LLMEngine design. This means that a request output will be returned, and logprobs must be calculated.

• Calculate logprobs
• Generate best_of responses instead of the current n to check from (sorted by the logprobs)
• Use Sampler
• Do not return a ChatChoices, return something like a RawResponse that contains the logprobs, results
• Streaming poses a problem. This will need to be handled by the LLMEngine.
• Overall, convert ModulePipeline to a pure-sequence-generating platform

LLM Engine

Manages calling the model pipeline and linking this with the KV cache.

Tasks

• Init cache by profiling model memory usage at maximum sequence length input. See: Add tracking of memory allocations to Device huggingface/candle#1412
• Write the .generate function, the decoding entrypoint. The rough call chain looks like this, from vllm:
  1. .generate: batch the input seqs, calling .add_request
  • Write .add_request
  • Write .generate
  2. .run_engine: .step through each unfinished request, recording output
  • Write a .has_unfinished_requests
  • Write .run_engine
  3. .step: 1) call Scheduler to manage the seqs to swap for this decoding phase with ._schedule, then 2) execute the model
  • Write a simple ._schedule
  • Write .step
  4. .execute_model: Follow the cache ops from ._schedule, finally run the model pipeline with the cache.
  • Looks like the cache passes to the model pipeline has shared ownership with the CacheEngine. Will need to use Arc+my interior immutability setup.
  • Write .execute_model

Completed tasks

SequenceGroup:

Sequences generated from the same prompt.

Tasks

• Add a new seq grp to the Scheduler when a req is recieved.

  https://github.com/vllm-project/vllm/blob/60dc62dc9e53428912953276e0d12a034b353fb6/vllm/engine/llm_engine.py#L252

• Add to this when a new seq is generated

  https://github.com/vllm-project/vllm/blob/60dc62dc9e53428912953276e0d12a034b353fb6/vllm/engine/llm_engine.py#L429-L435

Cache:

CacheEngine manages the KV cache.

Tasks

• Ensure initialization is called

BlockSpaceManager

Managed blocks and allocation

Deps

• SequenceGroup

Tasks

• BlockAllocator
• PhysicalTokenBlock
• BlockSpaceManager
  • Allocation, in conjuction with a Scheduler impl

Scheduler

Scheduler schedules blocks to swap in/out, copy.

Deps

• SequenceGroup
• BlockSpaceManager

Tasks

• SchedulerConfig
• BlockSpaceManager (with some deps: 1) Block Allocator)
  • Allocation via BlockSpaceManager

[EricLBuehler]

_memory_efficient_attention is flash_attn when compiled with --features cuda, otherwise just normal scaled_dot_product_attention.

[EricLBuehler]

#13 added PagedAttention.

[EricLBuehler]

this commit adds most of the Scheduler functionality.

[EricLBuehler]

Closing this pending a rewrite of the scheduler and cache, see the scheduler branch. A new tracking issue will be reopened.