kvbm_physical/layout/
config.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use derive_builder::Builder;
5use serde::{Deserialize, Serialize};
6use validator::{Validate, ValidationError};
7
8/// Configuration for block layouts.
9///
10/// The `#[validate]` attributes on fields are checked during layout construction
11/// (e.g., `FullyContiguousLayout::new_internal()`, `LayerSeparateLayout::new_internal()`),
12/// not at builder `.build()` time.
13#[derive(Debug, Clone, Builder, Validate, Serialize, Deserialize, PartialEq, Eq)]
14pub struct LayoutConfig {
15    /// Number of blocks
16    #[validate(range(min = 1))]
17    pub num_blocks: usize,
18
19    /// Number of layers
20    #[validate(range(min = 1))]
21    pub num_layers: usize,
22
23    /// Number of outer dimensions
24    #[validate(range(min = 1, max = 2))]
25    pub outer_dim: usize,
26
27    /// Page size
28    #[validate(range(min = 1))]
29    pub page_size: usize,
30
31    /// Inner dimension
32    #[validate(range(min = 1))]
33    pub inner_dim: usize,
34
35    /// Alignment
36    #[validate(custom(function = "validate_power_of_2"))]
37    #[builder(default = "1")]
38    pub alignment: usize,
39
40    /// Data type
41    #[validate(custom(function = "validate_dtype_width_bytes"))]
42    #[builder(default = "2")]
43    pub dtype_width_bytes: usize,
44
45    /// Number of attention heads (optional).
46    ///
47    /// When provided, enables KvBlockLayout support for universal formats.
48    /// The head dimension can be computed as: `inner_dim / (page_size * num_heads)`.
49    ///
50    /// Required for:
51    /// - Universal layout transformations
52    /// - Per-head memory region access
53    #[builder(default = "None")]
54    #[serde(default)]
55    pub num_heads: Option<usize>,
56}
57
58impl LayoutConfig {
59    /// Builder for LayoutConfig
60    pub fn builder() -> LayoutConfigBuilder {
61        LayoutConfigBuilder::default()
62    }
63
64    pub fn required_bytes(&self) -> usize {
65        self.num_blocks
66            .saturating_mul(self.num_layers)
67            .saturating_mul(self.outer_dim)
68            .saturating_mul(self.page_size)
69            .saturating_mul(self.inner_dim)
70            .saturating_mul(self.dtype_width_bytes)
71    }
72
73    /// Get the number of bytes per block.
74    ///
75    /// This is the total size of a single block across all layers and outer dimensions.
76    pub fn bytes_per_block(&self) -> usize {
77        self.num_layers
78            .saturating_mul(self.outer_dim)
79            .saturating_mul(self.page_size)
80            .saturating_mul(self.inner_dim)
81            .saturating_mul(self.dtype_width_bytes)
82    }
83
84    /// Get the head dimension if `num_heads` is specified.
85    ///
86    /// Computes `inner_dim / (page_size * num_heads)`.
87    ///
88    /// # Returns
89    /// `Some(head_dim)` if `num_heads` is set, `None` otherwise.
90    pub fn head_dim(&self) -> Option<usize> {
91        self.num_heads.map(|nh| {
92            let divisor = self.page_size * nh;
93            self.inner_dim.checked_div(divisor).unwrap_or(0)
94        })
95    }
96
97    /// Check if this config supports KvBlockLayout operations.
98    ///
99    /// Returns `true` if `num_heads` is set and the dimensions are valid
100    /// (inner_dim is evenly divisible by page_size * num_heads).
101    pub fn supports_kv_block_layout(&self) -> bool {
102        if let Some(nh) = self.num_heads {
103            let divisor = self.page_size * nh;
104            divisor > 0 && self.inner_dim.is_multiple_of(divisor)
105        } else {
106            false
107        }
108    }
109
110    /// Validate that this config supports KvBlockLayout operations.
111    ///
112    /// # Returns
113    /// `Ok(())` if valid, `Err` with details otherwise.
114    pub fn validate_for_kv_block_layout(&self) -> Result<(), ValidationError> {
115        let nh = match self.num_heads {
116            Some(nh) => nh,
117            None => {
118                return Err(ValidationError::new(
119                    "num_heads_required_for_kv_block_layout",
120                ));
121            }
122        };
123
124        if nh == 0 {
125            return Err(ValidationError::new("num_heads_must_be_positive"));
126        }
127
128        let divisor = self.page_size * nh;
129        if !self.inner_dim.is_multiple_of(divisor) {
130            return Err(ValidationError::new(
131                "inner_dim_must_be_divisible_by_page_size_times_num_heads",
132            ));
133        }
134
135        Ok(())
136    }
137}
138
139/// The first two dimensions of the tensor, `shape[0]` and `shape[1]`, one of those corresponds to the
140/// block dimension, while the other corresponds to the outer dimension.
141///
142/// The outer dimension is typically:
143/// - 1: MLA or K and V stored together,
144/// - 2: K and V stored separately,
145///
146/// The block dimension tell us the number of blocks.
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148pub enum BlockDimension {
149    /// The block dimension is the first dimension of the tensor, `[n_blocks, outer_dim, inner_dim]`
150    BlockIsFirstDim,
151
152    /// The block dimension is the second dimension of the tensor, `[outer_dim, n_blocks, inner_dim]`
153    /// This is a replacement for v1's `outer_contiguous` is true.
154    BlockIsSecondDim,
155}
156
157/// Validation function for Option<usize> to check if it's Some(power_of_2).
158pub fn validate_power_of_2(alignment: usize) -> Result<(), ValidationError> {
159    if !alignment.is_power_of_two() {
160        // Return validation error if alignment is not a power of 2
161        return Err(validator::ValidationError::new(
162            "alignment_must_be_power_of_2",
163        ));
164    }
165    // Passes validation if alignment is a power of 2
166    Ok(())
167}
168
169pub fn validate_dtype_width_bytes(dtype_width_bytes: usize) -> Result<(), ValidationError> {
170    if !dtype_width_bytes.is_power_of_two() || !(2..=8).contains(&dtype_width_bytes) {
171        return Err(validator::ValidationError::new(
172            "dtype_width_bytes_must_be_power_of_two_and_less_than_8_bytes",
173        ));
174    }
175    Ok(())
176}
kvbm_physical/layout/config.rs

kvbm_physical/layout/
config.rs
