changes to support jetstream

GoogleCloudPlatform · Sep 3, 2024 · 7983364 · 7983364
1 parent e7278e8
commit 7983364
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks/benchmark/tools/profile-generator/container/benchmark_serving.py b/benchmarks/benchmark/tools/profile-generator/container/benchmark_serving.py
@@ -181,7 +181,7 @@ async def send_request(
   elif backend == "jetstream":
     pload = {
         "prompt": prompt,
-        "max_tokens": 1,
+        "max_tokens": output_len,
     }
   else:
     raise ValueError(f"Unknown backend: {backend}")