Stacked cache mixtral. #155

wang2yn84 · 2024-07-20T03:00:00Z

No description provided.

… ring buffer support then fix the mask. Int8 updates also included but not tested.

…ing.

…o the end of existing cache attention.

…od enough.

… into jit function.

…quant version

…he. Fix the kv cache quantization test.

…im for quantization from 1,3 to -3,-1 to make it more robust;

…d for performance validation.

…ew cache.

…on. Refactor to use only 1 flash attention kernel. Changes the modified ring buffer ragged attention kernel with quantization, layer, etc.

…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.

wang2yn84 added 30 commits July 20, 2024 02:52

Almost working except mask, need to rebase to main to pick up the the…

b27f109

… 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 fail…

9fe9d08

…ing.

Adds lazy_cache_update and restructure the cache flags.

d7d6871

Disable all the prints. Fix create engine.

3e84343

Fix typos and minor errors.

6885b12

Fixes create engine.

c65744d

Adds new_cache_stacked and fixes cache update.

2094785

Fix cache update when new_cach_stacked is False.

2603e4f

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

8910fdf

Updates the exportable model to return cache.

d903e54

Removed the fori loop in cache finalize. Moves the cache.finalize() t…

65c5197

…o the end of existing cache attention.

Try to use shard_map for cache update.

9b8bf50

Fix update single cache line in cache.finalize()

ebb2fb8

Adds int8 support.

88b348a

Int8 left aligned lazy cache update working, performance still not go…

775afbe

…od enough.

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

151372e

Put original ragged attention back.

a59a700

Add the original ragged attention kernel.

9575fd3

Fixes the bf16/int8 cache stack.

cf79be5

Fix int8 stacked cache insertion in engine and finalization.

bbe8d90

Fixes int8 with lazy cache update.

e6b0cb9

Updates the int8 test.

8329d50

Fix the int8 ragged attention output sharding.

72c11c0

Fix group query attention broadcasting issue.

15e1387

Fix shard map input issue. Variables not listed as inputs are freezed…

c7b248a

… into jit function.

Fix the flash attention mask shape; Fix the update single cache line …

46791dc

…quant version

Adds the kv cache test.

51a5f0a

Replace quantized cache "pos" with "input_pos" to align with bf16 cac…

0f0deab

…he. Fix the kv cache quantization test.

Fix prefill cache insertion issue for stacked cache; Changes reduce d…

c5335b0

…im for quantization from 1,3 to -3,-1 to make it more robust;

Adds lazy cache update with generate cache stacked new cache unstacke…

53bc76a

…d for performance validation.

wang2yn84 and others added 21 commits July 20, 2024 02:52

Fix the shard map sharding for stacked generate cache and unstacked n…

a92191e

…ew cache.

Using Jax API to slicing instead of Pytorch index slicing.

6e1b35c

Adds stacked cache support in ragged attention reference kernel.

c2c3103

Adds stacked cache support for the modified ragged kernel.

08b63aa

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

86e0c86

Remove testing temp output files.

3e32dcb

Fix the llama 70b output accuracy resulting from gqa.

c52dd28

Fixes the attention output slicing issue when not using flash attenti…

90655d3

…on. 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

b28d3c1

Fix tests; Fix lint issues;

2dffb49

Fix the interactive script.

5855b3d

Add mlperf benchmark scripts in-tree. (#148)

dc0921e

Fix lint errors.

41f59a1

Fix errors.

d78a9bb

Fix the comments.

5035ce9

Fix based on comments; Fix all the unit tests.

e36833d

Fix the remaining pylint errors.

d263ff5

Default ring buffer back to true so that all the test_run_server and …

fadc777

…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.

703d71f

Fix run_offline script.

88f9ac8

Fix lint errors.

a5e47f8

wang2yn84 requested review from qihqi, lsy323 and FanhaiLu1 July 20, 2024 03:00

qihqi merged commit 2f77223 into mlperf-mixtral Jul 20, 2024
4 checks passed

qihqi deleted the stacked-cache-mixtral branch July 20, 2024 03:37

wang2yn84 removed request for qihqi, lsy323 and FanhaiLu1 July 20, 2024 17:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Stacked cache mixtral. #155

Stacked cache mixtral. #155

Uh oh!

wang2yn84 commented Jul 20, 2024

Uh oh!

Uh oh!

Uh oh!

Stacked cache mixtral. #155

Stacked cache mixtral. #155

Uh oh!

Conversation

wang2yn84 commented Jul 20, 2024

Uh oh!

Uh oh!

Uh oh!