Jelajahi Sumber

feat: per-recipe shm_size + DeepSeek V4 Flash recipe

Three changes:

1. Add `shm_size` to Recipe (default '10g').  Larger contexts need
   much more shared memory than the default 10g.

2. Make vLLM's `--tensor-parallel-size` derive from
   `recipe.tensor_parallel_size` (was hardcoded to 1).  This lets
   recipes configure multi-node TP.  SGLang already used this field.

3. Add `recipes/deepseek-v4-flash` translating the Docker Compose at
   https://github.com/MiaAI-Lab/DeepSeek-V4-Flash-Dual-DGX-Spark-1M-Context/blob/main/docker-compose.yml
   into a recipe (2 nodes, TP=2 PP=1, 1M context, 64gb shm).
Alexander Korolev 4 bulan lalu
induk
melakukan
7ba894be67
6 mengubah file dengan 82 tambahan dan 10 penghapusan
  1. 50 0
      recipes/deepseek-v4-flash/recipe.toml
  2. 1 0
      src/before_serve.rs
  3. 8 0
      src/config.rs
  4. 10 10
      src/docker/mod.rs
  5. 12 0
      src/docker/tests.rs
  6. 1 0
      src/doctor.rs

+ 50 - 0
recipes/deepseek-v4-flash/recipe.toml

@@ -0,0 +1,50 @@
+engine = "vllm"
+name = "DeepSeek V4 Flash (vLLM)"
+docker_image = "aidendle94/sparkrun-vllm-ds4-gb10:production-ready"
+model = "deepseek-ai/DeepSeek-V4-Flash"
+port = 8000
+required_nodes = 2
+tensor_parallel_size = 2
+pipeline_parallel_size = 1
+distributed_backend = "mp"
+
+# DeepSeek V4 with 1M context needs much more shared memory than the
+# default 10g.  KV cache + activations for 1M tokens fills up fast.
+shm_size = "64gb"
+
+args = [
+    "--kv-cache-dtype", "fp8",
+    "--block-size", "256",
+    "--max-model-len", "1000000",
+    "--max-num-seqs", "6",
+    "--max-num-batched-tokens", "8192",
+    "--gpu-memory-utilization", "0.82",
+    "--enable-prefix-caching",
+    "--speculative-config", "{\"method\":\"mtp\",\"num_speculative_tokens\":2}",
+    "--tokenizer-mode", "deepseek_v4",
+    "--tool-call-parser", "deepseek_v4",
+    "--enable-auto-tool-choice",
+    "--reasoning-parser", "deepseek_v4",
+    "--reasoning-config", "{\"reasoning_parser\":\"deepseek_v4\",\"reasoning_start_str\":\"\",\"reasoning_end_str\":\"\"}",
+    "--default-chat-template-kwargs.thinking", "true",
+    "--default-chat-template-kwargs.reasoning_effort", "high",
+    "--enable-flashinfer-autotune",
+    "--host", "0.0.0.0",
+    "--port", "8000",
+    "--trust-remote-code",
+    "--served-model-name", "deepseek-v4-flash",
+]
+
+[env]
+# HuggingFace / cache config
+HF_HUB_OFFLINE = "0"
+VLLM_CACHE_ROOT = "/cache/huggingface/vllm-cache"
+
+# vLLM features for DeepSeek V4 on Blackwell (GB10 / sm_120)
+VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"
+VLLM_USE_B12X_MOE = "1"
+VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"
+
+# Compile-time arch hints for PyTorch and FlashInfer on GB10
+TORCH_CUDA_ARCH_LIST = "12.1a"
+FLASHINFER_CUDA_ARCH_LIST = "12.1a"

+ 1 - 0
src/before_serve.rs

@@ -237,6 +237,7 @@ mod tests {
             entrypoint: vec![],
             distributed_backend: "mp".to_string(),
             before_serve: vec![],
+            shm_size: "10g".to_string(),
         }
     }
 

+ 8 - 0
src/config.rs

@@ -117,6 +117,10 @@ pub struct Recipe {
     pub entrypoint: Vec<String>,
     #[serde(default = "default_distributed_backend")]
     pub distributed_backend: String,
+    /// Shared memory size for the container (e.g. "10g", "64gb").
+    /// Default: "10g".  Larger contexts (1M tokens) need much more.
+    #[serde(default = "default_shm_size")]
+    pub shm_size: String,
     /// Recipe-level pre-serve commands (appended after cluster-level).
     #[serde(default)]
     pub before_serve: Vec<PreServeCommand>,
@@ -130,6 +134,10 @@ fn default_distributed_backend() -> String {
     "mp".to_string()
 }
 
+fn default_shm_size() -> String {
+    "10g".to_string()
+}
+
 fn default_tensor_parallel_size() -> u32 {
     1
 }

+ 10 - 10
src/docker/mod.rs

@@ -358,17 +358,17 @@ fn build_docker_command(
 ) -> Vec<String> {
     let is_sglang = recipe.engine == "sglang";
 
-    let mut flags = vec![
-        "--gpus all",
-        "--ipc=host",
-        "--privileged",
-        "--network=host",
-        "--shm-size 10g",
-        "--ulimit memlock=-1",
-        "--ulimit stack=67108864",
+    let mut flags: Vec<String> = vec![
+        "--gpus all".to_string(),
+        "--ipc=host".to_string(),
+        "--privileged".to_string(),
+        "--network=host".to_string(),
+        format!("--shm-size {}", recipe.shm_size.as_str()),
+        "--ulimit memlock=-1".to_string(),
+        "--ulimit stack=67108864".to_string(),
     ];
     if !keep {
-        flags.push("--rm");
+        flags.push("--rm".to_string());
     }
     let mut cmd = vec![
         format!("docker run -d {}", flags.join(" ")),
@@ -448,7 +448,7 @@ fn build_docker_command(
         ],
         "ray" | "mp" => vec![
             "--tensor-parallel-size".into(),
-            "1".into(),
+            recipe.tensor_parallel_size.to_string(),
             "--pipeline-parallel-size".into(),
             pp_size.to_string(),
             "--nnodes".into(),

+ 12 - 0
src/docker/tests.rs

@@ -52,6 +52,7 @@ fn trec() -> Recipe {
         entrypoint: vec![],
         distributed_backend: "mp".into(),
         before_serve: vec![],
+        shm_size: "10g".into(),
     }
 }
 
@@ -169,6 +170,16 @@ fn ray_backend() {
     assert!(cmd.contains("--distributed-executor-backend ray"));
 }
 
+#[test]
+fn custom_tensor_parallel_size() {
+    // Per-node TP=1 is the default. Custom TP (e.g. multi-node TP)
+    // can be set via recipe.tensor_parallel_size.
+    let mut r = trec();
+    r.tensor_parallel_size = 2;
+    let cmd = run(&r, &tcls().nodes[0], true);
+    assert!(cmd.contains("--tensor-parallel-size 2"), "custom TP");
+}
+
 #[test]
 fn auto_pp_size() {
     let cmd = run(&trec(), &tcls().nodes[0], true);
@@ -207,6 +218,7 @@ fn sglang_recipe() -> Recipe {
         entrypoint: vec!["sglang".into(), "serve".into()],
         distributed_backend: "mp".into(),
         before_serve: vec![],
+        shm_size: "10g".into(),
     }
 }
 

+ 1 - 0
src/doctor.rs

@@ -323,6 +323,7 @@ mod tests {
             entrypoint: vec![],
             distributed_backend: "mp".to_string(),
             before_serve: vec![],
+            shm_size: "10g".to_string(),
         }
     }