kvbm_physical/transfer/
options.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Transfer options for configuring block and layer transfers.
5
6use super::BounceBuffer;
7use crate::layout::KvBlockLayout;
8use cudarc::driver::CudaStream;
9use derive_builder::Builder;
10use derive_getters::Dissolve;
11use std::ops::Range;
12use std::sync::Arc;
13
14/// Options for configuring transfer operations.
15///
16/// This structure provides configuration for block and layer transfers,
17/// including layer ranges, NIXL write notifications, and bounce buffers.
18///
19/// # Examples
20///
21/// ```rust,ignore
22/// let options = TransferOptions::builder()
23///     .nixl_write_notification(42)
24///     .layer_range(0..10)
25///     .build();
26/// ```
27#[derive(Clone, Default, Builder, Dissolve)]
28#[builder(pattern = "owned", default)]
29pub struct TransferOptions {
30    /// Range of layers to transfer (None = all layers).
31    ///
32    /// When specified, only the layers in this range will be transferred.
33    /// This is useful for partial block transfers or layer-specific operations.
34    #[builder(default, setter(strip_option))]
35    pub layer_range: Option<Range<usize>>,
36
37    /// NIXL write notification value delivered after RDMA write completes.
38    ///
39    /// When specified, NIXL will deliver this notification value to the remote
40    /// node after the RDMA write operation completes. This enables efficient
41    /// notification of transfer completion without requiring polling.
42    #[builder(default, setter(strip_option))]
43    pub nixl_write_notification: Option<u64>,
44
45    /// Bounce buffer specification for multi-hop transfers.
46    ///
47    /// When direct transfers are not allowed or efficient, this specifies
48    /// an intermediate staging area. The transfer will be split into two hops:
49    /// source → bounce buffer → destination.
50    #[builder(default, setter(strip_option, into))]
51    pub bounce_buffer: Option<BounceBuffer>,
52
53    /// Optional caller-provided CUDA stream for the transfer.
54    ///
55    /// When provided, the transfer executor will use this stream instead of
56    /// acquiring one from the pool. The caller is responsible for synchronization -
57    /// no event is recorded by the executor.
58    ///
59    /// This is useful for layer-wise transfers where all layers must execute
60    /// on the same stream to allow proper event sequencing.
61    #[builder(default, setter(strip_option))]
62    pub cuda_stream: Option<Arc<CudaStream>>,
63
64    /// Override source block layout interpretation.
65    ///
66    /// When set, the transfer executor will treat source blocks as having
67    /// this layout instead of the layout's default block_layout().
68    /// This enables transferring blocks that are stored in one format
69    /// but should be interpreted as another (e.g., operational → universal).
70    #[builder(default, setter(strip_option))]
71    pub src_kv_layout: Option<KvBlockLayout>,
72
73    /// Override destination block layout interpretation.
74    ///
75    /// When set, the transfer executor will treat destination blocks as having
76    /// this layout instead of the layout's default block_layout().
77    /// This enables writing blocks in a different format than the destination
78    /// layout's native format.
79    #[builder(default, setter(strip_option))]
80    pub dst_kv_layout: Option<KvBlockLayout>,
81}
82
83impl TransferOptions {
84    /// Create a new builder for transfer options.
85    pub fn builder() -> TransferOptionsBuilder {
86        TransferOptionsBuilder::default()
87    }
88
89    /// Create transfer options from an optional layer range.
90    pub fn from_layer_range(layer_range: Option<Range<usize>>) -> Self {
91        Self {
92            layer_range,
93            ..Self::default()
94        }
95    }
96
97    /// Create default transfer options.
98    ///
99    /// This transfers all layers with no special configuration.
100    pub fn new() -> Self {
101        Self::default()
102    }
103}
104
105#[cfg(all(test, feature = "testing-kvbm"))]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn test_default_options() {
111        let options = TransferOptions::default();
112        assert!(options.layer_range.is_none());
113        assert!(options.nixl_write_notification.is_none());
114        assert!(options.bounce_buffer.is_none());
115    }
116
117    #[test]
118    fn test_builder_with_notification() {
119        let options = TransferOptions::builder()
120            .nixl_write_notification(42)
121            .build()
122            .unwrap();
123
124        assert_eq!(options.nixl_write_notification, Some(42));
125        assert!(options.layer_range.is_none());
126    }
127
128    #[test]
129    fn test_builder_with_layer_range() {
130        let options = TransferOptions::builder()
131            .layer_range(0..10)
132            .build()
133            .unwrap();
134
135        assert_eq!(options.layer_range, Some(0..10));
136        assert!(options.nixl_write_notification.is_none());
137    }
138
139    #[test]
140    fn test_builder_with_all_options() {
141        let options = TransferOptions::builder()
142            .nixl_write_notification(100)
143            .layer_range(5..15)
144            .build()
145            .unwrap();
146
147        assert_eq!(options.nixl_write_notification, Some(100));
148        assert_eq!(options.layer_range, Some(5..15));
149    }
150}
kvbm_physical/transfer/options.rs

kvbm_physical/transfer/
options.rs
