syntax = "proto3";
option optimize_for = SPEED;

package model_execute_data;

enum ExecuteType {
  MODEL_INIT = 0; // Initialize workers via IPC
  MODEL_INFER = 1;
  MODEL_FINALIZE = 2;
  MODEL_INFER_SECOND = 3;
  PD_LINK = 4;
  KV_TRANSFER = 5;
  TEXT_GENERATOR_CLEANUP = 6;
  LORA_OPERATION = 7;
  REMOTE_MODEL_INIT = 8; // Master node initializes slave nodes via gRPC
  EOS_CLEANUP = 9;
  PAUSE_COMMAND_EXEC = 10;
  CLEAR_COMMAND_EXEC = 11;
  RECOVER_COMMAND_EXEC = 12;
  START_COMMAND_EXEC = 13;
  PD_LINK_STATUS_QUERY = 14;
  EXECUTE_ERROR = 15;
  PAUSE_COMMAND_EXEC_ROCE = 16;
}

enum ForwardType {
  PREFILL = 0;
  DECODE = 1;
  EXTEND = 2;
  MIXED = 3;
  CLEAR = 4;
  DUMMY = 5;
}

// definition of these enum values: src/include/request_response/pd_role.h
enum PDOperationType {
  ASSIGN_ROLE = 0;
  UNLINK = 1;  // 废弃
  TRANSFER_KV = 2;
  DELETE_HOST_KV = 3;
}

enum PDRole {
  UNKNOWN_ROLE = 0;
  PREFILL_ROLE = 1;
  DECODE_ROLE = 2;
  STANDARD_ROLE = 3;
}

// definition of these error codes: mindie_llm/model_wrapper/utils/error.py
enum PDErrorCode {
  SUCCESS = 0;
  PD_UNLINK_ERROR = 2000;
  PD_LINK_ERROR = 2001;
  PD_SWITCH_ROLE_ERROR = 2002;
  PD_PULL_KV_ERROR = 2003;
  PD_PULL_KV_UNKNOWN_ERROR = 2004;
  PD_UNKNOWN_ERROR = 2005;
  TRANSFER_KV_ERROR = 2006;
  CLEAR_KV_ERROR = 2007;
}

enum LoraOperationType {
  LOAD = 0;
  UNLOAD = 1;
}

enum LoraOperationStatus {
  LORA_CMD_SUCCESS = 0;
  SLOTS_FULL = 1;
  DUPLICATED_LORA_ID = 2;
  INVALID_LORA_ID = 3;
  INVALID_LORA_PATH = 4;
  INVALID_LORA_RANK = 5;
  UNSUPPORT_CMD = 6;
}

message LoraOperationRequest {
  string master_model = 1;
  string lora_name = 2;
  string lora_path = 3;
  LoraOperationType lora_op_type = 4;
}

message LoraOperationResponse {
  string lora_name = 1;
  LoraOperationStatus lora_op_status = 2;
}

message SequenceData {
  repeated int64 prompt_token_ids = 1;
  repeated int64 output_token_ids = 2;
}

message SamplingParams {
  optional float repetition_penalty = 1;
  optional float frequency_penalty = 2;
  optional float presence_penalty = 3;
  optional float temperature = 4;
  optional int32 top_k = 5;
  optional float top_p = 6;
  optional float top_logprobs = 7;
  uint32 n = 8;
  optional int32 seed = 9;
  int32 max_output_len = 10;
  bool use_beam_search = 11;
  optional bool logprobs = 12;
  uint32 best_of = 13;
}

message IntList {
  repeated int32 values = 1; // 模拟 List<int>
}

message PromptAdapterRequest {
  string prompt_adapter_name = 1;
  int32 prompt_adapter_id = 2;
  string prompt_adapter_local_path = 3;
  int32 prompt_adapter_num_virtual_tokens = 4;
}

message DPBatchSeqLen {
  repeated int32 seq_lens = 1; // 分布式多DP的P节点，1个DP的batch的seq的长度。
}

message LwdCloudMetadata {
  int32 lwd_cloud_sp_rank_id = 1; // decode本轮调度使用的rank
  int32 lwd_cloud_append_block_rank_id = 2;
  repeated int32 lwd_cloud_sp_rank_token_num = 3; // 每个spcp rank的prompt的token 长度
  repeated int32 lwd_cloud_sp_rank_block_num = 4; // 每个spcp rank的block的个数
  bytes lwd_cloud_block_tables = 5;
}

message SequenceGroupMetadata {
  string request_id = 1;
  bool is_prompt = 2;
  bytes seqIds = 3;
  optional bytes prompt_lens = 4;
  optional bytes prompt_token_ids = 5;
  optional bool do_sample = 6;
  optional SamplingParams sampling_params = 7;
  // Block tables for KV cache. Each element is a serialized contiguous BlockId array (bytes),
  // corresponding to one block manager (manager0, manager1, ...).
  repeated bytes block_tables = 8;
  bytes computed_block_lens = 9;
  optional int32 token_chunk_size = 10;
  optional int32 num_speculative_tokens = 11;
  repeated string stop = 12;
  repeated int32 stop_token_ids = 13;
  optional bool include_stop_str_in_output = 14;
  optional bool ignore_eos = 15;
  optional bool skip_special_tokens = 16;
  repeated bool is_req_prefill = 17;
  repeated bool is_req_last_chunk = 18;
  repeated int32 split_start_pos = 19;
  repeated int32 split_end_pos = 20;
  optional string lora_id = 21;
  repeated int32 sp_rank_token_num = 22; // 每个spcp rank的prompt的token 长度
  repeated int32 sp_rank_block_num = 23; // 每个spcp rank的block的个数
  int32 sp_rank_id = 24; // decode本轮调度使用的rank
  repeated int64 reserved_seq_ids = 25;
  int32 dp_rank_id = 26;
  bytes remote_computed_block_lens = 27;
  optional bool is_append_block = 28;
  repeated int32 prefill_block_rank_id = 29;
  string server_id = 30;
  int32 append_block_rank_id = 31;
  int32 request_gap = 32; // 边云动态切块新增
  optional LwdCloudMetadata lwd_cloud_metadata = 33; //边云特性spcp云侧block管理
  optional string response_format = 34; // JSON 结构化输出约束
  repeated int64 predicted_token_ids = 35; // PD分离/重计算场景下已生成的output token IDs，用于同步grammar等状态
}

message IntPair {
  int32 num1 = 1;
  int32 num2 = 2;
}

// The following messages are defined for requests in normal scenario.
message ExecuteModelRequest {
  repeated SequenceGroupMetadata seq_group_metadata_list = 1;
  repeated IntPair blocks_to_swap_in = 2; // CPU->NPU block number
  repeated IntPair blocks_to_swap_out = 3;
  repeated IntPair blocks_to_copy = 4;
  int32 virtual_engine = 5;
  int32 num_lookahead_slots = 6;
  int32 running_queue_size = 7;
  int32 num_steps = 8;
  optional int32 spec_step_idx = 9;
  repeated string finished_requests_ids = 10;
  repeated int32 last_sampled_token_ids = 11;
  ForwardType forward_type = 12;
  repeated DPBatchSeqLen all_dp_batches_seq_lens = 13; // 分布式多DP的P节点，所有DP的batch的seq的长度。
  int32 wait_queue_len = 14; // 边云动态切块新增
}

message HostInfo {
  string host_ip = 1;
  string cluster_id = 2;
  optional int64 super_pod_id = 3;
}

message DeviceInfo {
  string device_ip = 1;
  int32 physical_id = 2;
  optional int64 super_device_id = 3;
}

message RemoteInfo {
  repeated HostInfo host_info = 1;
  repeated DeviceInfo device_info = 2;
}

// the following messages are defined for requests in PD disaggregation scenario.

// PD disaggregation request: establish/remove linking between P and D instances
message PDLinkRequest {
  message PDLinkInfo {
    PDRole pd_role = 1;
    bool change_role = 2;
    int64 link_num = 3;
    int64 unlink_num = 4;
    repeated RemoteInfo link_info = 5;
    repeated RemoteInfo unlink_info = 6;
    map<uint32, int64> instance2sp = 7; // key: pd instance id, value: sp size
    map<uint32, int64> instance2cp = 8; // key: pd instance id, value: cp size
    int64 host_ip_num = 9;
    int64 super_id_num = 10;
    int64 contains_dp_instance_ids = 11;
  }
  repeated PDLinkInfo pd_link_info = 1;
}

// PD status query request: query the status of P and D instances
message PDLinkStatusRequest {
  optional bool placeholder = 1;
}

// PD disaggregation: publish KV Cache (request message for P instance)
message PublishKVRequest {
  message PublishKVInfo {
    bytes src_block_tables = 1;
    string cluster_id = 2;
  }
  repeated PublishKVInfo publish_kv_infos = 1;
}

// PD disaggregation: pull KV Cache (request message for D instance)
message PullKVRequest {
  message PullKVInfo {
    repeated bytes src_block_tables = 1;
    repeated bytes dst_block_tables = 2;
    string cluster_id = 3;
    SequenceGroupMetadata seq_group_metadata = 4;
  }
  repeated PullKVInfo pull_kv_infos = 1;
}

message TGCleanupRequest { // TG: TextGenerator
  repeated int64 seq_ids = 1;
}

message RemoteModelInitRequest {
  map<string, string> pd_info = 1;
}

message RecoverCommandRequest {
  string command = 1;
}

message ExecuteRequest {
  ExecuteType execute_type = 1;
  map<string, string> config = 2;
  oneof request {
    ExecuteModelRequest execute_model_request = 3;
    PDLinkRequest pd_link_request = 4;
    PublishKVRequest publish_kv_request = 5;// 废弃
    PullKVRequest pull_kv_request = 6;
    TGCleanupRequest text_generator_cleanup_request = 7;
    LoraOperationRequest lora_operation_request = 8;
    RemoteModelInitRequest remote_model_init_request = 9;
    RecoverCommandRequest recover_command_request = 10;
    PDLinkStatusRequest pd_link_status_request = 11;
  }
}

// The following messages are defined for response.
message SequenceOutput {
  int64 seq_id = 1;
  int64 parent_seq_id = 2;
  repeated int64 output_token = 3 [packed=true];
  repeated float logprob = 4 [packed=true];
  int32 finish_reason = 5;
  int32 truncation_index = 6;
  int64 num_speculative_tokens = 7;
  float cumulative_logprobs = 8;
  repeated int64 top_token_ids = 9;
  repeated float top_logprobs = 10;
  int32 num_parallel_tokens = 11;
}

message SequenceGroupLogprob {
  message Logprob {
    int32 token_id = 1;
    float logprob = 2;
    optional int32 rank = 3;
    optional string decoded_token = 4;
  }
  repeated Logprob logprobs = 1;
}

//{token_id -> logprob} per each sequence group. None if the corresponding
message PromptLogprobs { repeated SequenceGroupLogprob prompt_logprobs = 1; }

message CompletionSequenceGroupOutput {
  repeated SequenceOutput samples = 1;
  optional PromptLogprobs prompt_logprobs = 2;
  optional int32 step_index = 3;
}

message ExecuteModelResponse {
  repeated CompletionSequenceGroupOutput outputs = 1;
  optional float model_forward_time = 2;
  optional float model_execute_time = 3;
  optional bool layerwise_is_prefill = 4;
  optional string err_msg = 5;
}

// PD disaggregation response: failed\success\running\waiting links between P and D instances
message PDLinkStatusResponse {
  message FailedLinkInfo {
    string cluster_id = 1;
    PDErrorCode pd_error_code = 2;
  }
  repeated FailedLinkInfo failed_link_info = 1;
  repeated string success_link_info = 2;
  repeated string running_link_info = 3;
  repeated string waiting_link_info = 4;
}

// PD disaggregation: publish KV Cache (response message for P instance)
message PublishKVResponse {
  message PublishKVResult{
    repeated int64 src_block_tables = 1;
    string cluster_id = 2;
  }
  repeated PublishKVResult publish_kv_results = 1;
}

// PD disaggregation: pull KV Cache (response message for D instance)
message PullKVResponse {
  message PullKVResult {
    string request_id = 1;
    PDErrorCode pd_error_code = 2;
  }
  repeated PullKVResult pull_kv_results = 1;
}

// KV cache description for multi block managers (e.g. different block sizes / compression ratios).
message KVCacheDesc {
  int32 npu_block_num = 1;
  int32 block_size = 2;
  uint32 compression_ratio = 3;
  int32 cache_type = 4;
}

message RemoteModelInitResults {
  int32 cpu_block_num = 1;
  int32 max_position_embeddings = 2;
  repeated KVCacheDesc kv_cache_descs = 3;
}

message RecoverCommandResponse {
  int32 npu_device_id = 1;
  int32 command_result = 2;
  string error_msg = 3;
}

message ExecuteResponse {
  message InitResults {
    map<string, string> init_result_map = 1;
    // Multi KV cache descriptors for creating multiple block managers on upper layer.
    // Backward compatible: old versions ignore this field.
    repeated KVCacheDesc kv_cache_descs = 2;
  }
  int32 msg_type = 1;
  int32 rank = 2;
  int32 status = 3;
  oneof result {
    // AsyncExecuteModel response.
    ExecuteModelResponse execute_model_response = 4;
    // ExecutorInstanceInit response.
    InitResults init_results = 5;
    // SetPDLink response
    PDLinkStatusResponse pd_link_status_response = 6;
    // ExecuteKVTransfer response
    PublishKVResponse publish_kv_response = 7;
    PullKVResponse pull_kv_response = 8;
    // Lora Load or Unload response
    LoraOperationResponse lora_operation_response = 9;
    // Model Init Results from Slave to Master
    RemoteModelInitResults remote_model_init_results = 10;
    RecoverCommandResponse recover_command_response = 11;
  }
}

message RegisterRequestMsg {
  string slave_ip = 1;
}

message NpuUtilizationReport {
  uint32 max_aicore_utilization_percent = 1;
}

message MasterToSlaveMsg{
  int32 source_dp_rank = 1;
  int32 target_dp_rank = 2;
  ExecuteRequest execute_request = 3;
}

message SlaveToMasterMsg{
  int32 source_dp_rank = 1;
  int32 target_dp_rank = 2;
  oneof result {
    RegisterRequestMsg register_request = 3;
    ExecuteResponse execute_response = 4;
    NpuUtilizationReport npu_util_report = 5;
  }
}

// GRPC service in master node.
service MasterService {
  rpc RegisterAndCommunicate(stream SlaveToMasterMsg) returns (stream MasterToSlaveMsg){}
}