uploaded production gemma with vulkan

Author: mkopcins

packages/react-native-executorch/common/runner/base_llm_runner.cpp

@@ -18,8 +18,8 @@ BaseLLMRunner::BaseLLMRunner(std::unique_ptr<Module> module,
       tokenizer_(std::make_unique<tokenizers::HFTokenizer>()),
       metadata_({
           {kEnableDynamicShape, false},
-          {kMaxSeqLen, 128},
-          {kMaxContextLen, 128},
+          {kMaxSeqLen, 2048},
+          {kMaxContextLen, 2048},
           {kUseKVCache, true},
       }) {}
 
@@ -56,11 +56,16 @@ Error BaseLLMRunner::load() {
             ? static_cast<int32_t>(metadata_.at(kMaxContextLen))
             : static_cast<int32_t>(metadata_.at(kMaxSeqLen));
   }
-  if (config_.max_new_tokens < 0)
-    config_.max_new_tokens =
-        std::min(config_.max_seq_len, config_.max_context_length);
   config_.enable_dynamic_shape =
       static_cast<bool>(metadata_.at(kEnableDynamicShape));
+  if (config_.max_new_tokens < 0) {
+    // For dynamic-shape PTEs, max_seq_len is the per-call decoder chunk
+    // size, not the generation budget — use max_context_length instead.
+    const int32_t seq_cap = config_.enable_dynamic_shape
+                                ? config_.max_context_length
+                                : config_.max_seq_len;
+    config_.max_new_tokens = std::min(seq_cap, config_.max_context_length);
+  }
   config_.enable_kv_cache = static_cast<bool>(metadata_.at(kUseKVCache));
 
   eos_ids_ = std::make_unique<std::unordered_set<uint64_t>>();

packages/react-native-executorch/common/runner/constants.h

@@ -17,23 +17,28 @@ inline constexpr auto kMaxSeqLen = "get_max_seq_len";
 inline constexpr auto kMaxContextLen = "get_max_context_len";
 inline constexpr auto kVocabSize = "get_vocab_size";
 inline constexpr auto kUseKVCache = "use_kv_cache";
-// PLE models only: token id that marks image placeholder slots in input_ids.
-// token_embedding run on this id produces the per-layer PLE signal for image
-// positions; the inputs_embeds output for those positions is discarded (the
-// vision encoder output replaces it).
-inline constexpr auto kImagePlaceholderId = "image_placeholder_id";
 
 // Multimodal method name conventions
 inline constexpr auto kVisionEncoderMethod = "vision_encoder";
 inline constexpr auto kAudioEncoderMethod = "audio_encoder";
-inline constexpr auto kTokenEmbeddingMethod = "text_encoder";
+inline constexpr auto kTokenEmbeddingMethod = "token_embedding";
 inline constexpr auto kTextModelMethod = "text_decoder";
-// Absolute ceiling on prefill length (in tokens) and the fallback value used
-// when a PTE doesn't bake `get_max_seq_len`. 2048 matches Gemma4 iter201's
-// PREFILL_LEN / get_max_context_len; legacy PTEs (e.g. LFM2-VL) typically
-// bake their own get_max_seq_len so this ceiling does not affect them.
-inline constexpr auto kMaxPrefillLen = 2048;
 inline constexpr auto numOfAddedBoSTokens = 0;
 inline constexpr auto numOfAddedEoSTokens = 0;
 
+// Gemma4
+// PLE models only: token id that marks image placeholder slots in input_ids.
+// token_embedding run on this id produces the per-layer PLE signal for image
+// positions; the inputs_embeds output for those positions is discarded (the
+// vision encoder output replaces it).
+inline constexpr auto kImagePlaceholderId = "image_placeholder_id";
+// True iff the model exposes a per-layer-embedding (PLE) signal alongside
+// inputs_embeds (Gemma4-style). When true, `token_embedding.execute()`
+// returns the tuple (inputs_embeds, ple_tok) and the runner must thread
+// ple_tok into text_decoder; when false (or absent), token_embedding returns
+// inputs_embeds alone. Text-only PTEs that ship a single `forward` method
+// omit this key entirely — it is meaningful only for multimodal PTEs that
+// expose a separate `token_embedding` method.
+inline constexpr auto kHasPLE = "has_ple";
+
 } // namespace executorch::extension::llm

packages/react-native-executorch/common/runner/multimodal_decoder_runner.h

@@ -14,14 +14,16 @@
 #include "text_decoder_runner.h"
 
 namespace executorch::extension::llm {
-// Supports two PTE contracts, selected automatically at load time from
-// `token_embedding`'s output arity:
+// Supports two PTE contracts, selected per-call from the kHasPLE metadata
+// key (mirrors how kEnableDynamicShape etc. are read — queried on demand,
+// not cached in a member). Callers that need it multiple times in a hot
+// path should snapshot into a local.
 //
-//  * Legacy (default):
+//  * Legacy (has_ple == false):
 //      token_embedding(ids) -> inputs_embeds
 //      text_decoder(inputs_embeds, input_pos)
 //
-//  * Gemma-style PLE (when token_embedding emits 2 outputs):
+//  * Gemma-style PLE (has_ple == true):
 //      token_embedding(ids) -> (inputs_embeds, ple_tok)
 //      text_decoder(inputs_embeds, ple_tok, input_pos)
 //    ple_tok carries Gemma4's per-layer PLE signal keyed on input_ids. It's
@@ -34,8 +36,15 @@ class MultimodalDecoderRunner : public TextDecoderRunner {
       : TextDecoderRunner(module, io_manager, config) {}
 
   // True iff the loaded PTE uses the Gemma-style PLE contract above.
-  // Meaningful only after load() has been called.
-  bool uses_ple() const { return uses_ple_; }
+  // Reads the kHasPLE constant_method every call; cheap, but callers in
+  // hot loops should snapshot into a local.
+  bool has_ple() const {
+    auto r = module_->get(kHasPLE);
+    if (r.error() != ::executorch::runtime::Error::Ok) {
+      return false;
+    }
+    return r->toScalar().to<bool>();
+  }
 
   inline ::executorch::runtime::Result<::executorch::aten::Tensor>
   step(TensorPtr &tokens, int64_t start_pos) override {
@@ -44,7 +53,7 @@ class MultimodalDecoderRunner : public TextDecoderRunner {
       return embed_result.error();
     }
     auto &embed_outputs = *embed_result;
-    if (uses_ple_) {
+    if (has_ple()) {
       ET_CHECK_MSG(embed_outputs.size() == 2,
                    "Expected 2 outputs (inputs_embeds, ple_tok) from "
                    "token_embedding, got %zu",
@@ -97,20 +106,13 @@ class MultimodalDecoderRunner : public TextDecoderRunner {
     }
     ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kTokenEmbeddingMethod));
     ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kTextModelMethod));
-
-    auto meta = module_->method_meta(kTokenEmbeddingMethod);
-    ET_CHECK_OK_OR_RETURN_ERROR(meta.error());
-    uses_ple_ = (meta->num_outputs() == 2);
     return ::executorch::runtime::Error::Ok;
   }
 
   inline bool is_method_loaded() override {
     return module_->is_method_loaded(kTokenEmbeddingMethod) &&
            module_->is_method_loaded(kTextModelMethod);
   }
-
-private:
-  bool uses_ple_ = true;
 };
 
 } // namespace executorch::extension::llm