protocols.rs 15.2 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
use derive_builder::Builder;
5
use serde::{Deserialize, Serialize};
6
use std::collections::{HashMap, HashSet};
7
8
use std::path::{Path, PathBuf};
use std::sync::Arc;
9
use uuid::Uuid;
10
use validator::Validate;
11

12
use crate::common::perf_model::PerfModel;
13
use dynamo_kv_router::protocols::KvCacheEvent;
14
15
use dynamo_tokens::blocks::UniqueBlock;
use dynamo_tokens::{BlockHash, SequenceHash, Token};
16

17
18
19
/// Trait for publishing KV cache events.
/// This abstracts the runtime dependency so mocker components can remain generic.
pub trait KvCacheEventSink: Send + Sync {
20
21
22
23
24
    fn publish(
        &self,
        event: KvCacheEvent,
        block_token_ids: Option<&[Vec<u32>]>,
    ) -> anyhow::Result<()>;
25
26
}

27
28
29
pub type NumBlocks = usize;

/// Represents different block movement operations in the cache
Yan Ru Pei's avatar
Yan Ru Pei committed
30
/// For Use and Promote variants, block hashes are included for KV event publishing
31
32
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum MoveBlock {
33
    Use(Vec<UniqueBlock>, Vec<BlockHash>, Option<Vec<Vec<u32>>>),
34
35
    Destroy(Vec<UniqueBlock>),
    Deref(Vec<UniqueBlock>),
36
    Promote(Uuid, SequenceHash, Option<u64>, BlockHash, Option<Vec<u32>>),
37
38
39
40
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum MoveBlockResponse {
41
42
    Store(Vec<SequenceHash>, Option<u64>),
    Remove(Vec<SequenceHash>),
43
44
45
46
47
48
49
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DirectRequest {
    pub tokens: Vec<Token>,
    pub max_output_tokens: usize,
    pub uuid: Option<Uuid>,
Yan Ru Pei's avatar
Yan Ru Pei committed
50
    pub dp_rank: u32,
51
52
53
54
55
}

/// Represents the cost of prefilling content in the cache
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrefillCost {
56
    pub new_blocks: usize,
57
    pub new_tokens: usize,
58
59
60
}

impl PrefillCost {
61
62
63
64
65
    pub fn predict_prefill_compute(
        &self,
        new_tokens: Option<usize>,
        perf_model: &PerfModel,
    ) -> f64 {
66
        let tokens = new_tokens.unwrap_or(self.new_tokens);
67
        perf_model.predict_prefill_time(tokens)
68
    }
69
70
}

71
72
73
74
75
76
77
/// Signal for output token generation with completion status
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OutputSignal {
    pub uuid: Uuid,
    pub completed: bool,
}

78
79
80
81
82
83
84
85
86
87
88
89
/// Worker type for disaggregated serving configurations
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum WorkerType {
    /// Standard aggregated worker handling both prefill and decode
    #[default]
    Aggregated,
    /// Dedicated prefill worker in disaggregated mode
    Prefill,
    /// Dedicated decode worker in disaggregated mode
    Decode,
}

Yan Ru Pei's avatar
Yan Ru Pei committed
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
/// Configuration for reasoning/thinking token output in the mocker.
///
/// When set, the mocker wraps the first portion of each response in thinking
/// boundary tokens: `[start_token, random..., end_token, random...]`.
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
pub struct ReasoningConfig {
    pub start_thinking_token_id: u32,
    pub end_thinking_token_id: u32,
    #[validate(range(min = 0.0, max = 1.0))]
    pub thinking_ratio: f64,
}

impl ReasoningConfig {
    /// Number of thinking tokens (including start/end boundaries) for a given osl.
    /// Returns 0 if osl < 2 (thinking disabled). Otherwise clamps to [2, osl].
    pub fn num_thinking_tokens(&self, max_output_tokens: usize) -> usize {
        if max_output_tokens < 2 {
            return 0;
        }
        let raw = (max_output_tokens as f64 * self.thinking_ratio).floor() as usize;
        if raw == 0 {
            return 0;
        }
        raw.max(2).min(max_output_tokens)
    }

    /// Number of response tokens after the thinking block.
    pub fn num_response_tokens(&self, max_output_tokens: usize) -> usize {
        max_output_tokens.saturating_sub(self.num_thinking_tokens(max_output_tokens))
    }
}

122
/// Configuration arguments for MockVllmEngine
123
#[derive(Debug, Clone, Serialize, Deserialize, Builder, Validate)]
124
125
126
#[builder(pattern = "owned", build_fn(public))]
pub struct MockEngineArgs {
    #[builder(default = "16384")]
127
    #[validate(range(min = 1))]
128
129
130
    pub num_gpu_blocks: usize,

    #[builder(default = "64")]
131
    #[validate(range(min = 2))]
132
133
134
135
    pub block_size: usize,

    // This was 1024 in the past but reverted back to 256
    #[builder(default = Some(256))]
136
    #[validate(range(min = 1))]
137
138
139
140
    pub max_num_seqs: Option<usize>,

    // default for open api server, for llm class it's 16384
    #[builder(default = Some(8192))]
141
    #[validate(range(min = 1))]
142
143
144
145
146
    pub max_num_batched_tokens: Option<usize>,

    #[builder(default = true)]
    pub enable_prefix_caching: bool,

147
148
149
    #[builder(default = true)]
    pub enable_chunked_prefill: bool,

150
    #[builder(default = "0.01")]
151
    #[validate(range(min = 0.0, max = 1.0))]
152
153
154
    pub watermark: f64,

    #[builder(default = "1.0")]
155
    #[validate(range(min = 0.0))]
156
157
158
    pub speedup_ratio: f64,

    #[builder(default = "1")]
159
    #[validate(range(min = 1))]
160
    pub dp_size: u32,
161
162
163

    /// Optional startup time in seconds to simulate engine initialization delay
    #[builder(default = "None")]
164
    #[validate(range(min = 0.0))]
165
    pub startup_time: Option<f64>,
166
167
168
169

    /// Worker type for disaggregated serving (Aggregated, Prefill, or Decode)
    #[builder(default = "WorkerType::Aggregated")]
    pub worker_type: WorkerType,
170
171
172
173
174

    /// Performance model for timing predictions (not serialized, loaded from planner_profile_data)
    #[serde(skip)]
    #[builder(default = "Arc::new(PerfModel::default())")]
    pub perf_model: Arc<PerfModel>,
175
176
177
178

    /// Enable worker-local KV indexer for tracking this worker's own KV cache state
    #[builder(default = "false")]
    pub enable_local_indexer: bool,
179
180
181
182
183
184

    /// Bootstrap port for disaggregated serving rendezvous.
    /// Prefill workers listen on this port; decode workers connect to it.
    /// If None, bootstrap rendezvous is disabled.
    #[builder(default = "None")]
    pub bootstrap_port: Option<u16>,
Yan Ru Pei's avatar
Yan Ru Pei committed
185

186
187
188
189
190
191
192
193
194
195
196
197
    /// KV cache bytes per token, auto-computed from model config by Python CLI.
    /// Formula: num_layers * 2 * num_kv_heads * head_dim * dtype_bytes
    #[builder(default = "None")]
    pub kv_bytes_per_token: Option<usize>,

    /// KV cache transfer bandwidth in GB/s for disaggregated serving latency simulation.
    /// Default: 64.0 (inter-node InfiniBand). Set to 0 to disable KV transfer delay.
    /// For intra-node NVLink, typical value is ~450.
    #[builder(default = "None")]
    #[validate(range(min = 0.0))]
    pub kv_transfer_bandwidth: Option<f64>,

Yan Ru Pei's avatar
Yan Ru Pei committed
198
199
200
201
    /// Reasoning/thinking token configuration.
    /// When set, the mocker wraps output in thinking boundary tokens.
    #[builder(default = "None")]
    pub reasoning: Option<ReasoningConfig>,
202
203
204
205
206
207

    /// ZMQ port for publishing KV events in vLLM's native wire format.
    /// When set, the scheduler publishes to a ZMQ PUB socket instead of directly to NATS.
    /// A KvEventPublisher relay subscribes to this socket and forwards events to NATS.
    #[builder(default = "None")]
    pub zmq_kv_events_port: Option<u16>,
208
209
}

210
211
212
213
214
215
216
217
impl Default for MockEngineArgs {
    fn default() -> MockEngineArgs {
        MockEngineArgsBuilder::default()
            .build()
            .expect("Failed to build default MockEngineArgs")
    }
}

218
219
220
221
impl MockEngineArgs {
    pub fn builder() -> MockEngineArgsBuilder {
        MockEngineArgsBuilder::default()
    }
222

223
224
225
226
227
228
229
230
231
232
233
234
    pub fn is_prefill(&self) -> bool {
        self.worker_type == WorkerType::Prefill
    }

    pub fn is_decode(&self) -> bool {
        self.worker_type == WorkerType::Decode
    }

    pub fn needs_kv_publisher(&self) -> bool {
        self.enable_prefix_caching && !self.is_decode()
    }

235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    /// Create MockEngineArgs from a JSON file containing extra engine arguments
    pub fn from_json_file(path: &Path) -> anyhow::Result<Self> {
        let mut builder = Self::builder();

        // Load and parse the JSON file
        let file_content = std::fs::read_to_string(path)?;
        let extra_args: HashMap<String, serde_json::Value> = serde_json::from_str(&file_content)?;

        // Define valid field names
        let valid_fields: HashSet<&str> = [
            "num_gpu_blocks",
            "block_size",
            "max_num_seqs",
            "max_num_batched_tokens",
            "enable_prefix_caching",
250
            "enable_chunked_prefill",
251
252
253
            "watermark",
            "speedup_ratio",
            "dp_size",
254
            "startup_time",
255
256
            "is_prefill",
            "is_decode",
257
            "planner_profile_data",
258
            "enable_local_indexer",
259
            "bootstrap_port",
260
261
            "kv_bytes_per_token",
            "kv_transfer_bandwidth",
Yan Ru Pei's avatar
Yan Ru Pei committed
262
            "reasoning",
263
            "zmq_kv_events_port",
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
        ]
        .iter()
        .cloned()
        .collect();

        // Check for invalid arguments
        let invalid_args: Vec<String> = extra_args
            .keys()
            .filter(|key| !valid_fields.contains(key.as_str()))
            .cloned()
            .collect();

        if !invalid_args.is_empty() {
            return Err(anyhow::anyhow!(
                "Invalid arguments found in JSON file: {}. Valid arguments are: {:?}",
                invalid_args.join(", "),
                valid_fields
            ));
        }

        // Apply each extra argument to the builder
285
286
287
288
        if let Some(value) = extra_args.get("num_gpu_blocks")
            && let Some(num) = value.as_u64()
        {
            builder = builder.num_gpu_blocks(num as usize);
289
290
        }

291
292
293
294
        if let Some(value) = extra_args.get("block_size")
            && let Some(num) = value.as_u64()
        {
            builder = builder.block_size(num as usize);
295
296
        }

297
298
299
300
        if let Some(value) = extra_args.get("max_num_seqs")
            && let Some(num) = value.as_u64()
        {
            builder = builder.max_num_seqs(Some(num as usize));
301
302
        }

303
304
305
306
        if let Some(value) = extra_args.get("max_num_batched_tokens")
            && let Some(num) = value.as_u64()
        {
            builder = builder.max_num_batched_tokens(Some(num as usize));
307
308
        }

309
310
311
312
        if let Some(value) = extra_args.get("enable_prefix_caching")
            && let Some(enabled) = value.as_bool()
        {
            builder = builder.enable_prefix_caching(enabled);
313
314
        }

315
316
317
318
        if let Some(value) = extra_args.get("enable_chunked_prefill")
            && let Some(enabled) = value.as_bool()
        {
            builder = builder.enable_chunked_prefill(enabled);
319
320
        }

321
322
323
324
        if let Some(value) = extra_args.get("watermark")
            && let Some(num) = value.as_f64()
        {
            builder = builder.watermark(num);
325
326
        }

327
328
329
330
        if let Some(value) = extra_args.get("speedup_ratio")
            && let Some(num) = value.as_f64()
        {
            builder = builder.speedup_ratio(num);
331
332
        }

333
334
335
336
        if let Some(value) = extra_args.get("dp_size")
            && let Some(num) = value.as_u64()
        {
            builder = builder.dp_size(num as u32);
337
338
        }

339
340
341
342
343
344
        if let Some(value) = extra_args.get("startup_time")
            && let Some(num) = value.as_f64()
        {
            builder = builder.startup_time(Some(num));
        }

345
346
347
348
349
350
        if let Some(value) = extra_args.get("enable_local_indexer")
            && let Some(enabled) = value.as_bool()
        {
            builder = builder.enable_local_indexer(enabled);
        }

351
352
353
354
355
356
        if let Some(value) = extra_args.get("bootstrap_port")
            && let Some(port) = value.as_u64()
        {
            builder = builder.bootstrap_port(Some(port as u16));
        }

357
358
359
360
361
362
363
364
365
366
367
368
        if let Some(value) = extra_args.get("kv_bytes_per_token")
            && let Some(num) = value.as_u64()
        {
            builder = builder.kv_bytes_per_token(Some(num as usize));
        }

        if let Some(value) = extra_args.get("kv_transfer_bandwidth")
            && let Some(num) = value.as_f64()
        {
            builder = builder.kv_transfer_bandwidth(Some(num));
        }

Yan Ru Pei's avatar
Yan Ru Pei committed
369
370
371
372
373
374
        if let Some(value) = extra_args.get("reasoning") {
            let cfg: ReasoningConfig = serde_json::from_value(value.clone())
                .map_err(|e| anyhow::anyhow!("Failed to parse reasoning config: {}", e))?;
            builder = builder.reasoning(Some(cfg));
        }

375
376
377
378
379
380
        if let Some(value) = extra_args.get("zmq_kv_events_port")
            && let Some(port) = value.as_u64()
        {
            builder = builder.zmq_kv_events_port(Some(port as u16));
        }

381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
        // Parse worker type from is_prefill and is_decode flags
        let is_prefill = extra_args
            .get("is_prefill")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);
        let is_decode = extra_args
            .get("is_decode")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        // Determine worker type based on flags
        let worker_type = match (is_prefill, is_decode) {
            (false, false) => WorkerType::Aggregated,
            (true, false) => WorkerType::Prefill,
            (false, true) => WorkerType::Decode,
            (true, true) => panic!(
                "Invalid worker configuration: is_prefill and is_decode cannot both be true. \
                 Worker must be either Aggregated (both false), Prefill (is_prefill=true), or Decode (is_decode=true)."
            ),
        };
        builder = builder.worker_type(worker_type);

403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
        // Load performance model from NPZ file if provided
        let perf_model = if let Some(path_str) = extra_args.get("planner_profile_data")
            && let Some(path_str) = path_str.as_str()
        {
            let npz_path = PathBuf::from(path_str);
            match PerfModel::from_npz(&npz_path) {
                Ok(model) => {
                    tracing::info!("Successfully loaded performance model from: {:?}", npz_path);
                    Arc::new(model)
                }
                Err(e) => {
                    tracing::error!(
                        "Failed to load performance model from {:?}: {}. Falling back to polynomial model.",
                        npz_path,
                        e
                    );
                    Arc::new(PerfModel::default())
                }
            }
        } else {
            Arc::new(PerfModel::default())
        };
        builder = builder.perf_model(perf_model);

427
428
429
430
431
        // Build the MockEngineArgs with either defaults or overridden values
        builder
            .build()
            .map_err(|e| anyhow::anyhow!("Failed to build MockEngineArgs: {}", e))
    }
432
433
}

434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_unique_block_default_uniqueness() {
        // Create 10 default UniqueBlock instances
        let blocks: Vec<UniqueBlock> = (0..10).map(|_| UniqueBlock::default()).collect();

        // Extract UUIDs from each block
        let mut uuids = Vec::new();
        for block in blocks {
            match block {
                UniqueBlock::PartialBlock(uuid) => uuids.push(uuid),
                _ => panic!("Expected UuidIdentifier variant"),
            }
        }

        // Check that all UUIDs are unique by comparing each with every other
        for i in 0..uuids.len() {
            for j in i + 1..uuids.len() {
                assert_ne!(
                    uuids[i], uuids[j],
                    "UUID at index {} and {} are identical",
                    i, j
                );
            }
        }
    }
}