dynamo-llm 1.0.2

Dynamo LLM Library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

pub mod local;
pub mod logical;
pub mod view;

pub use local::LocalBlockData as BlockData;

pub trait BlockDataExt<S: Storage>: Send + Sync + 'static + std::fmt::Debug {
    /// The index of the block in the block set
    fn block_id(&self) -> BlockId;

    /// The identifier of the block set within the worker
    fn block_set_id(&self) -> usize;

    /// The identifier of the worker that owns the block
    /// Note: If the block is a logical block, this will be the worker id of the worker
    /// that owns the logical block, not the worker id of the worker that owns the physical block
    /// because their could be multiple workers contributing to the same logical block.
    fn worker_id(&self) -> WorkerID;

    /// The storage type of the block
    fn storage_type(&self) -> &StorageType;

    /// Whether the block is fully contiguous
    fn is_fully_contiguous(&self) -> bool;

    /// Returns the number of layers in the block
    fn num_layers(&self) -> usize;

    /// The size of the page in the block
    fn page_size(&self) -> usize;

    /// Returns the number of outer dimensions in the block
    fn num_outer_dims(&self) -> usize;

    fn num_inner_dims(&self) -> usize;

    /// Whether or not one can acquire read-only views to the block's storage
    fn is_local(&self) -> Option<&dyn BlockDataViews<S>>;

    /// Whether or not one can acquire mutable views to the block's storage
    fn is_local_mut(&mut self) -> Option<&mut dyn BlockDataViews<S>>;

    /// Get a read-only view of this block's storage for a layer
    fn layer_view(
        &self,
        layer_idx: usize,
        outer_idx: usize,
    ) -> BlockResult<view::LayerView<'_, S>> {
        match self.is_local() {
            Some(views) => views.local_layer_view(layer_idx, outer_idx),
            None => Err(BlockError::ViewsNotAvailableOnLogicalBlocks),
        }
    }

    /// Get a mutable view of this block's storage for a layer
    fn layer_view_mut(
        &mut self,
        layer_idx: usize,
        outer_idx: usize,
    ) -> BlockResult<view::LayerViewMut<'_, S>> {
        match self.is_local_mut() {
            Some(views) => views.local_layer_view_mut(layer_idx, outer_idx),
            None => Err(BlockError::ViewsNotAvailableOnLogicalBlocks),
        }
    }

    /// Get a read-only view of this block's storage
    fn block_view(&self) -> BlockResult<view::BlockView<'_, S>> {
        match self.is_local() {
            Some(views) => views.local_block_view(),
            None => Err(BlockError::ViewsNotAvailableOnLogicalBlocks),
        }
    }

    /// Get a mutable view of this block's storage
    fn block_view_mut(&mut self) -> BlockResult<view::BlockViewMut<'_, S>> {
        match self.is_local_mut() {
            Some(views) => views.local_block_view_mut(),
            None => Err(BlockError::ViewsNotAvailableOnLogicalBlocks),
        }
    }
}

pub trait BlockDataViews<S: Storage> {
    /// Get a read-only view of this block's storage for a layer
    fn local_layer_view(
        &self,
        layer_idx: usize,
        outer_idx: usize,
    ) -> BlockResult<view::LayerView<'_, S>>;

    /// Get a mutable view of this block's storage for a layer
    fn local_layer_view_mut(
        &mut self,
        layer_idx: usize,
        outer_idx: usize,
    ) -> BlockResult<view::LayerViewMut<'_, S>>;

    /// Get a read-only view of this block's storage
    fn local_block_view(&self) -> BlockResult<view::BlockView<'_, S>>;

    /// Get a mutable view of this block's storage
    fn local_block_view_mut(&mut self) -> BlockResult<view::BlockViewMut<'_, S>>;
}

pub trait BlockDataProvider: StorageTypeProvider {
    type Locality: LocalityProvider;

    fn block_data(&self) -> &impl BlockDataExt<Self::StorageType>;
}

pub trait BlockDataProviderMut: BlockDataProvider {
    type Locality: LocalityProvider;

    fn block_data_mut(&mut self) -> &mut impl BlockDataExt<Self::StorageType>;
}