diff --git a/common/types/pd_planner.go b/common/types/pd_planner.go new file mode 100644 index 00000000..7f0714b1 --- /dev/null +++ b/common/types/pd_planner.go @@ -0,0 +1,428 @@ +// pd_planner_ee.go — PD Planner (parallelism strategy planner) +// +// Responsibility of Planner: +// Given model parameter count, expert count, precision and other characteristics, +// automatically computes the optimal PD disaggregated deployment parallelism config +// (TP/EP/DP), and estimates per-GPU weight memory usage and remaining KV cache space. +// +// Core functions: +// - PlanPD: enumerates all valid (TP, EP, DP) combinations and selects the optimal config +// using different heuristics for prefill vs decode +// - PlanPDRecommendation: estimates NonMoE parameter count from config.json metadata +// (hidden_size, num_hidden_layers, etc.), calls PlanPD to generate a PDRecommendation, +// and stores it in the metadata table +// +// Difference from pd_recommendation_ee.go: +// - Planner (this file): responsible for "computation" — given model specs and GPU specs, +// outputs the optimal parallelism strategy +// - Recommendation (pd_recommendation_ee.go): responsible for "description and matching" — +// defines the data structures for recommendation results and methods to match them +// against hardware resources + +package types + +import ( + "fmt" + "math" + "strings" +) + +// ===== Precision Constants ===== + +// precisionBytes maps precision strings to bytes per parameter. +var precisionBytes = map[string]float64{ + "fp16": 2.0, + "bf16": 2.0, + "fp8": 1.0, + "int8": 1.0, + "int4": 0.5, +} + +// systemOverheadFactor accounts for CUDA context, activations, and framework workspace. +const systemOverheadFactor = 1.15 + +// defaultKVRatioPrefill is the fraction of GPU VRAM reserved for KV cache in prefill. +const defaultKVRatioPrefill = 0.15 + +// defaultKVRatioDecode is the fraction of GPU VRAM reserved for KV cache in decode. +const defaultKVRatioDecode = 0.35 + +// ===== PD Planning Algorithm ===== + +// pdCandidate represents a valid (TP, EP, DP) configuration. +type pdCandidate struct { + TP int + EP int + DP int + WeightMem float64 + TotalGPUs int +} + +// PlanPD computes the optimal prefill and decode LWS configurations for PD disaggregation. +// +// Algorithm overview: +// 1. Enumerate all (TP, EP) combinations satisfying divisibility constraints. +// 2. For each candidate, estimate per-GPU weight memory using: +// M = (NonMoE/TP + Expert/(TP*EP)) * B * alpha +// 3. Select the best prefill config: maximize TP (within single node), minimize total GPUs. +// 4. Select the best decode config: minimize TP (<=2 preferred), maximize EP, minimize total GPUs. +// +// Returns (prefillResult, decodeResult, error). Returns an error if no valid config is found. +func PlanPD(input PDPlanInput) (prefill PDPlanResult, decode PDPlanResult, err error) { + if input.KVRatioPrefill <= 0 { + input.KVRatioPrefill = defaultKVRatioPrefill + } + if input.KVRatioDecode <= 0 { + input.KVRatioDecode = defaultKVRatioDecode + } + + spec := input.Model + bParam, ok := precisionBytes[strings.ToLower(spec.Precision)] + if !ok { + bParam = 1.0 // Default to fp8 if unknown + } + + isMoE := spec.TotalExperts > 0 + expertParamsB := spec.TotalParamsB - spec.NonMoEParamsB + + // Build TP candidates: powers of 2, limited to GPUsPerNode for NVLink efficiency + tpCandidates := buildTPCandidates(input.GPU.GPUsPerNode) + + // Build EP candidates: for MoE models, factors of TotalExperts; for dense, only EP=1 + epCandidates := buildEPCandidates(spec.TotalExperts, isMoE) + + // Build DP candidates: powers of 2, default [1]. + // DP replicates the model for higher throughput; it does not reduce per-GPU memory. + dpCandidates := []int{1} + + // Enumerate all valid candidates + candidates := make([]pdCandidate, 0, len(tpCandidates)*len(epCandidates)*len(dpCandidates)) + for _, tp := range tpCandidates { + for _, ep := range epCandidates { + weightMem := calcWeightMemPerGPU(spec.NonMoEParamsB, expertParamsB, tp, ep, bParam) + for _, dp := range dpCandidates { + candidates = append(candidates, pdCandidate{ + TP: tp, + EP: ep, + DP: dp, + WeightMem: weightMem, + TotalGPUs: tp * ep * dp, + }) + } + } + } + + // Apply MaxTotalGPUs limit (default: 32 = 4 nodes of 8 GPUs) + maxGPUs := input.MaxTotalGPUs + if maxGPUs <= 0 { + maxGPUs = 32 + } + + // Filter candidates by MaxTotalGPUs + filtered := make([]pdCandidate, 0, len(candidates)) + for _, c := range candidates { + if c.TotalGPUs <= maxGPUs { + filtered = append(filtered, c) + } + } + + // Plan prefill: maximize TP, prefer single node, minimize total GPUs + prefill, err = planPrefill(filtered, input.GPU.VRAMGB, input.KVRatioPrefill, input.GPU.GPUsPerNode) + if err != nil { + return PDPlanResult{}, PDPlanResult{}, fmt.Errorf("prefill planning failed: %w", err) + } + + // Plan decode: minimize TP (<=2 preferred), maximize EP, minimize total GPUs + decode, err = planDecode(filtered, input.GPU.VRAMGB, input.KVRatioDecode, input.GPU.GPUsPerNode) + if err != nil { + return PDPlanResult{}, PDPlanResult{}, fmt.Errorf("decode planning failed: %w", err) + } + + return prefill, decode, nil +} + +// calcWeightMemPerGPU estimates the model weight memory per GPU in GB. +// Formula: (NonMoE/TP + Expert/(TP*EP)) * B * alpha +// +// Preconditions: tp >= 1 and ep >= 1 (callers use buildTPCandidates/buildEPCandidates +// which always return values >= 1). Violating this would cause division by zero. +func calcWeightMemPerGPU(nonMoEParamsB, expertParamsB float64, tp, ep int, bParam float64) float64 { + return (nonMoEParamsB/float64(tp) + expertParamsB/(float64(tp)*float64(ep))) * bParam * systemOverheadFactor +} + +// buildTPCandidates returns TP values as powers of 2, not exceeding GPUsPerNode. +func buildTPCandidates(gpusPerNode int) []int { + powers := []int{1, 2, 4, 8, 16, 32, 64} + result := make([]int, 0, len(powers)) + for _, tp := range powers { + if tp <= gpusPerNode { + result = append(result, tp) + } + } + return result +} + +// buildEPCandidates returns EP values for expert parallelism. +// For dense models (totalExperts==0), returns [1]. +// For MoE models, returns all factors of totalExperts. +func buildEPCandidates(totalExperts int, isMoE bool) []int { + if !isMoE || totalExperts <= 0 { + return []int{1} + } + factors := make([]int, 0) + for i := 1; i <= totalExperts; i++ { + if totalExperts%i == 0 { + factors = append(factors, i) + } + } + return factors +} + +// planPrefill selects the best prefill configuration. +// Heuristic: maximize TP (for compute throughput), prefer single node (NVLink), +// but allow multi-node with EP if the model is too large for a single node. +func planPrefill(candidates []pdCandidate, gpuVRAMGB, kvRatio float64, gpusPerNode int) (PDPlanResult, error) { + maxWeight := gpuVRAMGB * (1 - kvRatio) + + // Primary: single-node candidates (TP*EP <= gpusPerNode) + var singleNode []pdCandidate + for _, c := range candidates { + if c.WeightMem <= maxWeight && c.TotalGPUs <= gpusPerNode { + singleNode = append(singleNode, c) + } + } + + if len(singleNode) > 0 { + // Sort: minimize total GPUs first, then maximize TP, then minimize EP + sortCandidates(singleNode, func(a, b pdCandidate) bool { + if a.TotalGPUs != b.TotalGPUs { + return a.TotalGPUs < b.TotalGPUs + } + if a.TP != b.TP { + return a.TP > b.TP + } + return a.EP < b.EP + }) + return buildPlanResult(singleNode[0], gpuVRAMGB, gpusPerNode), nil + } + + // Fallback: allow multi-node if single-node is not feasible + var multiNode []pdCandidate + for _, c := range candidates { + if c.WeightMem <= maxWeight { + multiNode = append(multiNode, c) + } + } + + if len(multiNode) == 0 { + return PDPlanResult{}, fmt.Errorf("no valid prefill config: model cannot fit with %.0fGB VRAM", gpuVRAMGB) + } + + // Sort: minimize total GPUs, then maximize TP + sortCandidates(multiNode, func(a, b pdCandidate) bool { + if a.TotalGPUs != b.TotalGPUs { + return a.TotalGPUs < b.TotalGPUs + } + return a.TP > b.TP + }) + + return buildPlanResult(multiNode[0], gpuVRAMGB, gpusPerNode), nil +} + +// planDecode selects the best decode configuration. +// Heuristic: minimize TP (TP<=2 preferred for small-batch efficiency), +// maximize EP (spread expert weights across more GPUs for KV cache headroom), +// minimize total GPUs. +func planDecode(candidates []pdCandidate, gpuVRAMGB, kvRatio float64, gpusPerNode int) (PDPlanResult, error) { + maxWeight := gpuVRAMGB * (1 - kvRatio) + + // Primary search: TP <= 2 + var primary []pdCandidate + for _, c := range candidates { + if c.WeightMem <= maxWeight && c.TP <= 2 { + primary = append(primary, c) + } + } + + var best pdCandidate + found := false + + if len(primary) > 0 { + // Sort: minimize TP, then minimize total GPUs, then maximize EP + sortCandidates(primary, func(a, b pdCandidate) bool { + if a.TP != b.TP { + return a.TP < b.TP + } + if a.TotalGPUs != b.TotalGPUs { + return a.TotalGPUs < b.TotalGPUs + } + return a.EP > b.EP + }) + best = primary[0] + found = true + } + + // Fallback: relax TP constraint, search all valid configs + if !found { + var backup []pdCandidate + for _, c := range candidates { + if c.WeightMem <= maxWeight { + backup = append(backup, c) + } + } + if len(backup) > 0 { + sortCandidates(backup, func(a, b pdCandidate) bool { + if a.TP != b.TP { + return a.TP < b.TP + } + return a.TotalGPUs < b.TotalGPUs + }) + best = backup[0] + found = true + } + } + + if !found { + return PDPlanResult{}, fmt.Errorf("no valid decode config: model cannot fit with %.0fGB VRAM", gpuVRAMGB) + } + + return buildPlanResult(best, gpuVRAMGB, gpusPerNode), nil +} + +// buildPlanResult converts a candidate into a PDPlanResult with LWS layout. +func buildPlanResult(c pdCandidate, gpuVRAMGB float64, gpusPerNode int) PDPlanResult { + var lwsSize int + var gpusPerPod int + totalGPUs := c.TotalGPUs + + if totalGPUs >= gpusPerNode { + lwsSize = int(math.Ceil(float64(totalGPUs) / float64(gpusPerNode))) + gpusPerPod = gpusPerNode + } else { + lwsSize = 1 + gpusPerPod = totalGPUs + } + + remaining := gpuVRAMGB - c.WeightMem + + return PDPlanResult{ + TP: c.TP, + EP: c.EP, + DP: c.DP, + TotalGPUs: totalGPUs, + LWSSize: lwsSize, + GPUsPerPod: gpusPerPod, + WeightMemPerGPU: c.WeightMem, + RemainingVRAMForKV: remaining, + } +} + +// sortCandidates sorts a slice of pdCandidate using a custom less function (insertion sort for small slices). +func sortCandidates(candidates []pdCandidate, less func(a, b pdCandidate) bool) { + // Use insertion sort since candidate lists are typically small + for i := 1; i < len(candidates); i++ { + for j := i; j > 0 && less(candidates[j], candidates[j-1]); j-- { + candidates[j], candidates[j-1] = candidates[j-1], candidates[j] + } + } +} + +// ===== PDRecommendation Generation from ModelInfo ===== + +// DefaultGPUUnitGB is the standard GPU VRAM unit (80GB) used for PD recommendations. +const DefaultGPUUnitGB = 80.0 + +// DefaultGPUsPerNode is the standard number of GPUs per node. +const DefaultGPUsPerNode = 8 + +// PlanPDRecommendation generates a PDRecommendation from model metadata parsed +// from config.json. It uses the 80GB GPU unit as the standard planning baseline. +// +// Parameters: +// - modelName: repository path (used for logging only, no longer looked up in a registry) +// - totalParamsB: total parameter count in billions (from config.json / safetensors) +// - totalExperts: number of routed experts (0 for dense models) +// - activeExperts: number of experts activated per token +// - precision: inference precision (fp8, bf16, etc.) +// - hiddenSize: hidden size from config.json (used to estimate NonMoEParamsB) +// - numHiddenLayers: number of hidden layers from config.json +// +// The function: +// 1. Estimates NonMoEParamsB from hidden_size and num_hidden_layers when available, +// otherwise falls back to the active/total expert ratio heuristic. +// 2. Plans prefill and decode configurations using PlanPD. +// 3. Returns a PDRecommendation suitable for storage as JSONB. +func PlanPDRecommendation(modelName string, totalParamsB float64, totalExperts, activeExperts int, precision string, hiddenSize, numHiddenLayers int) (*PDRecommendation, error) { + if totalParamsB <= 0 { + return nil, fmt.Errorf("total params must be positive, got %f", totalParamsB) + } + + spec := PDModelSpec{ + TotalParamsB: totalParamsB, + TotalExperts: totalExperts, + Precision: precision, + } + + // Estimate NonMoEParamsB (non-expert parameters: attention + shared FFN + embedding). + // When hidden_size and num_hidden_layers are available from config.json, + // use the standard transformer parameter formula: + // NonMoE ≈ num_layers * 12 * hidden_size^2 (Q/K/V/O + shared FFN + layer norms) + // This is more accurate than the active/total expert ratio heuristic. + if hiddenSize > 0 && numHiddenLayers > 0 { + spec.NonMoEParamsB = float64(numHiddenLayers) * 12.0 * float64(hiddenSize) * float64(hiddenSize) / 1e9 + // Cap at total params to avoid overestimation for small models + if spec.NonMoEParamsB > spec.TotalParamsB { + spec.NonMoEParamsB = spec.TotalParamsB + } + } else if spec.TotalExperts > 0 && activeExperts > 0 { + // Fallback: estimate non-MoE params from active/total expert ratio + spec.NonMoEParamsB = spec.TotalParamsB * float64(activeExperts) / float64(spec.TotalExperts) + } else { + // Dense model: all params are non-MoE + spec.NonMoEParamsB = spec.TotalParamsB + } + + // Plan with 80GB GPU unit + gpuConfig := PDGPUConfig{ + VRAMGB: DefaultGPUUnitGB, + GPUsPerNode: DefaultGPUsPerNode, + } + + prefill, decode, err := PlanPD(PDPlanInput{ + Model: spec, + GPU: gpuConfig, + }) + if err != nil { + return nil, fmt.Errorf("failed to plan PD for model %s: %w", modelName, err) + } + + return &PDRecommendation{ + ModelName: modelName, + TotalParamsB: spec.TotalParamsB, + NonMoEParamsB: spec.NonMoEParamsB, + TotalExperts: spec.TotalExperts, + ActiveExperts: activeExperts, + Precision: spec.Precision, + MinInferenceVRAMGB: DefaultGPUUnitGB, + Prefill: planResultToRoleConfig(prefill), + Decode: planResultToRoleConfig(decode), + }, nil +} + +// planResultToRoleConfig converts a PDPlanResult into a PDRoleConfig. +// TotalVRAMGB is computed as MinInferenceVRAMGB * TotalGPUs by the caller +// (PlanPDRecommendation sets MinInferenceVRAMGB = DefaultGPUUnitGB). +func planResultToRoleConfig(result PDPlanResult) PDRoleConfig { + pods := result.LWSSize + if pods < 1 { + pods = 1 + } + return PDRoleConfig{ + TP: result.TP, + EP: result.EP, + DP: result.DP, + TotalGPUs: result.TotalGPUs, + Pods: pods, + TotalVRAMGB: float64(result.TotalGPUs) * DefaultGPUUnitGB, + } +} diff --git a/common/types/pd_planner_test.go b/common/types/pd_planner_test.go new file mode 100644 index 00000000..81af3f72 --- /dev/null +++ b/common/types/pd_planner_test.go @@ -0,0 +1,354 @@ +package types + +import ( + "math" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestPlanPD_DeepSeekV3_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 671, NonMoEParamsB: 37, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.GreaterOrEqual(t, prefill.TP, 4) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.LessOrEqual(t, decode.TP, 4) + + t.Logf("DeepSeek-V3 H200: Prefill TP=%d EP=%d GPUs=%d LWSSize=%d weight=%.2fGB kv=%.2fGB", + prefill.TP, prefill.EP, prefill.TotalGPUs, prefill.LWSSize, prefill.WeightMemPerGPU, prefill.RemainingVRAMForKV) + t.Logf("DeepSeek-V3 H200: Decode TP=%d EP=%d GPUs=%d LWSSize=%d weight=%.2fGB kv=%.2fGB", + decode.TP, decode.EP, decode.TotalGPUs, decode.LWSSize, decode.WeightMemPerGPU, decode.RemainingVRAMForKV) +} + +func TestPlanPD_DeepSeekV3_A800(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 671, NonMoEParamsB: 37, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("DeepSeek-V3 A800: Prefill TP=%d EP=%d GPUs=%d weight=%.2fGB kv=%.2fGB", + prefill.TP, prefill.EP, prefill.TotalGPUs, prefill.WeightMemPerGPU, prefill.RemainingVRAMForKV) + t.Logf("DeepSeek-V3 A800: Decode TP=%d EP=%d GPUs=%d weight=%.2fGB kv=%.2fGB", + decode.TP, decode.EP, decode.TotalGPUs, decode.WeightMemPerGPU, decode.RemainingVRAMForKV) +} + +func TestPlanPD_DeepSeekV4Flash_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 284, NonMoEParamsB: 13, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("DeepSeek-V4-Flash H200: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("DeepSeek-V4-Flash H200: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_DeepSeekV4Pro_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 1600, NonMoEParamsB: 49, TotalExperts: 512, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("DeepSeek-V4-Pro H200: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("DeepSeek-V4-Pro H200: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_GLM51_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 744, NonMoEParamsB: 40, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("GLM-5.1 H200: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("GLM-5.1 H200: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_GLM52_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 1000, NonMoEParamsB: 50, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("GLM-5.2 H200: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("GLM-5.2 H200: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_GPT120B_A800(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 120, NonMoEParamsB: 120, TotalExperts: 0, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Equal(t, 1, prefill.EP) + require.Equal(t, 1, decode.EP) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("GPT-120B A800: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("GPT-120B A800: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_GPT120B_H200(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 120, NonMoEParamsB: 120, TotalExperts: 0, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Equal(t, 1, prefill.EP) + require.Equal(t, 1, decode.EP) + + t.Logf("GPT-120B H200: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("GPT-120B H200: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_Qwen3_30B_A3B_A800(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 30, NonMoEParamsB: 3, TotalExperts: 128, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("Qwen3-30B-A3B A800: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("Qwen3-30B-A3B A800: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_Qwen1_5MoE_A800(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 14.3, NonMoEParamsB: 2.5, TotalExperts: 60, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) + + t.Logf("Qwen1.5-MoE-A2.7B A800: Prefill TP=%d EP=%d GPUs=%d", prefill.TP, prefill.EP, prefill.TotalGPUs) + t.Logf("Qwen1.5-MoE-A2.7B A800: Decode TP=%d EP=%d GPUs=%d", decode.TP, decode.EP, decode.TotalGPUs) +} + +func TestPlanPD_DenseModelEPAlwaysOne(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 72, NonMoEParamsB: 72, TotalExperts: 0, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Equal(t, 1, prefill.EP) + require.Equal(t, 1, decode.EP) +} + +func TestPlanPD_WeightMemCalculation(t *testing.T) { + weight := calcWeightMemPerGPU(72, 0, 2, 1, 2.0) + require.InDelta(t, 82.8, weight, 0.01) +} + +func TestPlanPD_SmallModelSingleNode(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 7, NonMoEParamsB: 7, TotalExperts: 0, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + }) + require.NoError(t, err) + require.Equal(t, 1, prefill.TotalGPUs) + require.Equal(t, 1, prefill.TP) + require.Equal(t, 1, prefill.LWSSize) + require.Equal(t, 1, decode.TotalGPUs) + require.Equal(t, 1, decode.TP) +} + +func TestPlanPD_DefaultKVRatios(t *testing.T) { + prefill, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 72, NonMoEParamsB: 72, TotalExperts: 0, Precision: "bf16"}, + GPU: PDGPUConfig{VRAMGB: 80, GPUsPerNode: 8}, + KVRatioPrefill: 0, + KVRatioDecode: 0, + }) + require.NoError(t, err) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0) + require.Greater(t, decode.RemainingVRAMForKV, 0.0) +} + +func TestPlanPD_LWSSizeCalculation(t *testing.T) { + _, decode, err := PlanPD(PDPlanInput{ + Model: PDModelSpec{TotalParamsB: 671, NonMoEParamsB: 37, TotalExperts: 256, Precision: "fp8"}, + GPU: PDGPUConfig{VRAMGB: 141, GPUsPerNode: 8}, + }) + require.NoError(t, err) + if decode.TotalGPUs > 8 { + expectedSize := int(math.Ceil(float64(decode.TotalGPUs) / 8.0)) + require.Equal(t, expectedSize, decode.LWSSize) + require.Equal(t, 8, decode.GPUsPerPod) + } +} + +func TestBuildEPCandidates_Dense(t *testing.T) { + candidates := buildEPCandidates(0, false) + require.Equal(t, []int{1}, candidates) +} + +func TestBuildEPCandidates_MoE(t *testing.T) { + candidates := buildEPCandidates(256, true) + require.Contains(t, candidates, 1) + require.Contains(t, candidates, 256) + require.Contains(t, candidates, 8) + for _, ep := range candidates { + require.Equal(t, 0, 256%ep, "EP %d should be a factor of 256", ep) + } +} + +func TestPlanPD_FullMatrix(t *testing.T) { + models := map[string]PDModelSpec{ + "DeepSeek-V3": {TotalParamsB: 671, NonMoEParamsB: 37, TotalExperts: 256, Precision: "fp8"}, + "DeepSeek-V4-Flash": {TotalParamsB: 284, NonMoEParamsB: 13, TotalExperts: 256, Precision: "fp8"}, + "DeepSeek-V4-Pro": {TotalParamsB: 1600, NonMoEParamsB: 49, TotalExperts: 512, Precision: "fp8"}, + "GLM-5.1": {TotalParamsB: 744, NonMoEParamsB: 40, TotalExperts: 256, Precision: "fp8"}, + "GLM-5.2": {TotalParamsB: 1000, NonMoEParamsB: 50, TotalExperts: 256, Precision: "fp8"}, + "GPT-120B": {TotalParamsB: 120, NonMoEParamsB: 120, TotalExperts: 0, Precision: "bf16"}, + "Qwen3-30B-A3B": {TotalParamsB: 30, NonMoEParamsB: 3, TotalExperts: 128, Precision: "bf16"}, + "Qwen1.5-MoE": {TotalParamsB: 14.3, NonMoEParamsB: 2.5, TotalExperts: 60, Precision: "bf16"}, + } + + // infeasible combinations: model too large for GPU with 32-GPU limit + infeasible := map[string]bool{ + "DeepSeek-V4-Pro_A800": true, // 1600B model needs >32 A800 GPUs + } + + gpus := map[string]PDGPUConfig{ + "A800": {VRAMGB: 80, GPUsPerNode: 8}, + "H200": {VRAMGB: 141, GPUsPerNode: 8}, + } + + for modelName, spec := range models { + for gpuName, gpu := range gpus { + key := modelName + "_" + gpuName + t.Run(key, func(t *testing.T) { + if infeasible[key] { + t.Skipf("skipping infeasible combination: %s on %s", modelName, gpuName) + } + prefill, decode, err := PlanPD(PDPlanInput{Model: spec, GPU: gpu}) + require.NoError(t, err, "PlanPD should succeed for %s on %s", modelName, gpuName) + require.Greater(t, prefill.RemainingVRAMForKV, 0.0, "prefill should have KV headroom") + require.Greater(t, decode.RemainingVRAMForKV, 0.0, "decode should have KV headroom") + + t.Logf("%-20s %-5s | Prefill: TP=%2d EP=%3d GPUs=%2d Size=%d weight=%.1fGB | Decode: TP=%2d EP=%3d GPUs=%2d Size=%d weight=%.1fGB", + modelName, gpuName, + prefill.TP, prefill.EP, prefill.TotalGPUs, prefill.LWSSize, prefill.WeightMemPerGPU, + decode.TP, decode.EP, decode.TotalGPUs, decode.LWSSize, decode.WeightMemPerGPU) + }) + } + } +} + +func TestPlanPDRecommendation_DeepSeekV3(t *testing.T) { + rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61) + require.NoError(t, err) + require.NotNil(t, rec) + require.Equal(t, 671.0, rec.TotalParamsB) + require.Equal(t, 256, rec.TotalExperts) + require.Equal(t, 8, rec.ActiveExperts) + require.Equal(t, "fp8", rec.Precision) + require.Equal(t, 80.0, rec.MinInferenceVRAMGB) + require.Greater(t, rec.Prefill.TotalGPUs, 0) + require.Greater(t, rec.Decode.TotalGPUs, 0) + require.Greater(t, rec.Prefill.TotalVRAMGB, 0.0) + require.Greater(t, rec.Prefill.Pods, 0) + require.Greater(t, rec.Decode.Pods, 0) + + t.Logf("DeepSeek-V3 80G: Prefill TP=%d EP=%d GPUs=%d Pods=%d | Decode TP=%d EP=%d GPUs=%d Pods=%d", + rec.Prefill.TP, rec.Prefill.EP, rec.Prefill.TotalGPUs, rec.Prefill.Pods, + rec.Decode.TP, rec.Decode.EP, rec.Decode.TotalGPUs, rec.Decode.Pods) +} + +func TestPlanPDRecommendation_GLM52(t *testing.T) { + rec, err := PlanPDRecommendation("glm-5.2", 1000, 256, 8, "fp8", 8192, 80) + require.NoError(t, err) + require.NotNil(t, rec) + require.Equal(t, 1000.0, rec.TotalParamsB) + require.Greater(t, rec.Prefill.TotalGPUs, 0) + require.Greater(t, rec.Decode.TotalGPUs, 0) +} + +func TestPlanPDRecommendation_Qwen3_30B_A3B(t *testing.T) { + rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48) + require.NoError(t, err) + require.NotNil(t, rec) + require.Equal(t, 30.0, rec.TotalParamsB) + require.Equal(t, 128, rec.TotalExperts) +} + +func TestPlanPDRecommendation_UnknownModelMoE(t *testing.T) { + // Unknown MoE model with expert info from config.json + rec, err := PlanPDRecommendation("custom/moe-model", 200, 64, 8, "bf16", 4096, 32) + require.NoError(t, err) + require.NotNil(t, rec) + require.Equal(t, 200.0, rec.TotalParamsB) + require.Equal(t, 64, rec.TotalExperts) + require.Equal(t, 8, rec.ActiveExperts) + // NonMoEParamsB should be estimated from hidden_size and num_hidden_layers + require.Greater(t, rec.NonMoEParamsB, 0.0) + require.Less(t, rec.NonMoEParamsB, rec.TotalParamsB) +} + +func TestPlanPDRecommendation_DenseModel(t *testing.T) { + rec, err := PlanPDRecommendation("custom-dense-70b", 70, 0, 0, "bf16", 0, 0) + require.NoError(t, err) + require.NotNil(t, rec) + require.Equal(t, 70.0, rec.TotalParamsB) + require.Equal(t, 0, rec.TotalExperts) + require.Equal(t, 70.0, rec.NonMoEParamsB) +} + +func TestPlanPDRecommendation_ZeroParams(t *testing.T) { + _, err := PlanPDRecommendation("test", 0, 256, 8, "fp8", 0, 0) + require.Error(t, err) + require.Contains(t, err.Error(), "total params must be positive") +} + +func TestPlanPDRecommendation_NegativeParams(t *testing.T) { + _, err := PlanPDRecommendation("test", -10, 256, 8, "fp8", 0, 0) + require.Error(t, err) + require.Contains(t, err.Error(), "total params must be positive") +} + +func TestPlanPDRecommendation_PrecisionOverride(t *testing.T) { + // Known model with explicit precision override + // Use a smaller model that fits in 80GB with bf16 + rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48) + require.NoError(t, err) + require.Equal(t, "bf16", rec.Precision) +} + +func TestPlanPDRecommendation_ExpertOverride(t *testing.T) { + // Override expert count from config.json + rec, err := PlanPDRecommendation("deepseek-v3", 671, 512, 16, "fp8", 7168, 61) + require.NoError(t, err) + require.Equal(t, 512, rec.TotalExperts) + require.Equal(t, 16, rec.ActiveExperts) +} + +func TestPlanPDRecommendation_DPField(t *testing.T) { + rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61) + require.NoError(t, err) + // DP should default to 1 for all recommendations + require.Equal(t, 1, rec.Prefill.DP) + require.Equal(t, 1, rec.Decode.DP) + // TotalGPUs should equal TP * EP * DP + require.Equal(t, rec.Prefill.TP*rec.Prefill.EP*rec.Prefill.DP, rec.Prefill.TotalGPUs) + require.Equal(t, rec.Decode.TP*rec.Decode.EP*rec.Decode.DP, rec.Decode.TotalGPUs) +} diff --git a/docker/inference/sglang/serve.sh b/docker/inference/sglang/serve.sh index 654f1ce5..e88e5435 100644 --- a/docker/inference/sglang/serve.sh +++ b/docker/inference/sglang/serve.sh @@ -1,7 +1,7 @@ #!/bin/bash if [ -n "$PD_ROLE" ]; then - # PD disaggregation mode: delegate to pd-disaggregation.sh - bash /etc/csghub/pd-disaggregation.sh + # PD disaggregation mode: delegate to pd-disaggregation_ee.sh + bash /etc/csghub/pd-disaggregation_ee.sh elif [ -z "$LWS_WORKER_INDEX" ]; then bash /etc/csghub/single-node.sh else diff --git a/docker/inference/vllm/serve.sh b/docker/inference/vllm/serve.sh index 654f1ce5..e88e5435 100644 --- a/docker/inference/vllm/serve.sh +++ b/docker/inference/vllm/serve.sh @@ -1,7 +1,7 @@ #!/bin/bash if [ -n "$PD_ROLE" ]; then - # PD disaggregation mode: delegate to pd-disaggregation.sh - bash /etc/csghub/pd-disaggregation.sh + # PD disaggregation mode: delegate to pd-disaggregation_ee.sh + bash /etc/csghub/pd-disaggregation_ee.sh elif [ -z "$LWS_WORKER_INDEX" ]; then bash /etc/csghub/single-node.sh else