WebGPU 实战：在浏览器里跑 AI 模型，前端的下一个超能力

1
export async function checkWebGPU(): Promise<{
2
  supported: boolean;
3
  adapter: GPUAdapter | null;
4
  reason?: string;
5
}> {
6
  // 1. API 是否存在
7
  if (!navigator.gpu) {
8
    return { supported: false, adapter: null, reason: 'navigator.gpu 不存在，浏览器不支持 WebGPU' };
9
  }
10

11
  // 2. 能否拿到 adapter
12
  const adapter = await navigator.gpu.requestAdapter({
13
    powerPreference: 'high-performance', // 优先独显
14
  });
15

16
  if (!adapter) {
17
    return { supported: false, adapter: null, reason: '无法获取 GPU adapter，可能是驱动问题' };
18
  }
19

20
  // 3. 检查关键 feature
21
  const info = await adapter.requestAdapterInfo();
22
  console.log('GPU Info:', {
23
    vendor: info.vendor,
24
    architecture: info.architecture,
25
    device: info.device,
26
    description: info.description,
27
  });
28

29
  // 4. 检查内存限制
30
  const maxBufferSize = adapter.limits.maxBufferSize;
31
  const maxStorageSize = adapter.limits.maxStorageBufferBindingSize;
32
  console.log(`Max buffer: ${(maxBufferSize / 1024 / 1024).toFixed(0)}MB`);
33
  console.log(`Max storage binding: ${(maxStorageSize / 1024 / 1024).toFixed(0)}MB`);
34

35
  // 小于 256MB 的 GPU 跑不了大多数模型
36
  if (maxBufferSize < 256 * 1024 * 1024) {
37
    return { supported: false, adapter, reason: `GPU 内存不足: ${(maxBufferSize / 1024 / 1024).toFixed(0)}MB` };
38
  }
39

40
  return { supported: true, adapter };
41
}

1
import { pipeline, env } from '@huggingface/transformers';
2
import { checkWebGPU } from './gpu-detect';
3

4
// 配置：优先使用 WebGPU
5
env.backends.onnx.wasm.numThreads = 4; // WASM 兜底时的线程数
6

7
async function main() {
8
  const status = document.getElementById('status')!;
9
  const output = document.getElementById('output')!;
10

11
  // Step 1: 检测 GPU
12
  status.textContent = '正在检测 GPU...';
13
  const gpu = await checkWebGPU();
14

15
  if (!gpu.supported) {
16
    status.textContent = `WebGPU 不可用: ${gpu.reason}，将回退到 WASM`;
17
  }
18

19
  // Step 2: 加载模型（自动下载 + 缓存到 Cache API）
20
  status.textContent = '正在加载模型（首次需要下载约 200MB）...';
21
  const startLoad = performance.now();
22

23
  const summarizer = await pipeline(
24
    'summarization',
25
    'Xenova/distilbart-cnn-6-6', // ~200MB 量化模型
26
    {
27
      device: gpu.supported ? 'webgpu' : 'wasm',
28
      dtype: 'q4',  // 4-bit 量化，大幅减小体积
29
    }
30
  );
31

32
  const loadTime = ((performance.now() - startLoad) / 1000).toFixed(1);
33
  status.textContent = `模型加载完成 (${loadTime}s)，使用 ${gpu.supported ? 'WebGPU' : 'WASM'} 后端`;
34

35
  // Step 3: 推理
36
  const article = `
37
    WebGPU is a new web API that exposes modern GPU capabilities for rendering
38
    and computation. Unlike WebGL, WebGPU provides access to more advanced GPU
39
    features and enables more efficient interaction with the GPU for both
40
    graphics and general-purpose computations. It is designed to work across
41
    different GPU architectures and platforms, providing a unified API.
42
  `;
43

44
  const startInfer = performance.now();
45
  const result = await summarizer(article, {
46
    max_new_tokens: 100,
47
    do_sample: false,
48
  });
49
  const inferTime = ((performance.now() - startInfer) / 1000).toFixed(2);
50

51
  output.textContent = `[推理耗时 ${inferTime}s]\n${result[0].summary_text}`;
52
}
53

54
main().catch(console.error);

1
import { pipeline } from '@huggingface/transformers';
2

3
let model: any = null;
4

5
self.onmessage = async (e) => {
6
  const { type, payload } = e.data;
7

8
  switch (type) {
9
    case 'INIT':
10
      model = await pipeline(payload.task, payload.model, {
11
        device: 'webgpu',
12
        dtype: 'q4',
13
      });
14
      self.postMessage({ type: 'READY' });
15
      break;
16

17
    case 'INFER':
18
      const result = await model(payload.input, payload.options);
19
      self.postMessage({ type: 'RESULT', data: result });
20
      break;
21

22
    case 'DISPOSE':
23
      await model?.dispose();
24
      model = null;
25
      break;
26
  }
27
};
28

29
// 主线程调用
30
// main.ts
31
const worker = new Worker(
32
  new URL('./inference-worker.ts', import.meta.url),
33
  { type: 'module' }
34
);
35

36
function infer(input: string): Promise<any> {
37
  return new Promise((resolve) => {
38
    const handler = (e: MessageEvent) => {
39
      if (e.data.type === 'RESULT') {
40
        worker.removeEventListener('message', handler);
41
        resolve(e.data.data);
42
      }
43
    };
44
    worker.addEventListener('message', handler);
45
    worker.postMessage({ type: 'INFER', payload: { input } });
46
  });
47
}

1
class InferenceRouter {
2
  private localEngine: LocalInference | null = null;
3
  private remoteClient: RemoteAPI;
4

5
  async infer(input: string, options: InferOptions): Promise<InferResult> {
6
    // 决策逻辑
7
    const useLocal = await this.shouldUseLocal(options);
8

9
    if (useLocal) {
10
      try {
11
        return await this.localEngine!.run(input, options);
12
      } catch (e) {
13
        console.warn('本地推理失败，回退到远程:', e);
14
        return await this.remoteClient.infer(input, options);
15
      }
16
    }
17

18
    return await this.remoteClient.infer(input, options);
19
  }
20

21
  private async shouldUseLocal(options: InferOptions): Promise<boolean> {
22
    // 1. 本地引擎是否就绪
23
    if (!this.localEngine?.isReady()) return false;
24

25
    // 2. 任务复杂度是否在本地能力范围内
26
    if (options.maxTokens > 2000) return false;
27

28
    // 3. 是否需要隐私保护
29
    if (options.privacy === 'strict') return true;
30

31
    // 4. 网络状态
32
    if (!navigator.onLine) return true;
33

34
    // 5. 默认：小任务本地，大任务远程
35
    return options.expectedComplexity === 'low';
36
  }
37
}