Stacked cache for MLPerf (#154)

* Almost working except mask, need to rebase to main to pick up the the ring buffer support then fix the mask. Int8 updates also included but not tested.

* Fixed the test_model_impl for llama, but test_llama_e2e is still failing.

* Adds lazy_cache_update and restructure the cache flags.

* Disable all the prints. Fix create engine.

* Fix typos and minor errors.

* Fixes create engine.

* Adds new_cache_stacked and fixes cache update.

* Fix cache update when new_cach_stacked is False.

* Fix the cache manager and make unit tests pass except for 1.

* Updates the exportable model to return cache.

* Removed the fori loop in cache finalize. Moves the cache.finalize() to the end of existing cache attention.

* Try to use shard_map for cache update.

* Fix update single cache line in cache.finalize()

* Adds int8 support.

* Int8 left aligned lazy cache update working, performance still not good enough.

* Fix the stacked cache introduced in the previous couple of commits.

* Put original ragged attention back.

* Add the original ragged attention kernel.

* Fixes the bf16/int8 cache stack.

* Fix int8 stacked cache insertion in engine and finalization.

* Fixes int8 with lazy cache update.

* Updates the int8 test.

* Fix the int8 ragged attention output sharding.

* Fix group query attention broadcasting issue.

* Fix shard map input issue. Variables not listed as inputs are freezed into jit function.

* Fix the flash attention mask shape; Fix the update single cache line quant version

* Adds the kv cache test.

* Replace quantized cache "pos" with "input_pos" to align with bf16 cache. Fix the kv cache quantization test.

* Fix prefill cache insertion issue for stacked cache; Changes reduce dim for quantization from 1,3 to -3,-1 to make it more robust;

* Adds lazy cache update with generate cache stacked new cache unstacked for performance validation.

* Fix the shard map sharding for stacked generate cache and unstacked new cache.

* Using Jax API to slicing instead of Pytorch index slicing.

* Adds stacked cache support in ragged attention reference kernel.

* Adds stacked cache support for the modified ragged kernel.

* Llama2 70b int8 optimization done. Output not correct yet.

* Remove testing temp output files.

* Fix the llama 70b output accuracy resulting from gqa.

* Fixes the attention output slicing issue when not using flash attention. Refactor to use only 1 flash attention kernel. Changes the modified ring buffer ragged attention kernel with quantization, layer, etc.

* Fix the pallas kernel OOB issue

* Fix tests; Fix lint issues;

* Fix the interactive script.

* Fix lint errors.

* Fix errors.

* Fix the comments.

* Fix based on comments; Fix all the unit tests.

* Fix the remaining pylint errors.

* Default ring buffer back to true so that all the test_run_server and run_interactive in CPU mode can work. When we default ring buffer to false, should add additional flags to run_interactive CI to set test mode to true so that pallas kernel can run.

* Fix all the lint errors.

* Remove the deps/JetStream changes.

* Fix merge errors, fix lint errors.

Author: wang2yn84

benchmarks/run_offline.py

@@ -32,7 +32,7 @@
 flags.DEFINE_string("sharegpt_path", "", "path to sharegpt json file")
 
 
-def run_prefill_time(engine, params, decode_state, seqlen):
+def run_prefill_time(engine, params, decode_state, seqlen, profiler_started):
   """Run prefill and measure time."""
   metadata = engine.get_tokenizer()
   tokenizer = engine.build_tokenizer(metadata)
@@ -53,15 +53,20 @@ def run_prefill_time(engine, params, decode_state, seqlen):
   nums = 5
   start = time.perf_counter()
   for i in range(nums):
+    if i == nums - 1 and FLAGS.profiling_prefill and not profiler_started:
+      jax.profiler.start_trace(FLAGS.profiling_output)
+      profiler_started = True
+
     prefill_result, _ = engine.prefill(
         params=params, padded_tokens=tokens, true_length=true_length
     )
     decode_state = engine.insert(
         prefill_result, decode_state, slot=jnp.int32(i)
     )
   jax.block_until_ready(decode_state)
+
   end = time.perf_counter()
-  return (end - start) / nums, decode_state
+  return (end - start) / nums, decode_state, profiler_started
 
 
 MAXTEXT_PREFILL = {
@@ -86,9 +91,10 @@ def main(argv):
   prefill_times = {}
 
   decode_state = engine.init_decode_state()
+  profiler_started = False
   for batch, _ in MAXTEXT_PREFILL.items():
-    runtime, decode_state = run_prefill_time(
-        engine, params, decode_state, batch
+    runtime, decode_state, profiler_started = run_prefill_time(
+        engine, params, decode_state, batch, profiler_started
     )
     prefill_times[batch] = runtime
 
@@ -103,10 +109,12 @@ def main(argv):
 
   profiling_output = FLAGS.profiling_output
   print("======= decode starting ===")
+
   dec_times = []
   for i in range(10):
-    if profiling_output and i == 7:
+    if profiling_output and i == 7 and not profiler_started:
       jax.profiler.start_trace(profiling_output)
+      profiler_started = True
     start = time.perf_counter()
     # pylint: disable-next=all
     decode_state, sampled_tokens = engine.generate(params, decode_state)
@@ -116,7 +124,7 @@ def main(argv):
     dec_times.append(end - start)
     print(i, "decode time", (end - start))
 
-  if profiling_output:
+  if profiler_started:
     jax.profiler.stop_trace()
 
   print("prefill ", prefill_times)