Merge pull request #121 from leehack/fix/webgpu-batch-cascade

fix: cascade WebGPU batch defaults

Author: leehack

CHANGELOG.md

@@ -40,6 +40,10 @@
   * Forwarded native-compatible `ModelParams` load tuning knobs through the
     WebGPU bridge path, including `maxParallelSequences`, flash attention,
     KV-cache type, KV-unified, RoPE, split-mode, and main-GPU options.
+  * Matched native batch defaults on the WebGPU path so unset `batchSize` /
+    `microBatchSize` cascade to `n_batch = n_ctx` and `n_ubatch = n_batch`,
+    avoiding first-embedding aborts for BERT-class/non-causal encoder models
+    while preserving explicit caller values and Qwen3.5 web tuning.
 * **GPU device selection API**:
   * Added `ModelParams.mainGpu` and wired it to llama.cpp
     `llama_model_params.main_gpu`.

lib/src/backends/llama_cpp/llama_cpp_service.dart

@@ -325,36 +325,13 @@ class LlamaCppService {
 
   /// Resolves effective context batch parameters.
   ///
-  /// Legacy behavior is preserved when [ModelParams.batchSize] and
-  /// [ModelParams.microBatchSize] are not set:
-  ///
-  /// - `n_batch = n_ctx`
-  /// - `n_ubatch = n_batch`
-  ///
-  /// Values are clamped to safe bounds so `n_ubatch <= n_batch <= n_ctx`.
+  /// Uses the shared non-FFI helper so native and WebGPU batch semantics stay
+  /// in sync.
   static ({int batchSize, int microBatchSize}) resolveContextBatchSizes(
     ModelParams modelParams,
     int contextSize,
   ) {
-    final effectiveContextSize = contextSize > 0 ? contextSize : 1;
-
-    final configuredBatchSize = modelParams.batchSize > 0
-        ? modelParams.batchSize
-        : effectiveContextSize;
-    final cappedBatchSize = configuredBatchSize > effectiveContextSize
-        ? effectiveContextSize
-        : configuredBatchSize;
-    final batchSize = cappedBatchSize > 0 ? cappedBatchSize : 1;
-
-    final configuredMicroBatchSize = modelParams.microBatchSize > 0
-        ? modelParams.microBatchSize
-        : batchSize;
-    final cappedMicroBatchSize = configuredMicroBatchSize > batchSize
-        ? batchSize
-        : configuredMicroBatchSize;
-    final microBatchSize = cappedMicroBatchSize > 0 ? cappedMicroBatchSize : 1;
-
-    return (batchSize: batchSize, microBatchSize: microBatchSize);
+    return resolveModelContextBatchSizes(modelParams, contextSize);
   }
 
   /// Resolves whether multimodal projector init should use GPU.

lib/src/backends/webgpu/webgpu_backend.dart

@@ -591,29 +591,28 @@ class WebGpuLlamaBackend
     return attempts;
   }
 
-  ({int? nBatch, int? nUbatch}) _resolveWebBatchTuning({
+  ({int nBatch, int nUbatch}) _resolveWebBatchSizes({
     required String url,
     required ModelParams params,
+    required int contextSize,
   }) {
-    if (params.batchSize > 0 || params.microBatchSize > 0) {
-      return (nBatch: null, nUbatch: null);
-    }
-
-    if (params.preferredBackend == GpuBackend.cpu || params.gpuLayers == 0) {
-      return (nBatch: null, nUbatch: null);
-    }
-
     final normalizedUrl = url.toLowerCase();
     final isQwen35Small =
         normalizedUrl.contains('qwen3.5-0.8b') ||
         normalizedUrl.contains('qwen_qwen3.5-0.8b');
-    if (!isQwen35Small) {
-      return (nBatch: null, nUbatch: null);
+    final shouldUseQwen35SmallTuning =
+        params.batchSize <= 0 &&
+        params.microBatchSize <= 0 &&
+        params.preferredBackend != GpuBackend.cpu &&
+        params.gpuLayers != 0 &&
+        isQwen35Small;
+
+    if (shouldUseQwen35SmallTuning) {
+      return (nBatch: 32, nUbatch: 8);
     }
 
-    final tunedBatch = 32;
-    final tunedUbatch = 8;
-    return (nBatch: tunedBatch, nUbatch: tunedUbatch);
+    final resolved = resolveModelContextBatchSizes(params, contextSize);
+    return (nBatch: resolved.batchSize, nUbatch: resolved.microBatchSize);
   }
 
   int _webGpuFlashAttentionValue(ModelParams params) {
@@ -894,8 +893,6 @@ class WebGpuLlamaBackend
       requestedContextSize: params.contextSize,
       requestedGpuLayers: requestedGpuLayers,
     );
-    final batchTuning = _resolveWebBatchTuning(url: url, params: params);
-
     Object? lastError;
     Map<String, String> lastRuntimeHints = const <String, String>{};
     var retriedWithWasm32 = false;
@@ -909,6 +906,11 @@ class WebGpuLlamaBackend
     for (var index = 0; index < loadAttempts.length; index += 1) {
       final attempt = loadAttempts[index];
       _lastNCtx = attempt.contextSize;
+      final batchSizes = _resolveWebBatchSizes(
+        url: url,
+        params: params,
+        contextSize: attempt.contextSize,
+      );
       LlamaWebGpuBridge? bridgeForAttempt;
       bool? forceRemoteFetchBackend;
       final attemptThreads = switch (index) {
@@ -940,12 +942,8 @@ class WebGpuLlamaBackend
             nThreadsBatch: params.numberOfThreadsBatch > 0
                 ? params.numberOfThreadsBatch
                 : null,
-            nBatch: params.batchSize > 0
-                ? params.batchSize
-                : batchTuning.nBatch,
-            nUbatch: params.microBatchSize > 0
-                ? params.microBatchSize
-                : batchTuning.nUbatch,
+            nBatch: batchSizes.nBatch,
+            nUbatch: batchSizes.nUbatch,
             nGpuLayers: attempt.gpuLayers,
             nSeqMax: math.max(1, params.maxParallelSequences),
             flashAttention: _webGpuFlashAttentionValue(params),

lib/src/core/models/inference/model_params.dart

@@ -248,3 +248,37 @@ class ModelParams {
     );
   }
 }
+
+/// Resolves llama.cpp-compatible context batch parameters.
+///
+/// Preserves native defaults when [ModelParams.batchSize] and
+/// [ModelParams.microBatchSize] are unset:
+///
+/// - `n_batch = n_ctx`
+/// - `n_ubatch = n_batch`
+///
+/// Values are clamped to safe bounds so `n_ubatch <= n_batch <= n_ctx`.
+({int batchSize, int microBatchSize}) resolveModelContextBatchSizes(
+  ModelParams modelParams,
+  int contextSize,
+) {
+  final effectiveContextSize = contextSize > 0 ? contextSize : 1;
+
+  final configuredBatchSize = modelParams.batchSize > 0
+      ? modelParams.batchSize
+      : effectiveContextSize;
+  final cappedBatchSize = configuredBatchSize > effectiveContextSize
+      ? effectiveContextSize
+      : configuredBatchSize;
+  final batchSize = cappedBatchSize > 0 ? cappedBatchSize : 1;
+
+  final configuredMicroBatchSize = modelParams.microBatchSize > 0
+      ? modelParams.microBatchSize
+      : batchSize;
+  final cappedMicroBatchSize = configuredMicroBatchSize > batchSize
+      ? batchSize
+      : configuredMicroBatchSize;
+  final microBatchSize = cappedMicroBatchSize > 0 ? cappedMicroBatchSize : 1;
+
+  return (batchSize: batchSize, microBatchSize: microBatchSize);
+}