workspace.py

# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
from ..config import DeepSpeedInferenceConfig
from .base import BaseOp

minus_inf = -10000.0
key_idx = 0
value_idx = 1


class InferenceContext:

    __instance = None

    def __init__(self):
        self.kv_cache = None
        self.kv_cache_elem_dtype = None
        self.num_tokens = 1
        self.kv_cache_num_layers = None
        self.kv_cache_size = None
        self.max_out_tokens = None
        self.rotary = None
        self.allocate_called = False
        self.static_shapes = True

    @classmethod
    def Instance(cls):
        if InferenceContext.__instance is None:
            InferenceContext.__instance = InferenceContext()
        return InferenceContext.__instance

    def gen_workspace(self, num_layers, num_heads, batch_size, prompt_len, hidden_dim, mp_size, external_cache,
                      elem_dtype, rank, max_out_tokens, min_out_tokens):
        self.allocate_called = True
        self.kv_cache = None
        if not external_cache:
            self.kv_cache_num_layers = num_layers
            self.max_out_tokens = max_out_tokens
            head_size = hidden_dim // num_heads
            self.kv_cache_size = torch.Size([batch_size, (num_heads // mp_size), max_out_tokens, head_size])
            self.kv_cache_elem_dtype = elem_dtype
        self.num_tokens = 0
        self.static_shapes = True
        return True

    def retake_workspace(self):
        return True

    def _retake_workspace(self):
        assert self.allocate_called, "retake workspace called before allocate workspace"

        import deepspeed.accelerator as accelerator
        if self.kv_cache is None:
            self.kv_cache = []
            for layer in range(self.kv_cache_num_layers):
                self.kv_cache.append((torch.zeros(self.kv_cache_size,
                                                  dtype=self.kv_cache_elem_dtype,
                                                  device=accelerator.get_accelerator().device_name()),
                                      torch.zeros(self.kv_cache_size,
                                                  dtype=self.kv_cache_elem_dtype,
                                                  device=accelerator.get_accelerator().device_name())))

        return True

    def update_cache(self, layer_id, token_idx, is_prompt, bat_0213_key, bat_0213_value):
        has_workspace = self._retake_workspace()
        assert has_workspace, "Could not allocate workspace"

        # Update current token
        if is_prompt:
            self.static_shapes = True
            if token_idx is None:
                self.static_shapes = False
                InferenceContext.Instance().reset_tokens(bat_0213_key.shape[2])
            else:
                InferenceContext.Instance().reset_tokens(token_idx)

        if token_idx is None:
            token_idx = InferenceContext.Instance().current_tokens()

        bsz = bat_0213_key.shape[0]

        # Update cache content
        if is_prompt:
            cache_max_seq = self.kv_cache_size[2]
            cache_max_head_dim = self.kv_cache_size[3]
            seq = bat_0213_key.shape[2]

            mask = torch.arange(cache_max_seq, device=bat_0213_key.device)
            mask = mask.ge(token_idx)
            mask = mask.unsqueeze(-1)
            mask = mask.expand([cache_max_seq, cache_max_head_dim])

            self.kv_cache[layer_id][key_idx][:bsz, :, :seq, :].copy_(bat_0213_key)
            self.kv_cache[layer_id][key_idx][:bsz, :].masked_fill_(mask, 0)
            self.kv_cache[layer_id][value_idx][:bsz, :, :seq, :].copy_(bat_0213_value)
            self.kv_cache[layer_id][value_idx][:bsz, :].masked_fill_(mask, 0)
        else:
            if self.static_shapes:
                assert type(token_idx) == torch.Tensor, "token_idx is expected to be torch.Tensor"
                self.kv_cache[layer_id][key_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_key)
                self.kv_cache[layer_id][value_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_value)
            else:
                assert type(token_idx) == int, "token_idx is expected to be int"
                self.kv_cache[layer_id][key_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_key
                self.kv_cache[layer_id][value_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_value

        bat_0213_key = self.kv_cache[layer_id][key_idx][:bsz]
        bat_0213_value = self.kv_cache[layer_id][value_idx][:bsz]

        if not self.static_shapes:
            bat_0213_key = bat_0213_key[:, :, :token_idx, :]
            bat_0213_value = bat_0213_value[:, :, :token_idx, :]

        return bat_0213_key, bat_0213_value

    def release_workspace(self):
        self.kv_cache = None
        self.rotary = None

    def reset_tokens(self, initial_tokens=1):
        self.num_tokens = initial_tokens

    def current_tokens(self):
        return self.num_tokens

    def advance_tokens(self):
        self.num_tokens = self.num_tokens + 1

    def get_kv_cache(self):
        return self.kv_cache

    def get_rotary(self, rotary_dim, rope_theta, device=None):
        if self.rotary is None:
            from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding

            self.rotary = LlamaRotaryEmbedding(rotary_dim, base=rope_theta, device=device)

        return self.rotary

    def get_max_tokens_num(self):
        return self.max_out_tokens


class WorkspaceOp(BaseOp):

    def __init__(self, config: DeepSpeedInferenceConfig = None):
        if config is None:
            config = DeepSpeedInferenceConfig()
        self.inference_context = InferenceContext.Instance()
        self._is_allocated = False
        try:
            super(WorkspaceOp, self).__init__(config)
            if config.dtype == torch.float32:
                self.allocate_workspace_func = self.inference_module.allocate_workspace_fp32
                repeat_kv_cache_name = "repeat_kv_cache_fp32"
            elif config.dtype == torch.bfloat16:
                self.allocate_workspace_func = self.inference_module.allocate_workspace_bf16
                repeat_kv_cache_name = "repeat_kv_cache_bf16"
            else:
                self.allocate_workspace_func = self.inference_module.allocate_workspace_fp16
                repeat_kv_cache_name = "repeat_kv_cache_fp16"
            self.repeat_kv_cache_func = getattr(self.inference_module, repeat_kv_cache_name, None)
            self.release_workspace_func = self.inference_module.release_workspace
            self.retake_workspace_func = self.inference_module.retake_workspace
            self.reset_cache_func = self.inference_module.reset_cache
        except (ValueError, AttributeError) as e:
            print(f"Using fallback functions in workspace because of {e}")
            if config.dtype == torch.float32:
                self.allocate_workspace_func = self.allocate_workspace_fp32_fallback
            elif config.dtype == torch.bfloat16:
                self.allocate_workspace_func = self.allocate_workspace_bf16_fallback
            else:
                self.allocate_workspace_func = self.allocate_workspace_fp16_fallback
            self.release_workspace_func = self.release_workspace_fallback
            self.retake_workspace_func = self.retake_workspace_fallback
            self.reset_cache_func = self.reset_cache_fallback
            self.repeat_kv_cache_func = self.repeat_kv_cache_fallback

    def allocate_workspace(self, *args, **kwargs):
        self._is_allocated = True
        return self.allocate_workspace_func(*args, **kwargs)

    def release_workspace(self):
        self._is_allocated = False
        return self.release_workspace_func()

    def reset_cache(self):
        return self.reset_cache_func() if self.reset_cache_func else None

    def retake_workspace(self):
        return self.retake_workspace_func() if self.retake_workspace_func else None

    def repeat_kv_cache(self, source_batch_size, repeats):
        if self.repeat_kv_cache_func is None:
            raise RuntimeError("Shared prefill requires rebuilding the transformer inference extension")
        return self.repeat_kv_cache_func(source_batch_size, repeats)

    def allocate_workspace_fp32_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size,
                                         external_cache, rank, max_out_tokens, min_out_tokens):
        return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim,
                                                    mp_size, external_cache, torch.float, rank, max_out_tokens,
                                                    min_out_tokens)

    def allocate_workspace_bf16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size,
                                         external_cache, rank, max_out_tokens, min_out_tokens):
        return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim,
                                                    mp_size, external_cache, torch.bfloat16, rank, max_out_tokens,
                                                    min_out_tokens)

    def allocate_workspace_fp16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size,
                                         external_cache, rank, max_out_tokens, min_out_tokens):
        return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim,
                                                    mp_size, external_cache, torch.half, rank, max_out_tokens,
                                                    min_out_tokens)

    def reset_cache_fallback(self):
        return self.inference_context.reset_tokens()

    def release_workspace_fallback(self):
        return self.inference_context.release_workspace()

    def retake_workspace_fallback(self):
        return self.inference_context.retake_workspace()

    def repeat_kv_cache_fallback(self, source_batch_size, repeats):
        target_batch_size = source_batch_size * repeats
        cache_size = self.inference_context.kv_cache_size
        if cache_size is None or cache_size[0] != target_batch_size:
            raise RuntimeError("KV cache repeat does not match the allocated workspace batch size")
        if self.inference_context.kv_cache is None:
            raise RuntimeError("KV cache repeat requires a completed prompt forward")
        current_tokens = self.inference_context.current_tokens()
        if current_tokens <= 1:
            raise RuntimeError("KV cache repeat requires a completed prompt forward")
        prompt_tokens = current_tokens - 1
        repeated_cache = []
        for key_cache, value_cache in self.inference_context.kv_cache:
            # Backward copies preserve source rows that overlap the expanded destination range.
            for destination in range(target_batch_size - 1, -1, -1):
                source = destination // repeats
                if source == destination:
                    continue
                key_cache[destination, :, :prompt_tokens, :].copy_(key_cache[source, :, :prompt_tokens, :])
                value_cache[destination, :, :prompt_tokens, :].copy_(value_cache[source, :, :prompt_tokens, :])
            repeated_cache.extend((key_cache[:, :, :prompt_tokens, :], value_cache[:, :, :prompt_tokens, :]))
        return repeated_cache

    def is_allocated(self):
        return self._is_allocated