@@ -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>>();
 
@@ -23,12 +23,22 @@ inline constexpr auto kVisionEncoderMethod = "vision_encoder";
 inline constexpr auto kAudioEncoderMethod = "audio_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
@@ -98,7 +98,7 @@ Result<EValue> AudioEncoder::encode(const MultimodalInput &input) {
   std::memcpy(padded_wav_.data(), wav.samples.data(),
               static_cast<size_t>(n_valid) * sizeof(float));
 
-  num_blocks_scalar_ = k_blocks;
+  num_blocks_scalar_ = n_valid;
 
   auto wav_tensor = ::executorch::extension::from_blob(
       padded_wav_.data(), {1, static_cast<SizesType>(n_padded)},
 
@@ -14,19 +14,53 @@
 #include "text_decoder_runner.h"
 
 namespace executorch::extension::llm {
+// 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 (has_ple == false):
+//      token_embedding(ids) -> inputs_embeds
+//      text_decoder(inputs_embeds, input_pos)
+//
+//  * 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
+//    computed once in token_embedding and threaded through every decoder call
+//    so PLE fires at every position (including multimodal placeholder slots).
 class MultimodalDecoderRunner : public TextDecoderRunner {
 public:
   explicit MultimodalDecoderRunner(Module &module, IOManager *io_manager,
                                    const GenerationConfig &config)
       : TextDecoderRunner(module, io_manager, config) {}
 
+  // True iff the loaded PTE uses the Gemma-style PLE contract above.
+  // 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 {
     auto embed_result = module_->execute(kTokenEmbeddingMethod, tokens);
     if (!embed_result.ok()) {
       return embed_result.error();
     }
-    return decode((*embed_result)[0], start_pos);
+    auto &embed_outputs = *embed_result;
+    if (has_ple()) {
+      ET_CHECK_MSG(embed_outputs.size() == 2,
+                   "Expected 2 outputs (inputs_embeds, ple_tok) from "
+                   "token_embedding, got %zu",
+                   embed_outputs.size());
+      return decode(embed_outputs[0], embed_outputs[1], start_pos);
+    }
+    return decode(embed_outputs[0], start_pos);
   }
 
   inline ::executorch::runtime::Result<::executorch::aten::Tensor>