fix(ai): offer the open model only on a GPU with 16-bit floats

README.md

@@ -27,7 +27,8 @@ that **burns after reading**.
   a model that runs in your browser, so nothing leaves the end-to-end encryption. It uses Chrome's
   built-in Gemini Nano when the browser has it, and otherwise downloads an open model
   (Qwen2.5-1.5B-Instruct, about 1 GB, once, after you ask) and runs it on the GPU with
-  [WebLLM](https://github.com/mlc-ai/web-llm). Browsers with neither show no on-device AI.
+  [WebLLM](https://github.com/mlc-ai/web-llm), when the GPU supports 16-bit floats (WebGPU's
+  `shader-f16`). Browsers with neither show no on-device AI.
 - **Optional AI on the Internet Computer** through the IC's [LLM canister](https://github.com/dfinity/llm),
   **off by default**. Text sent to it leaves the end-to-end encryption, so a controller has to switch it
   on (`ai_enabled`); when on, it is a second choice in the assistant, nothing is sent until you choose an

frontend/src/lib/localAi.test.ts

@@ -10,13 +10,27 @@ vi.mock("@mlc-ai/web-llm", () => webllm)
 
 afterEach(() => {
   vi.unstubAllGlobals()
+  vi.clearAllMocks()
 })
 
 const SHA1 = "[0-9a-f]{40}"
 const SRI = /^sha256-[A-Za-z0-9+/]{43}=$/
 
+/** Gives the page a WebGPU adapter with these features. */
+function gpuWith(...features: string[]) {
+  const requestAdapter = vi.fn(async (_options?: unknown) => ({ features: new Set(features) }))
+  vi.stubGlobal("navigator", { gpu: { requestAdapter } })
+  return requestAdapter
+}
+
+/** The module afresh, since it remembers what it detected. */
+function freshLocalAi() {
+  vi.resetModules()
+  return import("./localAi")
+}
+
 describe("the open model", () => {
-  it("is web-llm's own record for it, pinned to commits and checked against hashes", async () => {
+  it("is web-llm's own record for it, pinned to commits, checked against hashes and needing 16-bit floats", async () => {
     const { WEBLLM_CONFIG, WEBLLM_MODEL } = await import("./localAi")
     const { prebuiltAppConfig } = await vi.importActual<{ prebuiltAppConfig: AppConfig }>("@mlc-ai/web-llm")
     const upstream = prebuiltAppConfig.model_list.find((m) => m.model_id === WEBLLM_MODEL) as ModelRecord
@@ -30,17 +44,16 @@ describe("the open model", () => {
       new RegExp(`^https://raw\\.githubusercontent\\.com/mlc-ai/binary-mlc-llm-libs/${SHA1}/`),
     )
     expect(model_lib.replace(new RegExp(`/${SHA1}/`), "/main/")).toBe(upstreamLib)
-    expect(rest).toEqual(upstreamRest)
+    expect(rest).toEqual({ ...upstreamRest, required_features: ["shader-f16"] })
     expect(integrity?.config).toMatch(SRI)
     expect(integrity?.model_lib).toMatch(SRI)
     expect(integrity?.tokenizer?.["tokenizer.json"]).toMatch(SRI)
     expect(integrity?.onFailure).toBe("error")
   })
 
   it("is downloaded and removed through the pinned record", async () => {
-    vi.stubGlobal("navigator", { gpu: { requestAdapter: async () => ({}) } })
-    vi.resetModules()
-    const { WEBLLM_CONFIG, WEBLLM_MODEL, forgetLocalModel, prepareLocalAi } = await import("./localAi")
+    gpuWith("shader-f16")
+    const { WEBLLM_CONFIG, WEBLLM_MODEL, forgetLocalModel, prepareLocalAi } = await freshLocalAi()
 
     await prepareLocalAi()
     expect(webllm.CreateMLCEngine).toHaveBeenCalledWith(
@@ -51,3 +64,38 @@ describe("the open model", () => {
     expect(webllm.deleteModelAllInfoInCache).toHaveBeenCalledWith(WEBLLM_MODEL, WEBLLM_CONFIG)
   })
 })
+
+describe("the on-device check", () => {
+  it("offers the open model on a GPU with 16-bit floats, checking the GPU web-llm would use", async () => {
+    const requestAdapter = gpuWith("shader-f16")
+    const { detectLocalAi, localAiStatus } = await freshLocalAi()
+    await detectLocalAi()
+    expect(localAiStatus()).toMatchObject({ status: "download", engine: "webllm" })
+    expect(requestAdapter).toHaveBeenCalledWith({ powerPreference: "high-performance" })
+  })
+
+  // Without them, web-llm 0.2.85 downloads the whole model and then fails to build its shaders.
+  it("does not offer it on a GPU without them", async () => {
+    gpuWith()
+    const { NO_F16_REASON, detectLocalAi, localAiStatus } = await freshLocalAi()
+    await detectLocalAi()
+    expect(localAiStatus()).toEqual({ status: "unavailable", reason: NO_F16_REASON })
+  })
+
+  it("says why if asked to run anyway", async () => {
+    gpuWith()
+    const { NO_F16_REASON, localAiError, prepareLocalAi } = await freshLocalAi()
+    const error = await prepareLocalAi().catch((e: unknown) => e)
+    expect(webllm.CreateMLCEngine).not.toHaveBeenCalled()
+    expect(localAiError(error)).toBe(NO_F16_REASON)
+  })
+
+  it("explains web-llm's error for a GPU without them", async () => {
+    const { NO_F16_REASON, localAiError } = await import("./localAi")
+    const error = new Error(
+      "This model requires WebGPU extension shader-f16, which is not enabled in this browser.",
+    )
+    error.name = "ShaderF16SupportError"
+    expect(localAiError(error)).toBe(NO_F16_REASON)
+  })
+})

frontend/src/lib/localAi.ts

@@ -31,8 +31,8 @@ export const WEBLLM_MODEL = "Qwen2.5-1.5B-Instruct-q4f16_1-MLC"
  * This record fixes both at a commit, and web-llm checks the config, the
  * tokenizer and the library against these hashes before it runs anything
  * (`onFailure: "error"`). The rest is web-llm 0.2.85's own record for the
- * model (localAi.test.ts compares them); a web-llm upgrade that moves to new
- * model libraries needs a new pin.
+ * model (localAi.test.ts compares them), plus `required_features`; a web-llm
+ * upgrade that moves to new model libraries needs a new pin.
  */
 export const WEBLLM_CONFIG: AppConfig = {
   model_list: [
@@ -44,6 +44,10 @@ export const WEBLLM_CONFIG: AppConfig = {
         "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/025bcaf3780fa8254f5e5efd3bfea0a5397248f4/web-llm-models/v0_2_84/base/Qwen2-1.5B-Instruct-q4f16_1_cs1k-webgpu.wasm",
       low_resource_required: true,
       vram_required_MB: 1629.75,
+      // web-llm's record leaves this out, though the library's shaders use
+      // 16-bit floats. With it, web-llm stops on a GPU without them before
+      // downloading the weights, instead of failing once all are downloaded.
+      required_features: ["shader-f16"],
       overrides: { context_window_size: 4096 },
       integrity: {
         config: "sha256-faSVQhLT8hqRzD5pPRefYyE/1XNWeWhH9l1/lUeKH54=",
@@ -67,6 +71,8 @@ const LANGUAGE = {
 }
 export const UNAVAILABLE_REASON =
   "This browser can't run a model on this device. It needs Chrome's built-in AI or WebGPU: try a recent Chrome or Edge on a laptop or desktop."
+export const NO_F16_REASON =
+  "This device's GPU can't run the model: it lacks the 16-bit float support the model needs. Try a recent Chrome or Edge on another laptop or desktop."
 
 interface Runner {
   generate(prompt: LocalPrompt, signal: AbortSignal, onText: (text: string) => void): Promise<string>
@@ -135,7 +141,13 @@ export function detectLocalAi(): Promise<void> {
       }
     }
     const gpu = (navigator as any).gpu
-    const adapter = gpu ? await gpu.requestAdapter().catch(() => null) : null
+    // The GPU web-llm asks for, so this checks the one the model would run on.
+    const adapter = gpu
+      ? await gpu.requestAdapter({ powerPreference: "high-performance" }).catch(() => null)
+      : null
+    if (adapter && !adapter.features?.has("shader-f16")) {
+      return set({ status: "unavailable", reason: NO_F16_REASON })
+    }
     if (adapter) {
       engine = "webllm"
       return set(
@@ -209,7 +221,7 @@ export function prepareLocalAi(): Promise<Runner> {
   preparing = (async () => {
     await detectLocalAi()
     const current = engine
-    if (!current) throw new Error(UNAVAILABLE_REASON)
+    if (!current) throw new Error(state.status === "unavailable" ? state.reason : UNAVAILABLE_REASON)
     const model = MODEL_NAME[current]
     try {
       if (current === "chrome") {
@@ -282,6 +294,9 @@ export async function forgetLocalModel(): Promise<void> {
 /** Explains a failed run in the app's words. */
 export function localAiError(error: unknown): string {
   const message = error instanceof Error ? error.message : String(error)
+  // Both name a GPU, so they would read as running out of memory below.
+  if (message === UNAVAILABLE_REASON || message === NO_F16_REASON) return message
+  if (error instanceof Error && error.name === "ShaderF16SupportError") return NO_F16_REASON
   if (/memory|OOM|device lost|GPU/i.test(message))
     return "The model ran out of memory on this device. Close other tabs and try again."
   if (/user activation|gesture/i.test(message))