Home
last modified time | relevance | path

Searched full:max_batch_size (Results 1 – 25 of 131) sorted by relevance

123456

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/batching_util/
H A Dadaptive_shared_batch_scheduler.h146 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 Dbatch_resource_base.cc172 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 Dserial_device_batch_scheduler.h108 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 Dbatch_resource_base.h150 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 Dbasic_batch_scheduler.h183 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 Dadaptive_shared_batch_scheduler_test.cc246 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 Drnn.cpp432 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 Dtrt_convert.py279 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 Dquantized_kv_cache.py33 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 Dattention.py34 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 Dsdpa.py316 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 Dmodel.py93 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 Dmixtral_moe_model.py72 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 Dbatch_ops.cc28 // '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 Dbatch_ops_test.py61 [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 Dbatch_ops.py30 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 Dstatic_llama.py223 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 Dconvert_graph.cc151 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 Dllama_transformer.py91 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 Drnn.py539 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 Dunion_find.cc85 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 Dcombined_nms_test.py114 """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 DPackedSequence.cpp147 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 Drnn.py658 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 Druntime_fallback_batch_tf_opkernels.cc115 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")

123456