()
| 910 | |
| 911 | |
| 912 | def main(): |
| 913 | device = args.device_id |
| 914 | rank = args.rank |
| 915 | set_device(device) |
| 916 | |
| 917 | paddle.set_default_dtype(args.default_dtype) |
| 918 | if args.cache_dtype == "block_wise_fp8": |
| 919 | cache_type = "uint8" |
| 920 | else: |
| 921 | cache_type = args.cache_dtype |
| 922 | |
| 923 | speculative_config = SpeculativeConfig(args.speculative_config) |
| 924 | num_extra_layers = speculative_config.num_extra_cache_layer |
| 925 | key_cache_shape_list = [int(i) for i in args.key_cache_shape.split(",")] |
| 926 | value_cache_shape_list = [] |
| 927 | if args.value_cache_shape: |
| 928 | value_cache_shape_list = [int(i) for i in args.value_cache_shape.split(",")] |
| 929 | total_gpu_blocks = key_cache_shape_list[0] |
| 930 | num_extra_layer_gpu_blocks = int(total_gpu_blocks * speculative_config.num_gpu_block_expand_ratio) |
| 931 | gpu_cache_kvs = {} |
| 932 | gpu_cache_k_tensors = [] |
| 933 | gpu_cache_v_tensors = [] |
| 934 | |
| 935 | logger.info(f"[rank {rank}/{args.mp_num}] Initializing kv cache for all layers.") |
| 936 | for i in range(args.num_layers + num_extra_layers): |
| 937 | num_gpu_blocks = total_gpu_blocks if i < args.num_layers else num_extra_layer_gpu_blocks |
| 938 | key_cache_shape = [ |
| 939 | num_gpu_blocks, |
| 940 | key_cache_shape_list[1], |
| 941 | key_cache_shape_list[2], |
| 942 | key_cache_shape_list[3], |
| 943 | ] |
| 944 | value_cache_shape = [] |
| 945 | if value_cache_shape_list: |
| 946 | value_cache_shape = [ |
| 947 | num_gpu_blocks, |
| 948 | value_cache_shape_list[1], |
| 949 | value_cache_shape_list[2], |
| 950 | value_cache_shape_list[3], |
| 951 | ] |
| 952 | logger.info( |
| 953 | f"[rank {rank}/{args.mp_num}] ..creating kv cache for layer {i}: {key_cache_shape} {value_cache_shape}" |
| 954 | ) |
| 955 | |
| 956 | gpu_cache_kvs[f"key_caches_{i}_rank{rank}_device{device}"] = paddle.full( |
| 957 | shape=key_cache_shape, |
| 958 | fill_value=0, |
| 959 | dtype=cache_type, |
| 960 | ) |
| 961 | gpu_cache_k_tensors.append(gpu_cache_kvs[f"key_caches_{i}_rank{rank}_device{device}"]) |
| 962 | set_data_ipc( |
| 963 | gpu_cache_kvs[f"key_caches_{i}_rank{rank}_device{device}"], |
| 964 | f"key_caches_{i}_rank{rank}.device{device}", |
| 965 | ) |
| 966 | if args.cache_dtype == "block_wise_fp8": |
| 967 | gpu_cache_kvs[f"key_cache_scales_{i}_rank{rank}_device{device}"] = paddle.full( |
| 968 | shape=[num_gpu_blocks, key_cache_shape[1], key_cache_shape[2]], |
| 969 | fill_value=0, |
no test coverage detected