From 59c39a2c73869c04d5a580e5363effe938b90197 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 24 Jun 2026 21:04:49 +0000 Subject: [PATCH] fix(config): per-device VRAM headroom for Blackwell defaults (#10485) The hardware-tuned defaults from #10411 were measured on a GB10 / DGX Spark (128 GiB unified memory) and over-provisioned multi-GPU consumer Blackwell (e.g. 2x16 GiB RTX 50-series) into CUDA OOM during model init: - The Blackwell physical batch (512 -> 2048) sets both n_batch and n_ubatch. The compute buffer scales ~n_ubatch * n_ctx and is allocated PER DEVICE (it can't be split across GPUs), so a large context turns ub2048 into multi-GiB of scratch that must fit one 16 GiB card. - The VRAM-scaled parallel-slot default tiered off TotalAvailableVRAM(), which SUMS all GPUs (2x16 -> "32 GiB" -> 8 slots), but the allocations are per-device. Make both decisions per-device and context-aware: - xsysinfo.MinPerGPUVRAM() reports the smallest device's VRAM; localGPU() uses it so the parallel tier and batch guard reason about one card. - PhysicalBatchForContext(gpu, ctx) raises the batch only when the extra compute buffer fits VRAM/4 at this model's context (16 GiB crosses over ~174k ctx, 32 GiB ~349k; GB10 reports system RAM so it still clears it). - Apply hardware defaults AFTER runBackendHooks in SetDefaults so the GGUF-guessed context is resolved before the batch decision. - The distributed router gates the node batch the same way. Unified-memory devices (GB10, Apple) report system RAM as their single device's VRAM, so they keep the prefill win. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude:opus-4.8 [Claude Code] --- core/config/hardware_defaults.go | 80 +++++++++++++++++-- .../config/hardware_defaults_internal_test.go | 19 ++++- core/config/hardware_defaults_test.go | 45 +++++++++-- core/config/model_config.go | 15 ++-- core/services/nodes/router.go | 5 +- .../nodes/router_hardware_internal_test.go | 13 ++- pkg/xsysinfo/gpu.go | 55 +++++++++++++ pkg/xsysinfo/minvram_internal_test.go | 37 +++++++++ 8 files changed, 244 insertions(+), 25 deletions(-) create mode 100644 pkg/xsysinfo/minvram_internal_test.go diff --git a/core/config/hardware_defaults.go b/core/config/hardware_defaults.go index 18c321639a58..b4e0e74c6f70 100644 --- a/core/config/hardware_defaults.go +++ b/core/config/hardware_defaults.go @@ -54,8 +54,35 @@ func (g GPU) IsNVIDIABlackwell() bool { return maj >= 12 } +// Compute-buffer headroom guard for the raised physical batch. +// +// Raising n_ubatch grows the CUDA *compute buffer* (the scratch for the forward +// graph), which is allocated PER DEVICE — it does not benefit from a second GPU +// the way weights or KV (which are split across devices) do. The buffer scales +// ~linearly with n_ubatch * n_ctx, so a large context turns the GB10-tuned +// ub2048 into multi-GiB of extra scratch that must fit on a SINGLE card. On a +// 16 GiB consumer Blackwell with a 200k context that overflows (issue #10485), +// even though the GB10 it was measured on (128 GiB unified memory) had room. +// +// These constants size a conservative guard: only raise the batch when the +// extra scratch fits the per-device VRAM ceiling. +const ( + // computeBufferBytesPerCell approximates the CUDA compute-buffer cost of one + // (n_ubatch * n_ctx) cell. Derived from an observed allocation (ub2048 * + // ctx204800 ~= 4.5 GiB => ~11 B/cell) and rounded up to 16 for margin, since + // the real cost also grows with model width (heads / embedding dim) which we + // don't know at config time. + computeBufferBytesPerCell = 16 + // blackwellBatchHeadroomDivisor caps the extra compute buffer from raising the + // physical batch at VRAM/divisor. /4 keeps the bulk of a device for weights + + // KV, which already dominate VRAM use. + blackwellBatchHeadroomDivisor = 4 +) + // PhysicalBatch returns the canonical physical batch (n_batch/n_ubatch) for the -// given hardware, used when the model config leaves batch unset. +// given hardware class, ignoring context/VRAM headroom. Use +// PhysicalBatchForContext when a model context and per-device VRAM are known +// (the load paths) so the raised batch can't overflow a single device. func PhysicalBatch(g GPU) int { if g.IsNVIDIABlackwell() { return BlackwellPhysicalBatch @@ -63,6 +90,32 @@ func PhysicalBatch(g GPU) int { return DefaultPhysicalBatch } +// PhysicalBatchForContext is PhysicalBatch gated on per-device VRAM headroom for +// the given context: it only raises the batch above the conservative default +// when the extra compute buffer (which is allocated on a single device and grows +// with n_ubatch * n_ctx) fits within blackwellBatchHeadroomDivisor of the GPU's +// VRAM. g.VRAM must be the PER-DEVICE ceiling (the smallest device on a +// multi-GPU host), not the summed total — the compute buffer can't be split. +// +// VRAM 0 (unknown) stays conservative rather than risk a per-device OOM; the +// GB10 / unified-memory path reports system RAM, so it still clears the guard. +func PhysicalBatchForContext(g GPU, ctx int) int { + if !g.IsNVIDIABlackwell() { + return DefaultPhysicalBatch + } + if ctx <= 0 { + ctx = DefaultContextSize + } + if g.VRAM == 0 { + return DefaultPhysicalBatch + } + extra := uint64(ctx) * uint64(BlackwellPhysicalBatch-DefaultPhysicalBatch) * computeBufferBytesPerCell + if extra <= g.VRAM/blackwellBatchHeadroomDivisor { + return BlackwellPhysicalBatch + } + return DefaultPhysicalBatch +} + // IsManagedPhysicalBatch reports whether n is a value PhysicalBatch assigns. // Callers that re-tune a value chosen by an upstream host (the distributed // router correcting the frontend's guess) use this to avoid clobbering an @@ -122,7 +175,12 @@ func hasParallelOption(opts []string) bool { // deterministic device — detection does a live nvidia-smi call. var localGPU = func() GPU { vendor, _ := xsysinfo.DetectGPUVendor() - vram, _ := xsysinfo.TotalAvailableVRAM() + // Use the SMALLEST device's VRAM, not the summed total: the parallel-slot + // tier and the batch headroom guard both reason about what fits on a single + // card, and per-device compute buffers can't be split across GPUs. Summing + // two 16 GiB cards into "32 GiB" is what over-provisioned multi-GPU hosts + // into OOM (issue #10485). + vram, _ := xsysinfo.MinPerGPUVRAM() return GPU{ Vendor: vendor, ComputeCapability: xsysinfo.NVIDIAComputeCapability(), @@ -137,10 +195,20 @@ func ApplyHardwareDefaults(cfg *ModelConfig, gpu GPU) { if cfg == nil { return } - if cfg.Batch == 0 && gpu.IsNVIDIABlackwell() { - cfg.Batch = BlackwellPhysicalBatch - xlog.Debug("[hardware_defaults] Blackwell GPU: defaulting physical batch", - "batch", cfg.Batch, "compute_cap", gpu.ComputeCapability) + // Raise the physical batch on Blackwell only when the resulting compute + // buffer fits the per-device VRAM at THIS model's context. Leaving Batch at 0 + // (rather than writing the default 512) preserves the downstream single-pass + // sizing in core/backend.EffectiveBatchSize for embedding/score/rerank. + if cfg.Batch == 0 { + ctx := DefaultContextSize + if cfg.ContextSize != nil { + ctx = *cfg.ContextSize + } + if PhysicalBatchForContext(gpu, ctx) == BlackwellPhysicalBatch { + cfg.Batch = BlackwellPhysicalBatch + xlog.Debug("[hardware_defaults] Blackwell GPU: defaulting physical batch", + "batch", cfg.Batch, "compute_cap", gpu.ComputeCapability, "context", ctx, "vram_gib", gpu.VRAM>>30) + } } // Enable concurrent serving by default on a capable GPU: without this the diff --git a/core/config/hardware_defaults_internal_test.go b/core/config/hardware_defaults_internal_test.go index 52c674c2d585..d6878c86e136 100644 --- a/core/config/hardware_defaults_internal_test.go +++ b/core/config/hardware_defaults_internal_test.go @@ -9,26 +9,37 @@ import ( // GPU. The detection seam (localGPU) is injected so the path is deterministic // without a real GPU. var _ = Describe("SetDefaults hardware defaults (single-instance)", func() { + const gib = uint64(1) << 30 + var orig func() GPU BeforeEach(func() { orig = localGPU }) AfterEach(func() { localGPU = orig }) - It("sets the physical batch on a local Blackwell GPU", func() { - localGPU = func() GPU { return GPU{ComputeCapability: "12.1"} } + It("sets the physical batch on a local Blackwell GPU with headroom", func() { + localGPU = func() GPU { return GPU{ComputeCapability: "12.1", VRAM: 119 * gib} } cfg := &ModelConfig{} cfg.SetDefaults() Expect(cfg.Batch).To(Equal(BlackwellPhysicalBatch)) }) + It("leaves batch unset when a large context would overflow the device", func() { + // Regression guard for issue #10485: 16 GiB consumer Blackwell + ~200k ctx. + localGPU = func() GPU { return GPU{ComputeCapability: "12.0", VRAM: 16 * gib} } + ctx := 204800 + cfg := &ModelConfig{LLMConfig: LLMConfig{ContextSize: &ctx}} + cfg.SetDefaults() + Expect(cfg.Batch).To(Equal(0)) + }) + It("leaves batch unset on a non-Blackwell local GPU", func() { - localGPU = func() GPU { return GPU{ComputeCapability: "8.9"} } + localGPU = func() GPU { return GPU{ComputeCapability: "8.9", VRAM: 119 * gib} } cfg := &ModelConfig{} cfg.SetDefaults() Expect(cfg.Batch).To(Equal(0)) }) It("never overrides an explicit batch", func() { - localGPU = func() GPU { return GPU{ComputeCapability: "12.1"} } + localGPU = func() GPU { return GPU{ComputeCapability: "12.1", VRAM: 119 * gib} } cfg := &ModelConfig{} cfg.Batch = 1024 cfg.SetDefaults() diff --git a/core/config/hardware_defaults_test.go b/core/config/hardware_defaults_test.go index ae7bf39647bf..3bc1bf29759b 100644 --- a/core/config/hardware_defaults_test.go +++ b/core/config/hardware_defaults_test.go @@ -7,6 +7,8 @@ import ( ) var _ = Describe("Hardware-driven config defaults", func() { + const gib = uint64(1) << 30 + DescribeTable("GPU.IsNVIDIABlackwell (sm_12x consumer family)", func(cc string, want bool) { Expect(GPU{ComputeCapability: cc}.IsNVIDIABlackwell()).To(Equal(want)) @@ -35,21 +37,54 @@ var _ = Describe("Hardware-driven config defaults", func() { }) }) + Describe("PhysicalBatchForContext (per-device VRAM headroom)", func() { + It("raises the batch when the compute buffer fits the device", func() { + // 16 GiB Blackwell with a small context: the extra scratch is tiny. + Expect(PhysicalBatchForContext(GPU{ComputeCapability: "12.0", VRAM: 16 * gib}, 8192)). + To(Equal(BlackwellPhysicalBatch)) + }) + It("keeps the default batch when a large context would overflow one device", func() { + // The issue #10485 case: 16 GiB consumer Blackwell, ~200k context. + Expect(PhysicalBatchForContext(GPU{ComputeCapability: "12.0", VRAM: 16 * gib}, 204800)). + To(Equal(DefaultPhysicalBatch)) + }) + It("still raises the batch on a large unified-memory device (GB10)", func() { + // GB10 reports system RAM (~119 GiB) as its single device's VRAM. + Expect(PhysicalBatchForContext(GPU{ComputeCapability: "12.1", VRAM: 119 * gib}, 204800)). + To(Equal(BlackwellPhysicalBatch)) + }) + It("stays conservative when VRAM is unknown", func() { + Expect(PhysicalBatchForContext(GPU{ComputeCapability: "12.1"}, 8192)). + To(Equal(DefaultPhysicalBatch)) + }) + It("never raises the batch on non-Blackwell", func() { + Expect(PhysicalBatchForContext(GPU{ComputeCapability: "9.0", VRAM: 80 * gib}, 8192)). + To(Equal(DefaultPhysicalBatch)) + }) + }) + Describe("ApplyHardwareDefaults", func() { - It("raises an unset batch to 2048 on Blackwell", func() { + It("raises an unset batch to 2048 on Blackwell with headroom", func() { cfg := &ModelConfig{} - ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "12.1"}) + ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "12.1", VRAM: 119 * gib}) Expect(cfg.Batch).To(Equal(BlackwellPhysicalBatch)) }) + It("leaves batch unset when a large context would overflow one device", func() { + // Regression guard for issue #10485: 16 GiB card + ~200k context. + ctx := 204800 + cfg := &ModelConfig{LLMConfig: LLMConfig{ContextSize: &ctx}} + ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "12.0", VRAM: 16 * gib}) + Expect(cfg.Batch).To(Equal(0)) + }) It("leaves batch unset on non-Blackwell", func() { cfg := &ModelConfig{} - ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "9.0"}) + ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "9.0", VRAM: 119 * gib}) Expect(cfg.Batch).To(Equal(0)) }) It("never overrides an explicit batch", func() { cfg := &ModelConfig{} cfg.Batch = 1024 - ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "12.1"}) + ApplyHardwareDefaults(cfg, GPU{ComputeCapability: "12.1", VRAM: 119 * gib}) Expect(cfg.Batch).To(Equal(1024)) }) It("no-ops on nil", func() { @@ -57,8 +92,6 @@ var _ = Describe("Hardware-driven config defaults", func() { }) }) - const gib = uint64(1) << 30 - DescribeTable("DefaultParallelSlots (by VRAM)", func(vramGiB uint64, want int) { Expect(DefaultParallelSlots(GPU{VRAM: vramGiB * gib})).To(Equal(want)) diff --git a/core/config/model_config.go b/core/config/model_config.go index 8886ddfd5a5d..2d1e18cc7954 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -1204,11 +1204,6 @@ func (cfg *ModelConfig) SetDefaults(opts ...ConfigLoaderOption) { // This ensures gallery-installed and runtime-loaded models get optimal parameters. ApplyInferenceDefaults(cfg, cfg.Name, cfg.Model) - // Apply hardware-driven defaults (e.g. a larger physical batch on Blackwell). - // Uses the local GPU here; in distributed mode the router re-applies the same - // heuristics for the selected node's GPU before loading. Explicit config wins. - ApplyHardwareDefaults(cfg, localGPU()) - // Apply serving-policy defaults (device-independent): cross-request prefix // caching. Propagates to distributed nodes via the model options. ApplyServingDefaults(cfg) @@ -1247,6 +1242,16 @@ func (cfg *ModelConfig) SetDefaults(opts ...ConfigLoaderOption) { cfg.ContextSize = &ctx } runBackendHooks(cfg, lo.modelPath) + + // Apply hardware-driven defaults (e.g. a larger physical batch on Blackwell) + // LAST, after the context size is fully resolved (explicit config, LoadOptions, + // then the GGUF guess inside runBackendHooks): the Blackwell batch guard sizes + // the per-device compute buffer against this model's context, so it must see + // the final value, not a pre-guess nil. Uses the local GPU here; in distributed + // mode the router re-applies the same heuristics for the selected node's GPU + // before loading. Explicit config always wins. + ApplyHardwareDefaults(cfg, localGPU()) + cfg.syncKnownUsecasesFromString() } diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index f26fea2b988b..6ad550cf1096 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -156,7 +156,10 @@ func applyNodeHardwareDefaults(opts *pb.ModelOptions, node *BackendNode) { VRAM: node.TotalVRAM, } if config.IsManagedPhysicalBatch(int(opts.NBatch)) { - opts.NBatch = int32(config.PhysicalBatch(gpu)) + // Gate the raised batch on the selected node's per-device VRAM at this + // model's context, so a large context can't overflow the node's compute + // buffer (issue #10485). node.TotalVRAM is the node's reported ceiling. + opts.NBatch = int32(config.PhysicalBatchForContext(gpu, int(opts.ContextSize))) } // Default concurrent serving for the selected node (the frontend that built // the options may have no GPU). Only adds when no parallel option is set. diff --git a/core/services/nodes/router_hardware_internal_test.go b/core/services/nodes/router_hardware_internal_test.go index 2418bf4440df..d8576c4e4c26 100644 --- a/core/services/nodes/router_hardware_internal_test.go +++ b/core/services/nodes/router_hardware_internal_test.go @@ -8,12 +8,19 @@ import ( ) var _ = Describe("applyNodeHardwareDefaults", func() { - It("raises a managed default batch on a Blackwell node", func() { - opts := &pb.ModelOptions{NBatch: config.DefaultPhysicalBatch} - applyNodeHardwareDefaults(opts, &BackendNode{GPUComputeCapability: "12.1"}) + It("raises a managed default batch on a Blackwell node with headroom", func() { + opts := &pb.ModelOptions{NBatch: config.DefaultPhysicalBatch, ContextSize: 8192} + applyNodeHardwareDefaults(opts, &BackendNode{GPUComputeCapability: "12.1", TotalVRAM: 119 << 30}) Expect(opts.NBatch).To(BeEquivalentTo(config.BlackwellPhysicalBatch)) }) + It("keeps the default batch when a large context would overflow the node", func() { + // Regression guard for issue #10485 on the distributed path. + opts := &pb.ModelOptions{NBatch: config.DefaultPhysicalBatch, ContextSize: 204800} + applyNodeHardwareDefaults(opts, &BackendNode{GPUComputeCapability: "12.0", TotalVRAM: 16 << 30}) + Expect(opts.NBatch).To(BeEquivalentTo(config.DefaultPhysicalBatch)) + }) + It("resets a Blackwell guess on a non-Blackwell node", func() { // frontend (Blackwell) guessed high, but the selected node is not Blackwell opts := &pb.ModelOptions{NBatch: config.BlackwellPhysicalBatch} diff --git a/pkg/xsysinfo/gpu.go b/pkg/xsysinfo/gpu.go index f0185ddeb6bc..da183212f46e 100644 --- a/pkg/xsysinfo/gpu.go +++ b/pkg/xsysinfo/gpu.go @@ -129,6 +129,61 @@ func TotalAvailableVRAM() (uint64, error) { return 0, nil } +// MinPerGPUVRAM returns the total VRAM of the SMALLEST GPU on the host (in +// bytes), or 0 when no per-device VRAM is known. Unlike TotalAvailableVRAM +// (which sums across devices) this reports a single device's ceiling, which is +// the right figure for decisions about what must fit on one card: the compute +// buffer (sized by n_ubatch) and the parallel-slot tier. Summing a multi-GPU +// host's VRAM over-provisions those into a per-device OOM (issue #10485). +// +// Unified-memory devices (GB10, Apple) report system RAM as their single +// device's VRAM, so they are unaffected. +func MinPerGPUVRAM() (uint64, error) { + // Prefer per-device binary detection (nvidia-smi/rocm-smi report true + // per-card VRAM); ghw's per-card memory can reflect NUMA node RAM on some + // hosts, which is why TotalAvailableVRAM treats it as a sum. + if infos := GetGPUMemoryUsage(); len(infos) > 0 { + if v := minNonZeroVRAM(infos); v > 0 { + return v, nil + } + } + + // Fallback: ghw per-card memory, taking the minimum non-zero card. + if gpus, err := GPUs(); err == nil { + var min uint64 + for _, gpu := range gpus { + if gpu == nil || gpu.Node == nil || gpu.Node.Memory == nil { + continue + } + if b := gpu.Node.Memory.TotalUsableBytes; b > 0 { + if u := uint64(b); min == 0 || u < min { + min = u + } + } + } + if min > 0 { + return min, nil + } + } + + return 0, nil +} + +// minNonZeroVRAM returns the smallest non-zero TotalVRAM across the given GPUs, +// or 0 when none report VRAM. +func minNonZeroVRAM(infos []GPUMemoryInfo) uint64 { + var min uint64 + for _, g := range infos { + if g.TotalVRAM == 0 { + continue + } + if min == 0 || g.TotalVRAM < min { + min = g.TotalVRAM + } + } + return min +} + func HasGPU(vendor string) bool { gpus, err := GPUs() if err != nil { diff --git a/pkg/xsysinfo/minvram_internal_test.go b/pkg/xsysinfo/minvram_internal_test.go new file mode 100644 index 000000000000..ccd72dd5e620 --- /dev/null +++ b/pkg/xsysinfo/minvram_internal_test.go @@ -0,0 +1,37 @@ +package xsysinfo + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("minNonZeroVRAM", func() { + const gib = uint64(1) << 30 + + It("returns the smallest device on a multi-GPU host", func() { + // Two unequal cards (e.g. RTX 5070 Ti + 5060 Ti, both 16 GiB, or a + // mixed pair): the smallest device is the per-card allocation ceiling. + infos := []GPUMemoryInfo{ + {TotalVRAM: 16 * gib}, + {TotalVRAM: 12 * gib}, + } + Expect(minNonZeroVRAM(infos)).To(Equal(12 * gib)) + }) + + It("ignores devices that report zero VRAM", func() { + infos := []GPUMemoryInfo{ + {TotalVRAM: 0}, + {TotalVRAM: 24 * gib}, + } + Expect(minNonZeroVRAM(infos)).To(Equal(24 * gib)) + }) + + It("returns the single device's VRAM on a one-GPU host", func() { + Expect(minNonZeroVRAM([]GPUMemoryInfo{{TotalVRAM: 16 * gib}})).To(Equal(16 * gib)) + }) + + It("returns 0 when no device reports VRAM", func() { + Expect(minNonZeroVRAM([]GPUMemoryInfo{{TotalVRAM: 0}})).To(BeZero()) + Expect(minNonZeroVRAM(nil)).To(BeZero()) + }) +})