Add files using upload-large-folder tool

5328bed verified about 1 month ago

3.28 kB

# Copyright (c) 2025 SandAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import List

import numpy as np
import torch


@dataclass(frozen=True)
class PackedCoreAttnParams:
    # Packed sequence parameters for core_attn
    q_range: torch.Tensor
    k_range: torch.Tensor
    np_q_range: np.ndarray
    np_k_range: np.ndarray
    max_seqlen_q: int
    max_seqlen_k: int


@dataclass(frozen=True)
class PackedCrossAttnParams:
    # Packed sequence parameters for cross_attn
    q_ranges: torch.Tensor = None
    kv_ranges: torch.Tensor = None
    cu_seqlens_q: torch.Tensor = None
    cu_seqlens_kv: torch.Tensor = None
    max_seqlen_q: int = None
    max_seqlen_kv: int = None


@dataclass(frozen=True)
class ModelMetaArgs:
    H: int
    W: int
    cp_pad_size: int
    cp_split_sizes: List[int]
    slice_point: int
    denoising_range_num: int
    range_num: int
    extract_prefix_video_feature: bool
    fwd_extra_1st_chunk: bool
    distill_nearly_clean_chunk: bool
    clip_token_nums: int
    enable_cuda_graph: bool
    core_attn_params: PackedCoreAttnParams
    cross_attn_params: PackedCrossAttnParams
    timestep: torch.Tensor
    get_attn_weights_layer_num: int
    save_kvcache_every_forward: bool
    cur_denoise_step: int
    # Includes all chunks of the current sequence
    start_chunk_id: int
    end_chunk_id: int
    compress_kv: bool    # use kv cache compression or not
    total_cache_len: int
    budget_cache_len: int
    chunk_num: int
    debug: bool
    near_clean_chunk_idx: int

class InferenceParams:
    """Inference parameters that are passed to the main model in order
    to efficienly calculate and store the context during inference."""

    def __init__(self, max_batch_size, max_sequence_length):
        self.max_sequence_length = max_sequence_length
        self.max_batch_size = max_batch_size
        self.sequence_len_offset = 0
        self.key_value_memory_dict = {}
        self.update_kv_cache = False

        self.kv_compressed = False

    def swap_key_value_dict(self, batch_idx):
        "swap between batches"
        if len(self.key_value_memory_dict) == 0:
            raise ValueError("should not swap when dict in empty")

        for layer_number in self.key_value_memory_dict.keys():
            inference_key_memory, inference_value_memory = self.key_value_memory_dict[layer_number]
            assert len(batch_idx) == inference_key_memory.shape[1]  # make sure batch size is the same
            new_inference_key_memory = inference_key_memory[:, batch_idx]
            new_inference_value_memory = inference_value_memory[:, batch_idx]
            self.key_value_memory_dict[layer_number] = (new_inference_key_memory, new_inference_value_memory)

	# Copyright (c) 2025 SandAI. All Rights Reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	from dataclasses import dataclass
	from typing import List

	import numpy as np
	import torch


	@dataclass(frozen=True)
	class PackedCoreAttnParams:
	# Packed sequence parameters for core_attn
	q_range: torch.Tensor
	k_range: torch.Tensor
	np_q_range: np.ndarray
	np_k_range: np.ndarray
	max_seqlen_q: int
	max_seqlen_k: int


	@dataclass(frozen=True)
	class PackedCrossAttnParams:
	# Packed sequence parameters for cross_attn
	q_ranges: torch.Tensor = None
	kv_ranges: torch.Tensor = None
	cu_seqlens_q: torch.Tensor = None
	cu_seqlens_kv: torch.Tensor = None
	max_seqlen_q: int = None
	max_seqlen_kv: int = None


	@dataclass(frozen=True)
	class ModelMetaArgs:
	H: int
	W: int
	cp_pad_size: int
	cp_split_sizes: List[int]
	slice_point: int
	denoising_range_num: int
	range_num: int
	extract_prefix_video_feature: bool
	fwd_extra_1st_chunk: bool
	distill_nearly_clean_chunk: bool
	clip_token_nums: int
	enable_cuda_graph: bool
	core_attn_params: PackedCoreAttnParams
	cross_attn_params: PackedCrossAttnParams
	timestep: torch.Tensor
	get_attn_weights_layer_num: int
	save_kvcache_every_forward: bool
	cur_denoise_step: int
	# Includes all chunks of the current sequence
	start_chunk_id: int
	end_chunk_id: int
	compress_kv: bool # use kv cache compression or not
	total_cache_len: int
	budget_cache_len: int
	chunk_num: int
	debug: bool
	near_clean_chunk_idx: int

	class InferenceParams:
	"""Inference parameters that are passed to the main model in order
	to efficienly calculate and store the context during inference."""

	def __init__(self, max_batch_size, max_sequence_length):
	self.max_sequence_length = max_sequence_length
	self.max_batch_size = max_batch_size
	self.sequence_len_offset = 0
	self.key_value_memory_dict = {}
	self.update_kv_cache = False

	self.kv_compressed = False

	def swap_key_value_dict(self, batch_idx):
	"swap between batches"
	if len(self.key_value_memory_dict) == 0:
	raise ValueError("should not swap when dict in empty")

	for layer_number in self.key_value_memory_dict.keys():
	inference_key_memory, inference_value_memory = self.key_value_memory_dict[layer_number]
	assert len(batch_idx) == inference_key_memory.shape[1] # make sure batch size is the same
	new_inference_key_memory = inference_key_memory[:, batch_idx]
	new_inference_value_memory = inference_value_memory[:, batch_idx]
	self.key_value_memory_dict[layer_number] = (new_inference_key_memory, new_inference_value_memory)