gemma4_31b: keep sampled token on device (int64 output + zero-copy alias)

Eliminate the per-decode-round token D2H->H2D round-trip:
- sampler.py: emit the sampled token as int64 (was float32) so the decode
  method's int64 token output can be aliased directly as the next forward's
  int64 token input (value-preserving: argmax index; token ids < 2^24).
- main.cpp: read_token reads int64; each forward's on-device output token is
  aliased via make_tensor_ptr and fed straight back as the next step's token
  input (prefill->decode and decode->decode). Only the per-round position
  H2D remains.

Measured (int6/gguf, cuda graph OFF, p19/d128): post-load HtoD 261->132
(token H2D removed; ~= decode length); DtoH/DtoD counts unchanged (129),
bytes 4B->8B (token now int64). Greedy output byte-identical to prior export.

Author: Gasoonjia

examples/models/gemma4_31b/main.cpp

@@ -82,7 +82,9 @@ DEFINE_bool(
 
 namespace llm = ::executorch::extension::llm;
 using ::executorch::extension::from_blob;
+using ::executorch::extension::make_tensor_ptr;
 using ::executorch::extension::Module;
+using ::executorch::extension::TensorPtr;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::EValue;
 #ifdef EXECUTORCH_BUILD_CUDA
@@ -91,18 +93,23 @@ using ::executorch::extension::clone_tensor_ptr_to;
 
 using SizesType = executorch::aten::SizesType;
 
-// Read a sampled token ID from a scalar float output (CUDA path).
+// Read a sampled token ID from a scalar int64 output (CUDA path).
+//
+// The model now emits the sampled token as int64 (see sampler.py), matching
+// the decode method's int64 token input so the on-device output buffer can be
+// aliased directly as the next step's input. We still copy the 8-byte scalar
+// back to the host here for EOS detection and detokenization.
 static uint64_t read_token(const executorch::aten::Tensor& output) {
   const void* ptr = output.const_data_ptr();
-  float val = 0.0f;
+  int64_t val = 0;
 
 #ifdef EXECUTORCH_BUILD_CUDA
   cudaPointerAttributes attrs{};
   bool on_device = cudaPointerGetAttributes(&attrs, ptr) == cudaSuccess &&
       attrs.type == cudaMemoryTypeDevice;
   if (on_device) {
     cudaError_t err =
-        cudaMemcpy(&val, ptr, sizeof(float), cudaMemcpyDeviceToHost);
+        cudaMemcpy(&val, ptr, sizeof(int64_t), cudaMemcpyDeviceToHost);
     if (err != cudaSuccess) {
       ET_LOG(
           Error,
@@ -111,13 +118,13 @@ static uint64_t read_token(const executorch::aten::Tensor& output) {
       return 0;
     }
   } else {
-    memcpy(&val, ptr, sizeof(float));
+    memcpy(&val, ptr, sizeof(int64_t));
   }
 #else
-  memcpy(&val, ptr, sizeof(float));
+  memcpy(&val, ptr, sizeof(int64_t));
 #endif
 
-  return static_cast<uint64_t>(llrintf(val));
+  return static_cast<uint64_t>(val);
 }
 
 int main(int argc, char** argv) {
@@ -295,6 +302,12 @@ int main(int argc, char** argv) {
   // ---------------------------------------------------------------
   uint64_t cur_token = 0;
   int64_t prefill_pos = 0;
+#ifdef EXECUTORCH_BUILD_CUDA
+  // Alias of the most recent forward's on-device int64 output token. The last
+  // prefill chunk's output seeds the first decode step (no token H2D); each
+  // decode step then re-aliases its own output for the next step.
+  TensorPtr device_out_token;
+#endif
   while (prefill_pos < num_prompt_tokens) {
     int64_t chunk_len =
         std::min(num_prompt_tokens - prefill_pos, max_prefill_chunk);
@@ -337,7 +350,11 @@ int main(int argc, char** argv) {
     }
 
 #ifdef EXECUTORCH_BUILD_CUDA
-    cur_token = read_token(result.get()[0].toTensor());
+    const auto& out_tensor = result.get()[0].toTensor();
+    cur_token = read_token(out_tensor);
+    // Keep the sampled token on device: alias the output buffer so it feeds
+    // straight into the next forward as the int64 token input (zero copy).
+    device_out_token = make_tensor_ptr(out_tensor);
 #else
     cur_token = static_cast<uint64_t>(
         llm::logits_to_token(result.get()[0].toTensor(), temp_val));
@@ -369,53 +386,48 @@ int main(int argc, char** argv) {
   // Decode loop
   // ---------------------------------------------------------------
   int64_t pos = num_prompt_tokens;
-  std::vector<int64_t> decode_token_data = {static_cast<int64_t>(cur_token)};
   std::vector<int64_t> decode_pos_data = {pos};
-  auto decode_tokens_cpu = from_blob(
-      decode_token_data.data(), {1, 1}, executorch::aten::ScalarType::Long);
   auto decode_pos_cpu = from_blob(
       decode_pos_data.data(), {1}, executorch::aten::ScalarType::Long);
 #ifdef EXECUTORCH_BUILD_CUDA
-  // skip_h2d: keep fixed device-resident input buffers across decode steps
-  // (seeded here with a one-time H2D). Their 8-byte contents are refreshed
-  // each step from the host (see loop), since the sampled id round-trips to
-  // the host for EOS detection.
-  auto decode_tokens = clone_tensor_ptr_to(decode_tokens_cpu, cuda_device);
+  // The token input is the aliased on-device output (device_out_token); only
+  // the position still needs a fixed device buffer refreshed by a per-round
+  // H2D (one int64 / round).
   auto decode_pos = clone_tensor_ptr_to(decode_pos_cpu, cuda_device);
 #else
-  auto decode_tokens = decode_tokens_cpu;
+  // Non-CUDA (MLX) path: keep host token/pos buffers; the backend stages them
+  // and the host samples from the returned logits.
+  std::vector<int64_t> decode_token_data = {static_cast<int64_t>(cur_token)};
+  auto decode_tokens = from_blob(
+      decode_token_data.data(), {1, 1}, executorch::aten::ScalarType::Long);
   auto decode_pos = decode_pos_cpu;
 #endif
 
   uint64_t prev_token = cur_token;
   bool hit_eos = eos_ids.find(cur_token) != eos_ids.end();
   for (int32_t step = 0; step < FLAGS_max_new_tokens && !hit_eos; step++) {
-    decode_token_data[0] = static_cast<int64_t>(cur_token);
     decode_pos_data[0] = pos;
 
 #ifdef EXECUTORCH_BUILD_CUDA
-    // skip_h2d: refresh the device-resident token/pos buffers ourselves (the
-    // backend no longer inserts the H2D copy). The prior step's sampled id
-    // already came back to the host via read_token, so re-upload the 8-byte
-    // token and position into the fixed device buffers.
-    ET_CHECK_MSG(
-        cudaMemcpy(
-            decode_tokens->mutable_data_ptr(),
-            decode_token_data.data(),
-            sizeof(int64_t),
-            cudaMemcpyHostToDevice) == cudaSuccess,
-        "Failed to upload decode token H2D");
+    // Token stays on device (aliased from the previous forward's output); only
+    // the 8-byte position is uploaded each round. No token D2H->H2D round-trip.
     ET_CHECK_MSG(
         cudaMemcpy(
             decode_pos->mutable_data_ptr(),
             decode_pos_data.data(),
             sizeof(int64_t),
             cudaMemcpyHostToDevice) == cudaSuccess,
         "Failed to upload decode position H2D");
+#else
+    decode_token_data[0] = static_cast<int64_t>(cur_token);
 #endif
 
     std::vector<EValue> inputs;
+#ifdef EXECUTORCH_BUILD_CUDA
+    inputs.push_back(EValue(device_out_token));
+#else
     inputs.push_back(EValue(decode_tokens));
+#endif
     inputs.push_back(EValue(decode_pos));
 
 #ifdef EXECUTORCH_BUILD_CUDA
@@ -432,7 +444,10 @@ int main(int argc, char** argv) {
 
     prev_token = cur_token;
 #ifdef EXECUTORCH_BUILD_CUDA
-    cur_token = read_token(result.get()[0].toTensor());
+    const auto& out_tensor = result.get()[0].toTensor();
+    cur_token = read_token(out_tensor);
+    // Alias this step's on-device output token as the next step's token input.
+    device_out_token = make_tensor_ptr(out_tensor);
 #else
     cur_token = static_cast<uint64_t>(
         llm::logits_to_token(result.get()[0].toTensor(), temp_val));

examples/models/gemma4_31b/model.py

@@ -484,7 +484,7 @@ def forward(
             temperature: 1-D float tensor for Gumbel-max sampling.
 
         Returns:
-            (B, 1) sampled token IDs as float.
+            (B, 1) sampled token IDs as int64.
         """
         x = self.embed_tokens(tokens) * self.embed_normalizer
 

examples/models/gemma4_31b/sampler.py

@@ -26,9 +26,12 @@ def sample(
             temperature still works ("near-greedy").
 
     Returns:
-        ``[B, 1]`` float32 token IDs (``argmax(logits/T + gumbel_noise)``).
+        ``[B, 1]`` int64 token IDs (``argmax(logits/T + gumbel_noise)``).
+        Emitting int64 (rather than casting to float) lets the runner alias the
+        on-device output token directly as the next decode step's int64 token
+        input — no D2H/H2D round-trip and no dtype cast.
     """
     logits = logits / temperature.clamp(min=1e-6)
     noise = torch.rand_like(logits)
     gumbel = -torch.log(-torch.log(noise + 1e-20) + 1e-20)
-    return (logits + gumbel).argmax(dim=-1, keepdim=True).float()
+    return (logits + gumbel).argmax(dim=-1, keepdim=True).to(torch.int64)