/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/batching_util/ |
H A D | adaptive_shared_batch_scheduler.h | 146 int max_batch_size = 1000; member 150 // If specified, it should be larger than or equal to 'max_batch_size'. 168 int max_batch_size, 319 size_t max_task_size() const override { return options_.max_batch_size; } in max_task_size() 446 if (options.max_batch_size <= 0) { in AddQueue() 447 return errors::InvalidArgument("max_batch_size must be positive; was ", in AddQueue() 448 options.max_batch_size); in AddQueue() 456 if (options.max_input_task_size.value() < options.max_batch_size) { in AddQueue() 458 "max_input_task_size must be larger than or equal to max_batch_size;" in AddQueue() 460 options.max_input_task_size.value(), " and max_batch_size as ", in AddQueue() [all …]
|
H A D | batch_resource_base.cc | 172 void RecordBatchParamMaxBatchSize(int64_t max_batch_size, in RecordBatchParamMaxBatchSize() argument 176 "/tensorflow/serving/batching/max_batch_size", in RecordBatchParamMaxBatchSize() 178 cell->GetCell(model_name, op_name)->Set(max_batch_size); in RecordBatchParamMaxBatchSize() 277 RecordBatchParamMaxBatchSize(adaptive_batcher_queue_options_.max_batch_size, in RegisterInput() 331 int32_t num_batch_threads, int32_t max_batch_size, in GetBatcherQueueOptions() argument 336 batcher_queue_options.input_batch_size_limit = max_batch_size; in GetBatcherQueueOptions() 344 int open_batch_remaining_slot, int max_batch_size, in GetBatcherQueueOptions() 347 max_batch_size, output_tasks); in GetBatcherQueueOptions() 351 batcher_queue_options.max_execution_batch_size = max_batch_size; in GetBatcherQueueOptions() 363 int32_t max_batch_size, int32_t batch_timeout_micros, in GetAdaptiveBatcherQueueOptions() argument [all …]
|
H A D | serial_device_batch_scheduler.h | 108 int max_batch_size = 1000; member 226 size_t max_task_size() const override { return options_.max_batch_size; } in max_task_size() 335 if (options.max_batch_size <= 0) { in AddQueue() 336 return errors::InvalidArgument("max_batch_size must be positive; was ", in AddQueue() 337 options.max_batch_size); in AddQueue() 490 if (size > options_.max_batch_size) { in Schedule() 493 options_.max_batch_size); in Schedule() 499 current_batch_->size() + size > options_.max_batch_size) { in Schedule() 540 current_batch_ ? options_.max_batch_size - current_batch_->size() : 0; in SchedulingCapacity() 543 return spare_batches * options_.max_batch_size + current_batch_capacity; in SchedulingCapacity()
|
H A D | batch_resource_base.h | 150 int32_t num_batch_threads, int32_t max_batch_size, 156 int32_t max_batch_size, int32_t batch_timeout_micros, 164 // 2) max_batch_size 169 // [open_batch_remaining_slot, max_batch_size, max_batch_size, ..., 177 int max_batch_size,
|
H A D | basic_batch_scheduler.h | 183 int max_batch_size = 1000; member 188 // even if the batch's size is below 'max_batch_size'. 219 // 2.`max_batch_size` specifies the limit of input and 239 // `max_batch_size`: Maximum size of each batch. 263 // less than or equal to `max_batch_size`. 265 // equal to `max_batch_size`. 325 options.max_batch_size; in Create()
|
H A D | adaptive_shared_batch_scheduler_test.cc | 246 queue_options.max_batch_size = 1000; in TEST() 248 queue_options.max_batch_size = 100; in TEST() 321 queue_options.max_batch_size = 1000; in TEST() 323 queue_options.max_batch_size = 100; in TEST() 436 queue_options.max_batch_size = 100; in TEST() 458 queue_options.max_batch_size = 100; in TEST() 499 queue_options.max_batch_size = 100; in TEST()
|
/aosp_15_r20/external/pytorch/torch/csrc/api/src/nn/modules/ |
H A D | rnn.cpp | 432 int64_t max_batch_size, in forward_helper() argument 438 max_batch_size, in forward_helper() 519 auto max_batch_size = in forward() local 525 input, batch_sizes, sorted_indices, max_batch_size, std::move(hx)); in forward() 538 auto max_batch_size = batch_sizes[0].item<int64_t>(); in forward_with_packed_input() local 541 input, batch_sizes, sorted_indices, max_batch_size, std::move(hx)); in forward_with_packed_input() 610 int64_t max_batch_size, in forward_helper() argument 619 max_batch_size, in forward_helper() 624 max_batch_size, in forward_helper() 670 auto max_batch_size = options.batch_first() ? input.size(0) : input.size(1); in forward() local [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/python/compiler/tensorrt/ |
H A D | trt_convert.py | 279 max_batch_size=None, argument 289 max_batch_size: maximum batch size for static engines. 308 if (is_dynamic_op is None or is_dynamic_op) and max_batch_size is not None: 309 raise ValueError("max_batch_size has to be None for TF2" 312 max_batch_size, int): 314 "max_batch_size has to be an integer for is_dynamic_op==False in TF1") 354 if max_batch_size is not None: 355 optimizer.parameter_map["max_batch_size"].i = max_batch_size 376 max_batch_size=None, argument 380 max_batch_size, is_v2, [all …]
|
/aosp_15_r20/external/executorch/examples/models/llama/source_transformation/ |
H A D | quantized_kv_cache.py | 33 max_batch_size, argument 57 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 58 scale_shape = (max_batch_size, n_heads, max_seq_length, 1) 60 cache_shape = (max_batch_size, max_seq_length, n_heads, head_dim) 61 scale_shape = (max_batch_size, max_seq_length, n_heads, 1) 197 max_batch_size, n_heads, max_seq_length, head_dim = cache_shape 199 max_batch_size, max_seq_length, n_heads, head_dim = cache_shape 201 max_batch_size,
|
H A D | attention.py | 34 max_batch_size: int, 43 cache_shape = (max_batch_size, max_seq_length, head_dim) 81 max_batch_size: int, 93 max_batch_size, max_seq_length, n_heads // n_rep, head_dim 136 self.max_batch_size = attention_mha.max_batch_size 140 self.max_batch_size,
|
H A D | sdpa.py | 316 max_batch_size: int, 324 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 326 self.max_batch_size = max_batch_size 351 child.max_batch_size, 366 max_batch_size: int, 373 cache_shape = (max_batch_size, max_seq_length, n_heads, head_dim) 403 child.max_batch_size,
|
/aosp_15_r20/external/pytorch/benchmarks/gpt_fast/ |
H A D | model.py | 93 self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16 argument 96 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 126 self.max_batch_size = -1 129 def setup_caches(self, max_batch_size, max_seq_length): argument 132 and self.max_batch_size >= max_batch_size 138 self.max_batch_size = max_batch_size 141 max_batch_size, max_seq_length, self.config.n_local_heads, head_dim
|
H A D | mixtral_moe_model.py | 72 self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16 argument 75 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 105 self.max_batch_size = -1 108 def setup_caches(self, max_batch_size, max_seq_length): argument 111 and self.max_batch_size >= max_batch_size 117 self.max_batch_size = max_batch_size 120 max_batch_size, max_seq_length, self.config.n_local_heads, head_dim
|
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/ |
H A D | batch_ops.cc | 28 // 'max_batch_size' denotes the maximum batch size acceptable, i.e., inputs 30 // By default, 'max_batch_size' must be equal to max value of 33 // 'max_batch_size' can be greater than or equal to max value of 35 // 1) input with size > 'max_batch_size' is still invalidated. 37 // a) size <= 'max_batch_size' 41 .Attr("max_batch_size: int") 68 .Attr("max_batch_size: int")
|
/aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
H A D | batch_ops_test.py | 61 [inp], num_batch_threads=1, max_batch_size=2, 103 [inp], num_batch_threads=1, max_batch_size=10, 138 max_batch_size=2, 177 [inp0, inp1], num_batch_threads=1, max_batch_size=2, 192 [inp], num_batch_threads=1, max_batch_size=10, 381 max_batch_size=10, 413 max_batch_size=10, 448 max_batch_size=10, 476 # max('allowed_batch_sizes') <= 'max_batch_size'. 478 max_batch_size=5, [all …]
|
H A D | batch_ops.py | 30 max_batch_size, argument 62 max_batch_size: Batch sizes will never be bigger than this. 68 monotonically, and the final entry must equal max_batch_size. 75 adding. 2.`max_batch_size` specifies the limit of input and 106 max_batch_size=max_batch_size,
|
/aosp_15_r20/external/executorch/examples/qualcomm/oss_scripts/llama2/model/ |
H A D | static_llama.py | 223 self.max_batch_size = config.max_batch_size 286 self.vocab_size, (self.max_batch_size, 1), dtype=torch.int32 288 pos_ids = torch.zeros((self.max_batch_size, 1), dtype=torch.int32) 290 atten_mask = torch.full((self.max_batch_size, self.max_seq_len), -255.0) 297 self.max_batch_size, 304 self.max_batch_size, 324 "get_max_batch_size": self.max_batch_size,
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ |
H A D | convert_graph.cc | 151 info->max_batch_size = segment_property.BatchSize().GetOptionalMaxBatchSize(); in GetEngineInfo() 461 int max_batch_size = info.max_batch_size.has_value() in CreateTRTNode() local 462 ? info.max_batch_size.value() in CreateTRTNode() 466 TF_RETURN_IF_ERROR(CreateStaticEngine(params, info, max_batch_size, in CreateTRTNode() 500 .Attr("max_batch_size", max_batch_size) in CreateTRTNode() 693 const EngineInfo& info, int max_batch_size, in CreateStaticEngine() argument 720 max_batch_size, info.max_workspace_size_bytes, input_shapes, trt_logger, in CreateStaticEngine() 779 segment_options.maximum_batch_size = params.max_batch_size; in ConvertGraph() 844 if (!curr_engine.max_batch_size.has_value()) { in ConvertGraph() 845 curr_engine.max_batch_size = params.max_batch_size; in ConvertGraph() [all …]
|
/aosp_15_r20/external/executorch/examples/models/llama/ |
H A D | llama_transformer.py | 91 max_batch_size: int = 32 variable in ModelArgs 149 max_batch_size: int, 161 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 163 cache_shape = (max_batch_size, max_seq_length, n_heads, head_dim) 165 self.max_batch_size = max_batch_size 276 self.max_batch_size = args.max_batch_size 298 args.max_batch_size,
|
/aosp_15_r20/external/pytorch/torch/ao/nn/quantized/dynamic/modules/ |
H A D | rnn.py | 539 max_batch_size: int, 546 max_batch_size, 598 max_batch_size = input.size(0) if self.batch_first else input.size(1) 603 input, hx, batch_sizes, max_batch_size, sorted_indices 613 max_batch_size = int(batch_sizes[0]) 616 input_, hx, batch_sizes, max_batch_size, sorted_indices 833 max_batch_size: int, 840 max_batch_size, 888 max_batch_size = input.size(0) if self.batch_first else input.size(1) 893 input, hx, batch_sizes, max_batch_size, sorted_indices [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/segment/ |
H A D | union_find.cc | 85 ClusterBatchSize& ClusterBatchSize::SetMaxBatchSize(int max_batch_size) { in SetMaxBatchSize() argument 86 SetBatchSize(static_cast<std::optional<int>>(max_batch_size)); in SetMaxBatchSize() 91 const std::optional<int>& max_batch_size) { in SetMaxBatchSize() argument 92 max_batch_size_ = MergeCompatible<int>(max_batch_size_, max_batch_size); in SetMaxBatchSize() 122 absl::StrAppendFormat(&s, ", max_batch_size="); in ToString()
|
/aosp_15_r20/external/tensorflow/tensorflow/python/compiler/tensorrt/test/ |
H A D | combined_nms_test.py | 114 """Returns the max_batch_size that the converter should use for tests.""" 118 # Build the engine with the allowed max_batch_size less than the actual 119 # max_batch_size, to fore the runtime to execute the native segment. This 127 # max_batch_size is only useful for selecting static engines. As such, 215 """Returns the max_batch_size that the converter should use for tests."""
|
/aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
H A D | PackedSequence.cpp | 147 int64_t max_batch_size = batch_sizes[0]; in _pad_packed_sequence() local 158 std::vector<int64_t> output_size; // == [max_seq_length, max_batch_size, *var_data.size()[1:]] in _pad_packed_sequence() 162 output_size.push_back(max_batch_size); in _pad_packed_sequence() 171 at::Tensor lengths_t = at::empty(max_batch_size, batch_sizes_t.options()); in _pad_packed_sequence() 172 int64_t * lengths = lengths_t.mutable_data_ptr<int64_t>() + max_batch_size - 1; in _pad_packed_sequence() 174 int64_t prev_batch_size = max_batch_size; in _pad_packed_sequence()
|
/aosp_15_r20/external/pytorch/torch/nn/modules/ |
H A D | rnn.py | 658 max_batch_size = batch_sizes[0] 659 … # script() is unhappy when max_batch_size is different type in cond branches, so we duplicate 663 max_batch_size, 693 max_batch_size = input.size(0) if self.batch_first else input.size(1) 699 max_batch_size, 1051 max_batch_size = batch_sizes[0] 1055 max_batch_size, 1062 max_batch_size, 1081 max_batch_size = input.size(0) if self.batch_first else input.size(1) 1087 max_batch_size, [all …]
|
/aosp_15_r20/external/tensorflow/tensorflow/core/runtime_fallback/runtime/ |
H A D | runtime_fallback_batch_tf_opkernels.cc | 115 static Status Create(int32_t num_batch_threads, int32_t max_batch_size, in Create() argument 130 GetBatcherQueueOptions(num_batch_threads, max_batch_size, in Create() 140 int32_t max_batch_size, int32_t batch_timeout_micros, in Create() argument 152 max_batch_size, batch_timeout_micros, max_enqueued_batches, in Create() 282 OP_REQUIRES_OK(c, c->GetAttr("max_batch_size", &max_batch_size_)); in BatchFunctionFallbackKernel() 419 "final entry in allowed_batch_sizes must equal max_batch_size when " in ValidateAllowedBatchSizes() 634 .Attr("max_batch_size: int")
|