From d85709f900f1cb38b3bad4e2f5c0de3c15c39c6f Mon Sep 17 00:00:00 2001 From: NJX <3771829673@qq.com> Date: Fri, 24 Apr 2026 23:33:09 +0800 Subject: [PATCH] feat(stage-ui): webgpu detect improved (#1681) --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../stage-shared/src/webgpu/detect.test.ts | 335 ++++++++++++++++++ packages/stage-shared/src/webgpu/detect.ts | 177 ++++++++- packages/stage-shared/src/webgpu/index.ts | 4 +- 3 files changed, 507 insertions(+), 9 deletions(-) create mode 100644 packages/stage-shared/src/webgpu/detect.test.ts diff --git a/packages/stage-shared/src/webgpu/detect.test.ts b/packages/stage-shared/src/webgpu/detect.test.ts new file mode 100644 index 000000000..8cd70448b --- /dev/null +++ b/packages/stage-shared/src/webgpu/detect.test.ts @@ -0,0 +1,335 @@ +import { check as gpuuCheck } from 'gpuu/webgpu' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { + detectWebGPU, + getCachedWebGPUCapabilities, + getEstimatedVRAMOverride, + resetWebGPUCache, + setEstimatedVRAMOverride, +} from './detect' + +// Mock gpuu/webgpu before importing detect.ts +vi.mock('gpuu/webgpu', () => ({ + check: vi.fn(), + isWebGPUSupported: vi.fn(), +})) + +const mockedCheck = vi.mocked(gpuuCheck) + +interface MockAdapterInfo { + vendor?: string + architecture?: string + device?: string + description?: string +} + +/** + * Build a minimal mock GPUAdapter for tests. + * `maxBufferSize` drives the heuristic VRAM calculation. + * `info` optionally populates adapterInfo. + */ +function makeMockAdapter(options: { + maxBufferSize?: number + info?: MockAdapterInfo +}): any { + const adapter: { limits: { maxBufferSize: number }, info?: MockAdapterInfo } = { + limits: { maxBufferSize: options.maxBufferSize ?? 0 }, + } + if (options.info) + adapter.info = options.info + return adapter +} + +describe('detectWebGPU', () => { + beforeEach(() => { + resetWebGPUCache() + setEstimatedVRAMOverride(null) + }) + + afterEach(() => { + resetWebGPUCache() + setEstimatedVRAMOverride(null) + vi.clearAllMocks() + }) + + it('should derive VRAM from the maxBufferSize * 4 heuristic', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 256 * 1024 * 1024 }), + }) + + const result = await detectWebGPU() + + expect(result.supported).toBe(true) + expect(result.fp16Supported).toBe(true) + expect(result.estimatedVRAM).toBe(256 * 1024 * 1024 * 4) + expect(result.estimatedVRAMSource).toBe('max-buffer-heuristic') + }) + + it('should report "none" when the adapter has no maxBufferSize', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: false, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 0 }), + }) + + const result = await detectWebGPU() + + expect(result.estimatedVRAM).toBe(0) + expect(result.estimatedVRAMSource).toBe('none') + }) + + it('should report "none" when WebGPU is unsupported', async () => { + mockedCheck.mockResolvedValue({ + supported: false, + fp16Supported: false, + isNode: false, + reason: 'not available', + }) + + const result = await detectWebGPU() + + expect(result.supported).toBe(false) + expect(result.estimatedVRAM).toBe(0) + expect(result.estimatedVRAMSource).toBe('none') + expect(result.reason).toBe('not available') + }) + + it('should extract adapter.info when available', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ + maxBufferSize: 1024 * 1024 * 1024, + info: { + vendor: 'apple', + architecture: 'apple-m1', + device: 'Apple M1', + description: 'Apple GPU', + }, + }), + }) + + const result = await detectWebGPU() + + expect(result.adapterInfo).toEqual({ + vendor: 'apple', + architecture: 'apple-m1', + device: 'Apple M1', + description: 'Apple GPU', + }) + }) + + it('should handle adapter.info with missing fields gracefully', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ + maxBufferSize: 1024 * 1024 * 1024, + info: { vendor: 'nvidia' }, // only vendor set + }), + }) + + const result = await detectWebGPU() + + expect(result.adapterInfo).toEqual({ + vendor: 'nvidia', + architecture: '', + device: '', + description: '', + }) + }) + + it('should set adapterInfo to null when adapter.info is not exposed', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 512 * 1024 * 1024 }), + }) + + const result = await detectWebGPU() + + expect(result.adapterInfo).toBeNull() + }) + + it('should fall back to requestAdapterInfo() when adapter.info is absent', async () => { + const legacyInfo: MockAdapterInfo = { + vendor: 'intel', + architecture: 'xe', + device: 'Iris Xe', + description: 'Intel Xe Graphics', + } + + const legacyAdapter: any = { + limits: { maxBufferSize: 256 * 1024 * 1024 }, + requestAdapterInfo: vi.fn(async () => legacyInfo), + } + + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: false, + isNode: false, + reason: '', + adapter: legacyAdapter, + }) + + const result = await detectWebGPU() + + expect(result.adapterInfo).toEqual({ + vendor: 'intel', + architecture: 'xe', + device: 'Iris Xe', + description: 'Intel Xe Graphics', + }) + }) + + it('should cache the detection result across calls', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 1024 * 1024 }), + }) + + const first = await detectWebGPU() + const second = await detectWebGPU() + + expect(first).toBe(second) + expect(mockedCheck).toHaveBeenCalledTimes(1) + }) + + it('should deduplicate concurrent calls', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 1024 * 1024 }), + }) + + const [a, b] = await Promise.all([detectWebGPU(), detectWebGPU()]) + + expect(a).toBe(b) + expect(mockedCheck).toHaveBeenCalledTimes(1) + }) + + it('should produce a safe fallback when gpuu throws', async () => { + mockedCheck.mockRejectedValue(new Error('internal')) + + const result = await detectWebGPU() + + expect(result.supported).toBe(false) + expect(result.estimatedVRAM).toBe(0) + expect(result.estimatedVRAMSource).toBe('none') + expect(result.adapterInfo).toBeNull() + expect(result.reason).toBe('Detection threw an exception') + }) +}) + +describe('vRAM override', () => { + beforeEach(() => { + resetWebGPUCache() + setEstimatedVRAMOverride(null) + }) + + afterEach(() => { + resetWebGPUCache() + setEstimatedVRAMOverride(null) + vi.clearAllMocks() + }) + + it('should apply the override when detection runs', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }), + }) + + setEstimatedVRAMOverride(8 * 1024 * 1024 * 1024) // 8 GB + + const result = await detectWebGPU() + + expect(result.estimatedVRAM).toBe(8 * 1024 * 1024 * 1024) + expect(result.estimatedVRAMSource).toBe('override') + }) + + it('should update cached result in-place when override is set after detection', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }), + }) + + await detectWebGPU() + expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('max-buffer-heuristic') + + setEstimatedVRAMOverride(4 * 1024 * 1024 * 1024) + + expect(getCachedWebGPUCapabilities()?.estimatedVRAM).toBe(4 * 1024 * 1024 * 1024) + expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('override') + }) + + it('should revert to heuristic when override is cleared with null', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }), + }) + + await detectWebGPU() + setEstimatedVRAMOverride(4 * 1024 * 1024 * 1024) + expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('override') + + setEstimatedVRAMOverride(null) + + expect(getCachedWebGPUCapabilities()?.estimatedVRAM).toBe(128 * 1024 * 1024 * 4) + expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('max-buffer-heuristic') + }) + + it('should expose the current override via getEstimatedVRAMOverride()', () => { + expect(getEstimatedVRAMOverride()).toBeNull() + setEstimatedVRAMOverride(2 * 1024 * 1024 * 1024) + expect(getEstimatedVRAMOverride()).toBe(2 * 1024 * 1024 * 1024) + }) + + it('should reject invalid override values', () => { + expect(() => setEstimatedVRAMOverride(-1)).toThrow() + expect(() => setEstimatedVRAMOverride(Number.NaN)).toThrow() + expect(() => setEstimatedVRAMOverride(Number.POSITIVE_INFINITY)).toThrow() + }) + + it('should accept zero as a no-op override (reverts to heuristic)', async () => { + mockedCheck.mockResolvedValue({ + supported: true, + fp16Supported: true, + isNode: false, + reason: '', + adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }), + }) + + setEstimatedVRAMOverride(0) + const result = await detectWebGPU() + + // Zero is considered "no override" for the purpose of estimation + expect(result.estimatedVRAMSource).toBe('max-buffer-heuristic') + }) +}) diff --git a/packages/stage-shared/src/webgpu/detect.ts b/packages/stage-shared/src/webgpu/detect.ts index be4d933a5..c07a2863e 100644 --- a/packages/stage-shared/src/webgpu/detect.ts +++ b/packages/stage-shared/src/webgpu/detect.ts @@ -3,17 +3,73 @@ * * Wraps `gpuu/webgpu` and caches the result so every consumer * gets the same answer without redundant adapter requests. + * + * ## VRAM estimation + * + * The web platform does not expose GPU memory usage or total VRAM. + * We approximate via three ordered sources: + * + * 1. User override (`setEstimatedVRAMOverride(bytes)`) + * 2. `adapter.limits.maxBufferSize * 4` heuristic (fallback) + * 3. Zero (unavailable) + * + * The provenance is reported via `WebGPUCapabilities.estimatedVRAMSource` + * so consumers can surface it for diagnostics. */ import { check as gpuuCheck, isWebGPUSupported as gpuuIsSupported } from 'gpuu/webgpu' +// Minimal structural subset of the WebGPU types we interact with. +// Avoids depending on `@webgpu/types` (which is shipped transitively via +// transformers.js but not declared by this package). +interface GPUAdapterInfoLike { + vendor?: string + architecture?: string + device?: string + description?: string +} + +interface GPUAdapterLike { + limits?: { maxBufferSize?: number } + info?: GPUAdapterInfoLike + requestAdapterInfo?: () => Promise +} + +/** + * Subset of `GPUAdapterInfo` that we surface to consumers. Values come + * directly from the browser's WebGPU implementation — treat them as + * opaque strings; vendor/architecture naming is not standardized. + */ +export interface WebGPUAdapterInfo { + /** Vendor name, e.g. "nvidia", "apple", "intel" */ + vendor: string + /** Architecture name, e.g. "ada-lovelace", "apple-m1" */ + architecture: string + /** Device description, e.g. "NVIDIA GeForce RTX 4090" */ + device: string + /** Free-form description string from the driver */ + description: string +} + +/** + * Source of the VRAM estimate, reported for observability. + * - `override`: user-provided value via `setEstimatedVRAMOverride()` + * - `max-buffer-heuristic`: derived from `adapter.limits.maxBufferSize * 4` + * - `none`: no estimate available (WebGPU unsupported or adapter query failed) + */ +export type VRAMSource = 'override' | 'max-buffer-heuristic' | 'none' + export interface WebGPUCapabilities { /** Whether WebGPU is available in this environment */ supported: boolean /** Whether fp16 shader operations are supported */ fp16Supported: boolean - /** Estimated VRAM in bytes (heuristic, 0 when unavailable) */ + /** Estimated VRAM in bytes (0 when unavailable) */ estimatedVRAM: number + /** Provenance of the VRAM estimate */ + estimatedVRAMSource: VRAMSource + /** Adapter-reported vendor/architecture/device info, when available */ + adapterInfo: WebGPUAdapterInfo | null /** Raw reason string from gpuu when unsupported */ reason: string } @@ -21,6 +77,71 @@ export interface WebGPUCapabilities { let cachedResult: WebGPUCapabilities | null = null let pendingDetection: Promise | null = null +// NOTICE: User override for VRAM estimation. When set, this value takes +// priority over all heuristics. Useful for users with known hardware where +// the heuristic is inaccurate (e.g. discrete GPUs with small maxBufferSize). +let vramOverride: number | null = null + +// Cached heuristic value so we can restore it when the override is cleared. +// Computed during detectWebGPU() and persisted for the lifetime of the cache. +let cachedHeuristicVRAM = 0 + +/** + * Best-effort extraction of `GPUAdapterInfo` from a `GPUAdapter`. Tries + * the modern synchronous `adapter.info` first, then falls back to the + * legacy `requestAdapterInfo()` promise API. Returns null if neither works. + * + * References: + * - https://www.w3.org/TR/webgpu/#gpu-adapterinfo + */ +async function extractAdapterInfo(adapter: GPUAdapterLike): Promise { + try { + // Modern API: synchronous `info` property (Chrome 114+, Safari 17.4+) + const info = adapter.info + if (info) { + return { + vendor: info.vendor ?? '', + architecture: info.architecture ?? '', + device: info.device ?? '', + description: info.description ?? '', + } + } + + // Legacy API: requestAdapterInfo() returns a Promise + const legacy = adapter.requestAdapterInfo + if (typeof legacy === 'function') { + const legacyInfo = await legacy.call(adapter) + return { + vendor: legacyInfo.vendor ?? '', + architecture: legacyInfo.architecture ?? '', + device: legacyInfo.device ?? '', + description: legacyInfo.description ?? '', + } + } + } + catch { + // Fall through to null — adapter info is best-effort, not required + } + return null +} + +/** Compute the heuristic VRAM estimate from `maxBufferSize`. */ +function computeHeuristicVRAM(adapter: GPUAdapterLike): number { + const maxBuffer = adapter.limits?.maxBufferSize ?? 0 + // Typical values: 256 MB on integrated GPUs, 2-4 GB on discrete. + // Multiply by 4 as a conservative total VRAM heuristic. + return maxBuffer > 0 ? maxBuffer * 4 : 0 +} + +/** Decide the VRAM estimate based on override > heuristic > none. */ +function resolveVRAM(heuristic: number): { bytes: number, source: VRAMSource } { + if (vramOverride !== null && vramOverride > 0) + return { bytes: vramOverride, source: 'override' } + if (heuristic > 0) + return { bytes: heuristic, source: 'max-buffer-heuristic' } + return { bytes: 0, source: 'none' } +} + /** * Detect WebGPU capabilities. The result is cached as a singleton * after the first successful call -- safe to call repeatedly. @@ -37,27 +158,33 @@ export async function detectWebGPU(): Promise { try { const result = await gpuuCheck() - let estimatedVRAM = 0 + let adapterInfo: WebGPUAdapterInfo | null = null + let heuristic = 0 if (result.supported && result.adapter) { - // Use maxBufferSize as a rough proxy -- typically 256 MB on - // integrated GPUs, 2-4 GB on discrete GPUs. - // Multiply by 4 as a conservative total VRAM heuristic. - const maxBuffer = result.adapter.limits?.maxBufferSize ?? 0 - estimatedVRAM = maxBuffer > 0 ? maxBuffer * 4 : 0 + heuristic = computeHeuristicVRAM(result.adapter) + adapterInfo = await extractAdapterInfo(result.adapter) } + cachedHeuristicVRAM = heuristic + const vram = resolveVRAM(heuristic) + cachedResult = { supported: result.supported, fp16Supported: result.fp16Supported ?? false, - estimatedVRAM, + estimatedVRAM: vram.bytes, + estimatedVRAMSource: vram.source, + adapterInfo, reason: result.reason ?? '', } } catch { + cachedHeuristicVRAM = 0 cachedResult = { supported: false, fp16Supported: false, estimatedVRAM: 0, + estimatedVRAMSource: 'none', + adapterInfo: null, reason: 'Detection threw an exception', } } @@ -86,10 +213,44 @@ export async function isWebGPUSupported(): Promise { return gpuuIsSupported() } +/** + * Override the estimated VRAM value. Pass `null` to clear the override and + * revert to the heuristic. The override applies to future detections, and if + * a result is already cached its VRAM fields are updated immediately, so + * `resetWebGPUCache()` is not required. + * + * Intended for user preference UI ("I have 8 GB VRAM") and testing. + */ +export function setEstimatedVRAMOverride(bytes: number | null): void { + if (bytes !== null && (!Number.isFinite(bytes) || bytes < 0)) + throw new Error(`Invalid VRAM override: ${bytes} (expected null or non-negative finite number)`) + + vramOverride = bytes + + // If we already have a cached result, update it in-place so consumers + // see the new value without needing to call resetWebGPUCache(). The + // original heuristic value is preserved in `cachedHeuristicVRAM` so we + // can revert when the override is cleared. + if (cachedResult) { + const vram = resolveVRAM(cachedHeuristicVRAM) + cachedResult = { + ...cachedResult, + estimatedVRAM: vram.bytes, + estimatedVRAMSource: vram.source, + } + } +} + +/** Read the current VRAM override, or null if unset. */ +export function getEstimatedVRAMOverride(): number | null { + return vramOverride +} + /** * Reset the cached detection result. Intended for tests only. */ export function resetWebGPUCache(): void { cachedResult = null pendingDetection = null + cachedHeuristicVRAM = 0 } diff --git a/packages/stage-shared/src/webgpu/index.ts b/packages/stage-shared/src/webgpu/index.ts index 287e25d88..d855540d6 100644 --- a/packages/stage-shared/src/webgpu/index.ts +++ b/packages/stage-shared/src/webgpu/index.ts @@ -1,7 +1,9 @@ export { detectWebGPU, getCachedWebGPUCapabilities, + getEstimatedVRAMOverride, isWebGPUSupported, resetWebGPUCache, + setEstimatedVRAMOverride, } from './detect' -export type { WebGPUCapabilities } from './detect' +export type { VRAMSource, WebGPUAdapterInfo, WebGPUCapabilities } from './detect'