From 3c7c8625bf14308302f020c001c36363199cb4a0 Mon Sep 17 00:00:00 2001 From: eplatero Date: Wed, 29 Jul 2026 14:19:35 -0500 Subject: [PATCH 1/4] Add dynamic batching on QAIC via batch_index KV-slicing Compile one QPC that serves multiple live decode batch sizes b <= B_max by riding the continuous-batching export path: every decode specialization pins the retained KV cache at full_batch_size (B_max) while only the input batch axis varies. This is the only path where the input batch axis (ONNX symbol "batch_size") is decoupled from the retained KV axis ("full_batch_size"), so varying the input batch is a non-retained-axis change the compiler accepts -- unlike the earlier per-batch-KV attempt, which asked one QPC to hold multiple retained-state shapes and was rejected with "inconsistent retained state". Compile side (modeling_auto.py): - batch_size accepts List[int]; a list requires continuous_batching=True and every value must be <= full_batch_size (B_max). - build_decode_specialization gains decode_input_batch_size so decode specs vary the input batch while full_batch_size/kv_cache_batch_size pin the KV. - Reject combining a batch list with CCL (comp_ctx_lengths): the two vary different identifying inputs and the compiler cannot disambiguate them. Runtime (text_generation_inference.py): - Drive a live batch via decode_batch_size; batch_index = arange(b) routes b sequences into b of the B_max KV slots. - generate(execution_batch_size=b), validated against the compiled batches. - Fix _fetch_full_batch_size to handle 3+ allowed shapes; add _fetch_compiled_batch_sizes and _resolve_execution_batch_size. Validated: 206/206 CPU spec tests; single-QPC compile confirmed with qaic-exec; 5/5 on_qaic tests pass on device, including token-id parity of the dynamic-batch QPC at live batch b vs a natively-compiled batch-b QPC. Co-Authored-By: Claude Opus 4.8 Signed-off-by: eplatero --- .../generation/text_generation_inference.py | 162 +++++++++++--- .../transformers/models/modeling_auto.py | 183 ++++++++++------ tests/transformers/spd/test_spd_inference.py | 176 ++++++++++++++++ .../models/test_modeling_auto_cpu.py | 199 ++++++++++++++++++ tests/unit_test/utils/test_auto_model_api.py | 40 ++++ tests/unit_test/utils/test_generation.py | 106 ++++++++++ 6 files changed, 777 insertions(+), 89 deletions(-) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 17c992064c..8e0d3b3dab 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -482,6 +482,12 @@ def __init__( self.full_batch_size = ( full_batch_size if full_batch_size else self._fetch_full_batch_size() ) # Check and fetch full batch size if CB is enabled + # Decode input batch actually driven each step. Defaults to full_batch_size (B_max) so plain + # continuous batching is unchanged; dynamic batching sets it to a smaller compiled batch to + # run a live batch b <= B_max over the fixed-B_max KV cache via a length-b batch_index. + self.decode_batch_size = self.full_batch_size + # Decode input batch sizes the QPC was compiled for (dynamic batching exposes several). + self._compiled_batch_sizes = self._fetch_compiled_batch_sizes() # Initialize the storage variables. self.batch_index = None @@ -517,23 +523,68 @@ def _fetch_full_batch_size( self, ): """ - Fetches the full batch size from the session's bindings or allowed shapes. + Fetches the full batch size (B_max, the retained KV-cache capacity) from the session. + + B_max is the batch dim (axis 0) of the retained-state KV cache, which is pinned to a single + value across every specialization in the QPC. Reading it from a retained-state binding + (rather than the ``batch_index`` input axis) is correct even for dynamic batching, where the + decode *input* batch varies over ``batch_size=[...]`` and ``max(batch_size)`` may be smaller + than the true KV capacity ``full_batch_size``. Returns: - full_batch_size: The full batch size fetched from the session's bindings or allowed shapes. If "batch_index" is not - in the session's binding index map, full_batch_size will be None. + full_batch_size: The full batch size fetched from the session's bindings or allowed shapes. If + neither a "batch_index" input nor a retained-state binding is present (i.e. continuous batching + is not enabled), full_batch_size will be None. """ - full_batch_size = None - if "batch_index" in self._session.binding_index_map: + if "batch_index" not in self._session.binding_index_map: + return None + # Prefer the retained KV-cache batch dim (true B_max, constant across specs). + kv_index = next( + (idx for name, idx in self._session.binding_index_map.items() if is_retained_state_name(name)), + None, + ) + if kv_index is not None: if self._session.allowed_shapes: - full_batch_size, _ = [ - x[self._session.binding_index_map["batch_index"]][1][0] for x in self._session.allowed_shapes - ] - else: - full_batch_size, _ = self._session.bindings[self._session.binding_index_map["batch_index"]].dims + return self._session.allowed_shapes[0][kv_index][1][0] + return self._session.bindings[kv_index].dims[0] + # Fallback: no retained-state binding exposed. The batch_index input axis equals B_max for + # plain continuous batching; take the max across decode specs as a best effort. + if self._session.allowed_shapes: + return max(x[self._session.binding_index_map["batch_index"]][1][0] for x in self._session.allowed_shapes) + full_batch_size, _ = self._session.bindings[self._session.binding_index_map["batch_index"]].dims return full_batch_size + def _fetch_compiled_batch_sizes(self): + """ + Fetches the sorted set of decode *input* batch sizes the QPC was compiled for. + + For a dynamic-batching QPC (``batch_size=[...]`` at compile time) this returns each compiled + decode input batch (e.g. ``[1, 2, 4]``); for an ordinary QPC it returns the single batch size. + Used to pick and validate the live execution batch at runtime. + + Only decode specializations are considered. Under continuous batching the prefill spec has an + ``input_ids`` batch dim of 1 that is never a selectable decode batch, so specs whose + ``input_ids`` seq_len equals ``prefill_seq_len`` are excluded to avoid admitting a batch size + that has no decode specialization. + + Returns: + list[int]: Sorted, de-duplicated decode input batch sizes, or None if unavailable. + """ + if not self._session.allowed_shapes: + return None + input_ids_idx = self._session.binding_index_map["input_ids"] + decode_batch_sizes = { + shape[input_ids_idx][1][0] + for shape in self._session.allowed_shapes + if shape[input_ids_idx][1][1] != self._prefill_seq_len + } + # Fall back to all specs if every shape looks like a prefill spec (e.g. decode-only QPC where + # prefill_seq_len == decode seq_len), so the method never returns an empty set. + if not decode_batch_sizes: + decode_batch_sizes = {shape[input_ids_idx][1][0] for shape in self._session.allowed_shapes} + return sorted(decode_batch_sizes) + def _fetch_batch_size_prefill_seq_len( self, ): @@ -622,7 +673,7 @@ def prepare_decode_inputs(self): Returns: dict: The decode inputs. """ - batch_size = self.full_batch_size if self.full_batch_size is not None else self.batch_size + batch_size = self.decode_batch_size if self.decode_batch_size is not None else self.batch_size decode_inputs = {} if self.is_tlm: position_ids = np.full((batch_size, self._decode_seq_len), -1, dtype=np.int64) @@ -647,10 +698,12 @@ def prepare_decode_inputs(self): if self._prompt_to_lora_id_mapping_decode: if self.full_batch_size: - first_batch_lora_ids = [self._prompt_to_lora_id_mapping_decode[i] for i in range(self.full_batch_size)] - decode_inputs["lora_ids"] = np.array(first_batch_lora_ids, dtype=np.int64).reshape( - self.full_batch_size, 1 - ) + # Size lora_ids to the live decode batch (decode_batch_size), matching input_ids / + # position_ids / batch_index above. Under dynamic batching decode_batch_size may be a + # smaller compiled batch than full_batch_size (B_max); using B_max here would emit a + # lora_ids batch dim that mismatches the other decode inputs. + first_batch_lora_ids = [self._prompt_to_lora_id_mapping_decode[i] for i in range(batch_size)] + decode_inputs["lora_ids"] = np.array(first_batch_lora_ids, dtype=np.int64).reshape(batch_size, 1) else: batch_lora_ids = [self._prompt_to_lora_id_mapping_decode.popleft() for i in range(self.batch_size)] decode_inputs["lora_ids"] = np.array(batch_lora_ids, dtype=np.int64).reshape(self.batch_size, 1) @@ -737,7 +790,7 @@ def run_prefill_for_all_inputs(self, prompt_queue, generation_len): generation_len (int): The generation length. """ - for decode_batch_id in range(self.full_batch_size): + for decode_batch_id in range(self.decode_batch_size): next_prompt = prompt_queue.popleft() # run prefill for num_chunks @@ -884,19 +937,19 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len): # Set output placeholders for decode self._set_output_buffers( - batch_size=self.full_batch_size, + batch_size=self.decode_batch_size, sequence_length=self._decode_seq_len, ) # Generate flag for tracking progress for each batch ID - current_decode_ongoing = np.full((self.full_batch_size, 1), True) + current_decode_ongoing = np.full((self.decode_batch_size, 1), True) # Generate an array for maintaining the tokens generated in each batch ID - generated_id_current_index = np.ones((self.full_batch_size, 1), np.int64) + generated_id_current_index = np.ones((self.decode_batch_size, 1), np.int64) # Generate a batch ID map for mapping the batch ID if input > full_batch_size. # This ID map will be used for storing all generated tokens - batch_id_map = {i: i for i in range(self.full_batch_size)} + batch_id_map = {i: i for i in range(self.decode_batch_size)} decode_pause_time = 0 # Prepare decode inputs inputs. decode_inputs = self.prepare_decode_inputs() @@ -911,7 +964,7 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len): # Prepare inputs for next iteration next_token_id = self._fetch_next_token_id(outputs) - for decode_batch_id in range(self.full_batch_size): + for decode_batch_id in range(self.decode_batch_size): if ( next_token_id[decode_batch_id, -1] == self.tokenizer.eos_token_id or generated_id_current_index[decode_batch_id] >= self.generation_len[decode_batch_id] @@ -932,7 +985,7 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len): generated_id_current_index[decode_batch_id] = 1 self._set_output_buffers( - batch_size=self.full_batch_size, + batch_size=self.decode_batch_size, sequence_length=self._decode_seq_len, ) decode_pause_time += perf_counter() - start @@ -1116,11 +1169,46 @@ def __init__( def perf_metrics(self): return self._perf_metrics + def _resolve_execution_batch_size(self, execution_batch_size: Optional[int]) -> int: + """ + Resolves the live decode batch for continuous-batching execution. + + For a dynamic-batching QPC compiled with several decode input batches, an explicit + ``execution_batch_size`` must be one of the compiled batches; when omitted, the largest + compiled decode batch is used (all decode slots active). For an ordinary continuous-batching + QPC this returns ``full_batch_size``. + + Args: + execution_batch_size (Optional[int]): Requested live decode batch, or None to auto-select. + + Returns: + int: The decode batch to drive each step. + """ + compiled = self._qaic_model._compiled_batch_sizes + if execution_batch_size is None: + # Default to the largest compiled *decode* batch (all decode slots active). This is not + # necessarily full_batch_size: dynamic batching only requires max(batch_size) <= B_max, so + # a QPC compiled with batch_size=[1,2], full_batch_size=4 has no decode spec at batch 4. + # Falling back to full_batch_size only when the compiled set is unavailable keeps plain + # continuous batching unchanged. + return max(compiled) if compiled else self._full_batch_size + if compiled is not None and execution_batch_size not in compiled: + raise ValueError( + f"execution_batch_size={execution_batch_size} is not a compiled decode batch size " + f"{compiled}. Recompile with this batch size, or pass one of the compiled values." + ) + if execution_batch_size > self._full_batch_size: + raise ValueError( + f"execution_batch_size={execution_batch_size} exceeds full_batch_size={self._full_batch_size}." + ) + return execution_batch_size + def _setup_model_execution_inputs( self, prompt: List[str], generation_len: Optional[int] = None, prompt_to_lora_id_mapping: Optional[List[int]] = None, + execution_batch_size: Optional[int] = None, ): """ This method should be called to set/reset inputs @@ -1128,10 +1216,13 @@ def _setup_model_execution_inputs( :prompt (List[str]): prompts for the model text generation :generation_len (Optional[int], optional): Number of tokens to be generated. :prompt_to_lora_id_mapping (Optional[List[int]], optional): Mapping to associate prompts with their respective LoRA adapter. + :execution_batch_size (Optional[int], optional): Number of decode slots driven per step. + Defaults to full_batch_size (continuous batching) or the compiled batch size. """ - execution_batch_size = ( - self._full_batch_size if self._full_batch_size is not None else self._qaic_model.batch_size - ) + if execution_batch_size is None: + execution_batch_size = ( + self._full_batch_size if self._full_batch_size is not None else self._qaic_model.batch_size + ) max_gen_length = self._ctx_len if not generation_len else max(self._ctx_len, generation_len) # Create a prompt queue. @@ -1193,6 +1284,7 @@ def _continuous_batching_execution( prompt: List[str], generation_len: Optional[int] = None, prompt_to_lora_id_mapping: Optional[List[int]] = None, + execution_batch_size: Optional[int] = None, ): """ Executes the model using continuous batching. @@ -1202,12 +1294,18 @@ def _continuous_batching_execution( :prompt (List[str]): The list of prompts for the model. :generation_len (Optional[int], optional): The generation length. :prompt_to_lora_id_mapping (Optional[List[int]], optional): Mapping to associate prompts with their respective LoRA adapter. + :execution_batch_size (Optional[int], optional): Live decode batch b <= full_batch_size for + dynamic batching. Defaults to full_batch_size (plain continuous batching). Returns: :tuple: A tuple containing performance metrics and generated texts. """ - self._setup_model_execution_inputs(prompt, generation_len, prompt_to_lora_id_mapping) - self._qaic_model.batch_index = np.arange(self._full_batch_size).reshape(-1, 1) + exec_bs = self._resolve_execution_batch_size(execution_batch_size) + self._qaic_model.decode_batch_size = exec_bs + self._setup_model_execution_inputs( + prompt, generation_len, prompt_to_lora_id_mapping, execution_batch_size=exec_bs + ) + self._qaic_model.batch_index = np.arange(exec_bs).reshape(-1, 1) start = perf_counter() self._qaic_model.run_prefill_for_all_inputs(self._prompt_queue, generation_len) @@ -1282,6 +1380,7 @@ def generate( stream: bool = True, automation: Optional[bool] = False, prompt_to_lora_id_mapping: Optional[List[int]] = None, + execution_batch_size: Optional[int] = None, ): """ Executes the model for a given list of prompts and a specified generation length. @@ -1291,6 +1390,9 @@ def generate( generation_len (Optional[int], optional): The generation length. stream (Optional[bool], optional): Boolean flag to enable stream output to console. prompt_to_lora_id_mapping (Optional[List[int]], optional): Mapping to associate prompts with their respective LoRA adapter. + execution_batch_size (Optional[int], optional): Dynamic batching only. Live decode batch + b <= full_batch_size to drive, which must be one of the compiled decode batch sizes. + Defaults to full_batch_size (all decode slots active). Returns: latency_stats (tuple): A tuple containing the generated texts, performance metrics. """ @@ -1298,9 +1400,13 @@ def generate( if self._full_batch_size is not None: logger.warning("Streamer is currently unavailable for continuous batch execution.") perf_metrics, generated_texts = self._continuous_batching_execution( - prompt, generation_len, prompt_to_lora_id_mapping + prompt, generation_len, prompt_to_lora_id_mapping, execution_batch_size ) else: + if execution_batch_size is not None: + raise ValueError( + "execution_batch_size is only supported for continuous-batching (dynamic-batching) QPCs." + ) if stream: print("\nPrompt : " + prompt[0] + "\nCompletion :", flush=True, end="") perf_metrics, generated_texts = self._regular_model_execution( diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 25fac3b646..fef508a7bd 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -4225,6 +4225,7 @@ def build_decode_specialization( kv_cache_batch_size: Optional[int] = None, full_batch_size: Optional[int] = None, num_speculative_tokens: Optional[int] = None, + decode_input_batch_size: Optional[int] = None, **kwargs, ): """ @@ -4244,6 +4245,12 @@ def build_decode_specialization( Continuous batching batch size. Used if `continuous_batching` is enabled. Default is None. num_speculative_tokens : int, optional Number of speculative tokens for Speculative Decoding Target Language Model. Default is None. + decode_input_batch_size : int, optional + Dynamic-batching only (continuous batching): the live decode *input* batch for this + specialization. When set, the input axis (`input_ids`/`position_ids`/`batch_index`) is + sized to this value while the retained KV cache stays pinned at `full_batch_size` (B_max). + This is what lets one QPC hold several decode input batches without varying retained state. + Default is None (input batch = `full_batch_size`, i.e. plain continuous batching). Returns ------- @@ -4254,15 +4261,22 @@ def build_decode_specialization( decode_seq_len = (num_speculative_tokens + 1) if self.is_tlm else 1 if decode_seq_len == prefill_seq_len and not self.continuous_batching: return None + # Under continuous batching the decode input batch is normally `full_batch_size`. For dynamic + # batching we override it per specialization with `decode_input_batch_size` so the input axis + # varies while the retained KV cache (set via `full_batch_size`/`kv_cache_batch_size`) stays fixed. + if self.continuous_batching: + exec_input_bs = decode_input_batch_size if decode_input_batch_size is not None else full_batch_size + else: + exec_input_bs = batch_size if hasattr(self.model, "get_specializations"): spec = self.model.get_specializations( - batch_size=full_batch_size if self.continuous_batching else batch_size, + batch_size=exec_input_bs, prefill_seq_len=(num_speculative_tokens + 1) if self.is_tlm else 1, ctx_len=ctx_len, )[1] else: spec = { - "batch_size": full_batch_size if self.continuous_batching else batch_size, + "batch_size": exec_input_bs, "seq_len": (num_speculative_tokens + 1) if self.is_tlm else 1, "ctx_len": ctx_len, } @@ -4288,7 +4302,7 @@ def compile( ctx_len: int = 128, comp_ctx_lengths_prefill: Optional[List[int]] = None, comp_ctx_lengths_decode: Optional[List[int]] = None, - batch_size: int = 1, + batch_size: Union[int, List[int]] = 1, full_batch_size: Optional[int] = None, kv_cache_batch_size: Optional[int] = None, num_devices: int = 1, @@ -4324,8 +4338,16 @@ def compile( Length of the prefill prompt. Default is 32. ctx_len : int, optional Maximum context length the compiled model can remember. Default is 128. - batch_size : int, optional - Batch size. Default is 1. + batch_size : int or list[int], optional + Batch size. Default is 1. A list requests **dynamic batching**: one decode + specialization is emitted per batch size into a single QPC, and the runtime selects + the matching decode specialization by input tensor batch shape. A list is only + supported when `continuous_batching=True`; the retained KV cache is pinned at + `full_batch_size` (B_max) for every specialization while only the input batch axis + varies, so each value must satisfy `1 <= batch_size <= full_batch_size`. It may be + combined with speculative-decoding `num_speculative_tokens` (a full cartesian product), + but NOT with CCL (`comp_ctx_lengths_decode`) — that combination is rejected because the + two features vary different identifying inputs that the compiler cannot disambiguate. full_batch_size : int, optional Continuous batching batch size. Required if `continuous_batching=True` was set during `from_pretrained`. @@ -4443,8 +4465,38 @@ def compile( "Please pass valid integer for kv_cache_batch_size or full_batch_size, both have same meaning, as continuous_batching is enabled for prefill-only model" ) - # Infer kv_cache_batch_size if not provided - kv_cache_batch_size = kv_cache_batch_size or full_batch_size or batch_size + # Normalize batch_size. A list requests dynamic batching: one decode specialization per + # batch size in a single QPC. A plain int N behaves exactly as before (single-element list). + _is_dynamic_batch = isinstance(batch_size, (list, tuple)) + _decode_bs = sorted(set(batch_size)) if _is_dynamic_batch else [batch_size] + if _is_dynamic_batch: + if any((not isinstance(b, int)) or b < 1 for b in _decode_bs): + raise ValueError(f"All `batch_size` values must be integers >= 1, got {batch_size}.") + # Dynamic batching only works on the continuous-batching export path: there the input + # batch axis (`input_ids`/`position_ids`/`batch_index`) is a distinct ONNX symbol from the + # retained KV-cache batch axis (`full_batch_size`), so decode specializations can vary the + # input batch while sharing one retained-state shape. On the non-CB path the two axes are + # the same symbol, so a list would ask the compiler for multiple retained-state shapes in + # one QPC, which QAIC rejects ("inconsistent retained state"). See + # docs/qaic/dynamic_batching/finding_and_pivot.md. + if not self.continuous_batching: + raise ValueError( + "`batch_size` as a list (dynamic batching) requires `continuous_batching=True` in " + "`from_pretrained`. On the non-continuous-batching path the input batch axis and the " + "retained KV-cache batch axis are the same ONNX symbol, so multiple batch sizes would " + "require multiple retained-state shapes in one QPC, which QAIC rejects." + ) + if full_batch_size is None: + raise TypeError("`full_batch_size` (B_max, the KV-cache capacity) is required for dynamic batching.") + if max(_decode_bs) > full_batch_size: + raise ValueError( + f"Every `batch_size` value must be <= `full_batch_size` (B_max={full_batch_size}); " + f"got {batch_size}." + ) + + # Infer kv_cache_batch_size if not provided. Under continuous batching the retained KV cache + # is always sized at full_batch_size (B_max) regardless of the decode input batch(es). + kv_cache_batch_size = kv_cache_batch_size or full_batch_size or max(_decode_bs) # if ccl_enabled is True read Compute-Context-Length lists if self.ccl_enabled: @@ -4498,6 +4550,31 @@ def compile( ) _decode_ks = [validated_k] + # CCL cannot be crossed with any other co-varying decode axis. CCL distinguishes decode + # specializations by the length of the separate `comp_ctx_lengths` tensor, while dynamic + # batching varies the `input_ids` batch dim and multi-spec SpD varies the `input_ids` + # seq_len. When CCL co-varies with either of those, no single input uniquely identifies a + # specialization and the QAIC compiler rejects the QPC with "No input that uniquely + # identifies specialization" (see docs/qaic/dynamic_batching/finding_and_pivot.md §8). + # Reject early with a clear error instead of deferring to an opaque compiler failure. + # (Checked after the speculative_config collapse above so a list reduced to one effective + # K is not falsely rejected.) + if self.comp_ctx_lengths_decode is not None: + _multi_spec_ks = _decode_ks is not None and self.is_tlm and len(_decode_ks) > 1 + if _is_dynamic_batch: + raise ValueError( + "`batch_size` as a list (dynamic batching) cannot be combined with compute-context-length " + "specializations (`comp_ctx_lengths_decode`/CCL): the two vary different identifying inputs " + "(input_ids batch dim vs the comp_ctx_lengths tensor), which the QAIC compiler cannot " + "disambiguate. Use one or the other." + ) + if _multi_spec_ks: + raise ValueError( + "Multi-spec speculative decoding (`num_speculative_tokens` as a list) cannot be combined " + "with compute-context-length specializations (`comp_ctx_lengths_decode`/CCL): the two vary " + "different identifying inputs (input_ids seq_len vs the comp_ctx_lengths tensor), which the " + "QAIC compiler cannot disambiguate. Pass a plain int for num_speculative_tokens when using CCL." + ) if ( self.model.qaic_config is not None and self.model.qaic_config.get("include_sampler", False) @@ -4515,6 +4592,17 @@ def compile( retain_full_kv = False # --- Specializations --- + # Decode specializations are a full cartesian product of batch_size × CCL × spec_len. + # Warn (do not block) on large counts, mirroring the soft seq_len guidance elsewhere. + _n_ccl = len(self.comp_ctx_lengths_decode) if self.comp_ctx_lengths_decode is not None else 1 + _n_ks = len(_decode_ks) if (self.is_tlm and _decode_ks is not None) else 1 + _projected_decode_specs = len(_decode_bs) * _n_ccl * _n_ks + if _projected_decode_specs > 15: + logger.warning( + f"Compiling {_projected_decode_specs} decode specializations " + f"(batch×CCL×spec_len = {len(_decode_bs)}×{_n_ccl}×{_n_ks}); " + "large counts increase compile time and QPC size." + ) specializations = [] if prefill_only is None or prefill_only or prefill_seq_len == 1: # TODO: we are handling decode-only case inside prefill call which is utterly mis-leading @@ -4527,7 +4615,7 @@ def compile( prefill_seq_len=prefill_seq_len, ctx_len=ctx_len, comp_ctx_lengths=ccl_lengths[i], - batch_size=batch_size, + batch_size=max(_decode_bs), kv_cache_batch_size=kv_cache_batch_size, full_batch_size=full_batch_size, prefill_only=prefill_only, @@ -4539,7 +4627,7 @@ def compile( self.build_prefill_specialization( prefill_seq_len=prefill_seq_len, ctx_len=ctx_len, - batch_size=batch_size, + batch_size=max(_decode_bs), kv_cache_batch_size=kv_cache_batch_size, full_batch_size=full_batch_size, prefill_only=prefill_only, @@ -4548,58 +4636,31 @@ def compile( ) if (prefill_only is None or not prefill_only) and prefill_seq_len != 1: - if _decode_ks is not None and self.is_tlm: - # TLM multi-spec path: one decode specialization per K in num_speculative_tokens. - # CCL (comp_ctx_lengths) + multi-spec TLM is not yet supported: the per-K call - # to build_decode_specialization would need to iterate over CCL values, producing - # len(decode_ks) × len(comp_ctx_lengths_decode) decode specializations whose - # naming and ordering is untested. Reject early so users get a clear error - # instead of a silently wrong QPC. - if self.comp_ctx_lengths_decode is not None: - raise NotImplementedError( - "TLM multi-spec (num_speculative_tokens as a list) combined with " - "comp_ctx_lengths_decode is not yet supported. Pass a plain int for " - "num_speculative_tokens when using CCL." - ) - for k in _decode_ks: - spec = self.build_decode_specialization( - num_speculative_tokens=k, - prefill_seq_len=prefill_seq_len, - ctx_len=ctx_len, - batch_size=batch_size, - kv_cache_batch_size=kv_cache_batch_size, - full_batch_size=full_batch_size, - ) - if spec is not None: - specializations.append(spec) - - elif self.comp_ctx_lengths_decode is not None: - # CCL loop (non-TLM) - for i in range(0, len(self.comp_ctx_lengths_decode)): - decode_spec = self.build_decode_specialization( - prefill_seq_len=prefill_seq_len, - ctx_len=ctx_len, - comp_ctx_lengths=self.comp_ctx_lengths_decode[i], - batch_size=batch_size, - kv_cache_batch_size=kv_cache_batch_size, - full_batch_size=full_batch_size, - num_speculative_tokens=None, - ) - if decode_spec: - specializations.append(decode_spec) - - else: - decode_spec = self.build_decode_specialization( - prefill_seq_len=prefill_seq_len, - ctx_len=ctx_len, - batch_size=batch_size, - kv_cache_batch_size=kv_cache_batch_size, - full_batch_size=full_batch_size, - num_speculative_tokens=None, - prefill_only=prefill_only, - ) - if decode_spec: - specializations.append(decode_spec) + # Decode specializations expand as a full cartesian product over + # (batch_size × comp_ctx_lengths_decode × num_speculative_tokens). [None] sentinels + # stand in for "no CCL" / "no SpD" so the single loop covers every combination, + # including the legacy single-decode-spec case. For dynamic batching (continuous + # batching), each spec varies only its decode *input* batch (`decode_input_batch_size`) + # while `kv_cache_batch_size` stays pinned at full_batch_size (B_max) so the retained + # KV-cache shape is identical across all specs. + _ccl_decode_vals = self.comp_ctx_lengths_decode if self.comp_ctx_lengths_decode is not None else [None] + _ks_vals = _decode_ks if (self.is_tlm and _decode_ks is not None) else [None] + for bs in _decode_bs: + for ccl in _ccl_decode_vals: + for k in _ks_vals: + decode_spec = self.build_decode_specialization( + prefill_seq_len=prefill_seq_len, + ctx_len=ctx_len, + comp_ctx_lengths=ccl, + batch_size=bs, + kv_cache_batch_size=(kv_cache_batch_size if self.continuous_batching else bs), + full_batch_size=full_batch_size, + num_speculative_tokens=k, + decode_input_batch_size=(bs if (self.continuous_batching and _is_dynamic_batch) else None), + prefill_only=prefill_only, + ) + if decode_spec is not None: + specializations.append(decode_spec) if kw_spec := compiler_options.pop("specializations", None): specializations = kw_spec diff --git a/tests/transformers/spd/test_spd_inference.py b/tests/transformers/spd/test_spd_inference.py index feb0153e3c..2580f03f4b 100644 --- a/tests/transformers/spd/test_spd_inference.py +++ b/tests/transformers/spd/test_spd_inference.py @@ -16,6 +16,7 @@ from transformers import AutoConfig, AutoTokenizer from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.generation.text_generation_inference import TextGeneration from QEfficient.utils.constants import Constants from QEfficient.utils.test_utils import load_qeff_causal_lm_model @@ -585,3 +586,178 @@ def test_multi_spec_qpc_logit_correctness(decode_ks, manual_cleanup): assert total_assertions > 0 manual_cleanup([vanilla.onnx_path, tlm.onnx_path]) + + +# --------------------------------------------------------------------------- +# Dynamic-batching — hardware-level QPC compile + per-batch execution test +# --------------------------------------------------------------------------- + +_DYN_BATCH_MODEL = "JackFram/llama-68m" +_DYN_BATCH_NUM_LAYERS = 2 +_DYN_BATCH_PREFILL_LEN = 32 +_DYN_BATCH_CTX_LEN = 128 + + +@pytest.mark.on_qaic +@pytest.mark.feature +@pytest.mark.parametrize("batch_sizes", [[1, 2, 4]]) +def test_dynamic_batch_qpc_per_batch_execution(batch_sizes, manual_cleanup): + """ + Compile ONE continuous-batching QPC carrying one decode specialization per input batch size + while the retained KV cache is pinned at ``full_batch_size`` (B_max), then verify each live + batch executes. + + This is the direct regression against the finding-doc hardware failure ("inconsistent retained + state"): the compile must succeed with a single retained-state shape. For each batch size ``b`` + in the list, run a decode step with ``(b, 1)`` ``input_ids`` and a length-``b`` ``batch_index`` + that routes those ``b`` sequences into ``b`` of the ``B_max`` KV slots, and assert finite logits + with batch dimension ``b``. + """ + b_max = max(batch_sizes) + tokenizer = AutoTokenizer.from_pretrained(_DYN_BATCH_MODEL, padding_side="right") + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + vocab_size = len(tokenizer) + + qeff_model = load_qeff_causal_lm_model( + _DYN_BATCH_MODEL, num_hidden_layers=_DYN_BATCH_NUM_LAYERS, continuous_batching=True + ) + qpc_path = qeff_model.compile( + num_cores=2, + prefill_seq_len=_DYN_BATCH_PREFILL_LEN, + ctx_len=_DYN_BATCH_CTX_LEN, + aic_enable_depth_first=True, + batch_size=batch_sizes, + full_batch_size=b_max, + ) + assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) + + for bs in batch_sizes: + session = QAICInferenceSession(qpc_path) + session.skip_buffers([x for x in session.input_names if x.startswith("past_")]) + session.skip_buffers([x for x in session.output_names if x.endswith("_RetainedState")]) + + ph = np.zeros((bs, 1, vocab_size), dtype=np.float32) + session.set_buffers({"logits": ph}) + out = session.run( + { + "input_ids": np.zeros((bs, 1), dtype=np.int64), + "position_ids": np.zeros((bs, 1), dtype=np.int64), + "batch_index": np.arange(bs, dtype=np.int64).reshape(-1, 1), + } + ) + logits = out["logits"] + assert logits.shape[0] == bs, f"batch_size={bs}: expected batch dim {bs}, got {logits.shape}" + assert np.isfinite(logits).all(), f"batch_size={bs}: non-finite logits" + del session + + manual_cleanup([qeff_model.onnx_path]) + + +@pytest.mark.on_qaic +@pytest.mark.feature +@pytest.mark.parametrize("exec_bs", [1, 2]) +def test_dynamic_batch_generation_parity(exec_bs, manual_cleanup): + """ + End-to-end parity: a dynamic-batching QPC (``batch_size=[1,2,4]``, ``full_batch_size=4``) run at + a live batch ``exec_bs`` must produce the same generated token ids as a QPC compiled natively at + ``continuous_batching=True, full_batch_size=exec_bs`` (single decode spec). + + Guards against the KV-slicing corrupting active-slot outputs when the live batch is smaller than + B_max (CLAUDE.md: never claim parity from compile success alone). + """ + prompts = ["The capital of France is", "My favorite color is", "Once upon a time", "In the year 2050"][:exec_bs] + gen_len = 16 + tokenizer = AutoTokenizer.from_pretrained(_DYN_BATCH_MODEL, padding_side="right") + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + # Dynamic-batching QPC: many decode input batches, KV pinned at B_max=4. + dyn = load_qeff_causal_lm_model(_DYN_BATCH_MODEL, num_hidden_layers=_DYN_BATCH_NUM_LAYERS, continuous_batching=True) + dyn_qpc = dyn.compile( + num_cores=2, + prefill_seq_len=_DYN_BATCH_PREFILL_LEN, + ctx_len=_DYN_BATCH_CTX_LEN, + aic_enable_depth_first=True, + batch_size=[1, 2, 4], + full_batch_size=4, + ) + + # Reference QPC: natively compiled at full_batch_size == exec_bs (single decode spec). + ref = load_qeff_causal_lm_model(_DYN_BATCH_MODEL, num_hidden_layers=_DYN_BATCH_NUM_LAYERS, continuous_batching=True) + ref_qpc = ref.compile( + num_cores=2, + prefill_seq_len=_DYN_BATCH_PREFILL_LEN, + ctx_len=_DYN_BATCH_CTX_LEN, + aic_enable_depth_first=True, + full_batch_size=exec_bs, + ) + + dyn_gen = TextGeneration(tokenizer=tokenizer, qpc_path=dyn_qpc, full_batch_size=4, ctx_len=_DYN_BATCH_CTX_LEN) + dyn_gen.generate(prompts, generation_len=gen_len, execution_batch_size=exec_bs, stream=False) + dyn_ids = dyn_gen._qaic_model.generated_ids[:exec_bs] + + ref_gen = TextGeneration(tokenizer=tokenizer, qpc_path=ref_qpc, full_batch_size=exec_bs, ctx_len=_DYN_BATCH_CTX_LEN) + ref_gen.generate(prompts, generation_len=gen_len, stream=False) + ref_ids = ref_gen._qaic_model.generated_ids[:exec_bs] + + assert np.array_equal(dyn_ids, ref_ids), ( + f"exec_bs={exec_bs}: dynamic-batch generation diverged from native batch-{exec_bs} reference.\n" + f"dyn={dyn_ids.tolist()}\nref={ref_ids.tolist()}" + ) + manual_cleanup([dyn.onnx_path, ref.onnx_path]) + + +@pytest.mark.on_qaic +@pytest.mark.feature +def test_dynamic_batch_times_spec_len_compiles(manual_cleanup): + """ + Regression: dynamic batching stacks with speculative decoding. A TLM compiled with + ``batch_size=[1,2]`` × ``num_speculative_tokens=[1,3]`` must produce ONE continuous-batching QPC + whose decode specializations vary input batch and seq_len while the retained KV stays pinned at + B_max (full_batch_size). + + Note: dynamic batching is NOT stacked with CCL here — combining ``comp_ctx_lengths`` with + speculative-decoding specializations is a pre-existing QAIC compiler limitation ("No input that + uniquely identifies specialization") that reproduces even without a batch list, so it is out of + scope for this feature. + """ + tlm = load_qeff_causal_lm_model( + _DYN_BATCH_MODEL, + num_hidden_layers=_DYN_BATCH_NUM_LAYERS, + continuous_batching=True, + qaic_config={"speculative_model_type": "target"}, + ) + qpc_path = tlm.compile( + num_cores=2, + prefill_seq_len=_DYN_BATCH_PREFILL_LEN, + ctx_len=_DYN_BATCH_CTX_LEN, + aic_enable_depth_first=True, + batch_size=[1, 2], + full_batch_size=2, + num_speculative_tokens=[1, 3], + ) + assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) + manual_cleanup([tlm.onnx_path]) + + +@pytest.mark.on_qaic +@pytest.mark.feature +def test_dynamic_batch_with_ccl_rejected(): + """Dynamic batching combined with CCL must be rejected at compile time (not deferred to the + QAIC compiler): the two vary different identifying inputs and cannot be disambiguated.""" + model = load_qeff_causal_lm_model( + _DYN_BATCH_MODEL, + num_hidden_layers=_DYN_BATCH_NUM_LAYERS, + continuous_batching=True, + ) + with pytest.raises(ValueError, match="comp_ctx_lengths"): + model.compile( + num_cores=2, + prefill_seq_len=_DYN_BATCH_PREFILL_LEN, + ctx_len=2048, + aic_enable_depth_first=True, + batch_size=[1, 2], + full_batch_size=2, + comp_ctx_lengths_decode=[1024, 2048], + ) diff --git a/tests/unit_test/models/test_modeling_auto_cpu.py b/tests/unit_test/models/test_modeling_auto_cpu.py index 82291ef022..ab4e81796a 100644 --- a/tests/unit_test/models/test_modeling_auto_cpu.py +++ b/tests/unit_test/models/test_modeling_auto_cpu.py @@ -1283,3 +1283,202 @@ def test_compile_int_zero_backward_compat(self): decode_specs = [s for s in specs if s.get("seq_len", 0) != 32] assert len(decode_specs) == 1, f"Expected 1 decode spec for scalar 0, got: {decode_specs}" assert decode_specs[0]["seq_len"] == 1 # k=0 → seq_len=1 + + +# --------------------------------------------------------------------------- +# Dynamic-batching specialization unit tests (batch_size as list) +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +@pytest.mark.causal_lm +class TestDynamicBatchSpecializations: + """Tests for dynamic-batching decode specializations (batch_size as a list). + + Dynamic batching rides the continuous-batching export path: every decode specialization keeps + the retained KV cache pinned at ``full_batch_size`` (B_max) while only the decode *input* batch + axis varies over the requested list. This is what makes a single QPC legal on QAIC (one retained + state shape) — see docs/qaic/dynamic_batching/finding_and_pivot.md. + """ + + @staticmethod + def _capture_specializations(qeff, **compile_kwargs): + """Run compile() with _compile mocked and return the captured specializations list.""" + from unittest.mock import patch + + captured = {} + with patch.object( + type(qeff), + "_compile", + side_effect=lambda *args, **kw: ( + captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" + ), + ): + qeff.compile(**compile_kwargs) + assert captured.get("specializations") is not None, "_compile was not reached" + return captured["specializations"] + + @staticmethod + def _decode_specs(specs, prefill_seq_len=32): + return [s for s in specs if s.get("seq_len", 0) != prefill_seq_len] + + def test_list_batch_size_produces_one_decode_spec_per_batch(self): + """batch_size=[1, 2, 4] → 3 decode specs whose input batch is {1, 2, 4}.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + specs = self._capture_specializations( + qeff, prefill_seq_len=32, ctx_len=128, batch_size=[1, 2, 4], full_batch_size=4 + ) + decode_specs = self._decode_specs(specs) + assert len(decode_specs) == 3 + assert {s["batch_size"] for s in decode_specs} == {1, 2, 4} + + def test_all_decode_specs_share_kv_batch_bmax(self): + """Retained KV batch (full_batch_size) must be identical (=B_max) across every decode spec. + + This is the regression guard for the hardware failure in finding_and_pivot.md: the earlier + attempt sized KV per-batch, giving inconsistent retained state. Here KV is pinned at B_max. + """ + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + specs = self._capture_specializations( + qeff, prefill_seq_len=32, ctx_len=128, batch_size=[1, 2, 4], full_batch_size=4 + ) + # Every specialization (prefill + decode) must carry full_batch_size == 4. + assert {s["full_batch_size"] for s in specs} == {4} + # Prefill runs at input batch 1 under continuous batching. + prefill_specs = [s for s in specs if s.get("seq_len", 0) == 32] + assert len(prefill_specs) == 1 + assert prefill_specs[0]["batch_size"] == 1 + assert prefill_specs[0]["full_batch_size"] == 4 + + def test_int_batch_size_backward_compat(self): + """batch_size=2 as a plain int (continuous batching) → one decode spec, KV=full_batch_size.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + specs = self._capture_specializations(qeff, prefill_seq_len=32, ctx_len=128, batch_size=2, full_batch_size=4) + decode_specs = self._decode_specs(specs) + assert len(decode_specs) == 1 + # Plain int under continuous batching keeps the legacy behavior: input batch == full_batch_size. + assert decode_specs[0]["batch_size"] == 4 + assert decode_specs[0]["full_batch_size"] == 4 + + def test_default_batch_size_legacy_shape(self): + """No batch_size, no continuous batching → 1 prefill + 1 decode spec, decode batch 1.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model) + specs = self._capture_specializations(qeff, prefill_seq_len=32, ctx_len=128) + decode_specs = self._decode_specs(specs) + prefill_specs = [s for s in specs if s.get("seq_len", 0) == 32] + assert len(prefill_specs) == 1 + assert len(decode_specs) == 1 + assert decode_specs[0]["batch_size"] == 1 + + def test_batch_size_times_spec_len_product(self): + """TLM batch_size=[1,2] × num_speculative_tokens=[1,3] → 4 decode specs (2×2 product).""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM( + model, continuous_batching=True, qaic_config={"speculative_model_type": "target"} + ) + specs = self._capture_specializations( + qeff, + prefill_seq_len=32, + ctx_len=128, + batch_size=[1, 2], + full_batch_size=2, + num_speculative_tokens=[1, 3], + ) + decode_specs = self._decode_specs(specs) + assert len(decode_specs) == 4 + # seq_len = k+1 → {2, 4}; input batch_size → {1, 2}; KV pinned at B_max=2. + assert {(s["batch_size"], s["seq_len"]) for s in decode_specs} == {(1, 2), (1, 4), (2, 2), (2, 4)} + assert {s["full_batch_size"] for s in decode_specs} == {2} + + def test_batch_size_with_ccl_rejected(self): + """Dynamic batching + CCL is rejected (compiler cannot disambiguate identifying inputs).""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + with pytest.raises(ValueError, match="comp_ctx_lengths"): + qeff.compile( + prefill_seq_len=32, + ctx_len=2048, + batch_size=[1, 2], + full_batch_size=2, + comp_ctx_lengths_decode=[1024, 2048], + ) + + def test_tlm_multispec_with_ccl_rejected(self): + """TLM multi-spec (num_speculative_tokens list) + CCL must be rejected at compile time. + + CCL crossed with a co-varying seq_len axis produces decode specs the QAIC compiler cannot + disambiguate ("No input that uniquely identifies specialization" — finding_and_pivot.md §8, + Case 3). We reject early with a clear ValueError instead of deferring to the compiler. + """ + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM( + model, continuous_batching=True, qaic_config={"speculative_model_type": "target"} + ) + with pytest.raises(ValueError, match="comp_ctx_lengths"): + qeff.compile( + prefill_seq_len=32, + ctx_len=2048, + full_batch_size=1, + comp_ctx_lengths_decode=[1024, 2048], + num_speculative_tokens=[1, 3], + ) + + def test_scalar_spec_len_with_ccl_still_allowed(self): + """A plain int num_speculative_tokens still combines with CCL (single seq_len, CCL varies alone).""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM( + model, continuous_batching=True, qaic_config={"speculative_model_type": "target"} + ) + specs = self._capture_specializations( + qeff, + prefill_seq_len=32, + ctx_len=2048, + full_batch_size=1, + comp_ctx_lengths_decode=[1024, 2048], + num_speculative_tokens=3, + ) + decode_specs = self._decode_specs(specs) + # One seq_len (k=3 → 4) crossed with two CCL values → 2 decode specs, CCL is the sole discriminator. + assert {s["seq_len"] for s in decode_specs} == {4} + assert {s["comp_ctx_lengths"] for s in decode_specs} == {1024, 2048} + + def test_batch_ccl_spec_len_combo_rejected(self): + """batch_size list + CCL + num_speculative_tokens → ValueError (CCL cannot combine with a batch list).""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM( + model, continuous_batching=True, qaic_config={"speculative_model_type": "target"} + ) + with pytest.raises(ValueError, match="comp_ctx_lengths"): + qeff.compile( + prefill_seq_len=32, + ctx_len=2048, + batch_size=[1, 2], + full_batch_size=2, + comp_ctx_lengths_decode=[1024, 2048], + num_speculative_tokens=[1, 3], + ) + + def test_list_batch_size_rejected_without_continuous_batching(self): + """batch_size list without continuous_batching → ValueError (non-CB cannot decouple axes).""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model) + with pytest.raises(ValueError, match="continuous_batching=True"): + qeff.compile(prefill_seq_len=32, ctx_len=128, batch_size=[1, 2]) + + def test_list_batch_size_requires_full_batch_size(self): + """batch_size list + continuous_batching but no full_batch_size → TypeError.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + with pytest.raises(TypeError, match="full_batch_size"): + qeff.compile(prefill_seq_len=32, ctx_len=128, batch_size=[1, 2]) + + def test_batch_size_exceeding_full_batch_size_rejected(self): + """Any batch_size value > full_batch_size (B_max) → ValueError.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + with pytest.raises(ValueError, match="full_batch_size"): + qeff.compile(prefill_seq_len=32, ctx_len=128, batch_size=[1, 2, 8], full_batch_size=4) diff --git a/tests/unit_test/utils/test_auto_model_api.py b/tests/unit_test/utils/test_auto_model_api.py index 07428e94e7..cd9677ccd4 100644 --- a/tests/unit_test/utils/test_auto_model_api.py +++ b/tests/unit_test/utils/test_auto_model_api.py @@ -222,6 +222,46 @@ def test_build_decode_specialization_with_num_speculative_tokens(self): # The result should reflect the speculative tokens in some way assert "ctx_len" in result + def test_build_decode_specialization_batch_size_from_kv_cache_batch_size(self): + """Non-CB decode spec batch_size is taken from kv_cache_batch_size.""" + qeff = self._make_qeff() + result = qeff.build_decode_specialization(ctx_len=32, batch_size=4, kv_cache_batch_size=4, full_batch_size=None) + assert result["batch_size"] == 4 + + def test_build_decode_specialization_ccl_and_num_speculative_tokens_together(self): + """build_decode_specialization accepts comp_ctx_lengths and num_speculative_tokens simultaneously.""" + qeff = self._make_qeff() + qeff.is_tlm = True + result = qeff.build_decode_specialization( + ctx_len=128, + batch_size=1, + kv_cache_batch_size=1, + full_batch_size=None, + comp_ctx_lengths=64, + num_speculative_tokens=3, + prefill_seq_len=32, + ) + assert result is not None + assert result["seq_len"] == 4 # k+1 + assert result["comp_ctx_lengths"] == 64 + + def test_build_decode_specialization_dynamic_batch_pins_kv_at_bmax(self): + """Dynamic batching: decode input batch == decode_input_batch_size, KV batch == full_batch_size.""" + from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM + + qeff = QEFFAutoModelForCausalLM(make_tiny_gpt2(), continuous_batching=True) + result = qeff.build_decode_specialization( + ctx_len=128, + batch_size=2, + kv_cache_batch_size=4, + full_batch_size=4, + decode_input_batch_size=2, + ) + assert result is not None + # Input axis follows the live batch; retained KV batch stays pinned at B_max. + assert result["batch_size"] == 2 + assert result["full_batch_size"] == 4 + # --------------------------------------------------------------------------- # Tests: QEFFAutoModelForCausalLM prefill toggle diff --git a/tests/unit_test/utils/test_generation.py b/tests/unit_test/utils/test_generation.py index e683e849e7..645fb53659 100644 --- a/tests/unit_test/utils/test_generation.py +++ b/tests/unit_test/utils/test_generation.py @@ -966,6 +966,112 @@ def test_run_prefill_for_all_inputs_empties_queue(self): assert len(prompt_queue) == 0 +# --------------------------------------------------------------------------- +# Tests: Dynamic batching runtime helpers (batch_index KV-slicing) +# Exercises _fetch_full_batch_size / _fetch_compiled_batch_sizes / +# _resolve_execution_batch_size for a QPC compiled with several decode input +# batches while the retained KV cache is pinned at B_max. All CPU-only via a +# mocked session (allowed_shapes carry the compiled specialization shapes). +# --------------------------------------------------------------------------- + + +class TestDynamicBatchRuntime: + """Runtime resolution of the live decode batch for a dynamic-batching QPC. + + Regression guards for the code-review fixes: the largest *decode* batch (not B_max) is the + default; the prefill spec's batch-1 input must not be admitted as a decode batch; and B_max is + read from the retained KV cache, not the input-axis batch_index. + See docs/qaic/dynamic_batching/code_review_findings.md. + """ + + def _make_dynamic_batch_instance(self, decode_batches, full_batch_size): + """Mocked QEffTextGenerationBase for a QPC with decode specs over `decode_batches`, + KV pinned at `full_batch_size` (B_max). allowed_shapes carry one CB prefill spec (input + batch 1) plus one decode spec per requested decode batch; the retained past_key binding + carries B_max on its batch axis.""" + from QEfficient.generation.text_generation_inference import QEffTextGenerationBase, TextGeneration + + tok = _make_tokenizer() + mock_session = _make_mock_session( + batch_size=max(decode_batches), + prefill_seq_len=PREFILL_LEN, + ctx_len=CTX_LEN, + full_batch_size=full_batch_size, + force_seq_len=1, + ) + # Append a retained-state KV binding (batch axis = B_max) and a batch_index input. + for name, dims in (("past_key.0", [full_batch_size, 4, CTX_LEN, 8]), ("batch_index", [full_batch_size, 1])): + b = MagicMock() + b.name, b.dims, b.dir, b.type = name, dims, "input", 1 + b.size = int(np.prod(dims)) * 4 + mock_session.bindings.append(b) + mock_session.binding_index_map[name] = len(mock_session.bindings) - 1 + mock_session.input_names.append("batch_index") + + ii = mock_session.binding_index_map["input_ids"] + pk = mock_session.binding_index_map["past_key.0"] + bi = mock_session.binding_index_map["batch_index"] + lg = mock_session.binding_index_map["logits"] + + def _shape_row(input_batch, seq_len): + row = [None] * len(mock_session.bindings) + row[ii] = (4, [input_batch, seq_len]) + row[pk] = (4, [full_batch_size, 4, CTX_LEN, 8]) # KV batch axis = B_max, constant + row[bi] = (4, [input_batch, 1]) + row[lg] = (4, [input_batch, seq_len, VOCAB_SIZE]) + return row + + # CB prefill spec has input batch 1; one decode spec per compiled decode batch. + allowed = [_shape_row(1, PREFILL_LEN)] + allowed += [_shape_row(b, 1) for b in decode_batches] + mock_session.allowed_shapes = allowed + + with patch( + "QEfficient.generation.text_generation_inference.QAICInferenceSession", + return_value=mock_session, + ): + base = QEffTextGenerationBase( + tokenizer=tok, qpc_path="/fake/path/model.qpc", ctx_len=CTX_LEN, full_batch_size=full_batch_size + ) + gen = object.__new__(TextGeneration) + gen._qaic_model = base + gen._full_batch_size = base.full_batch_size + return base, gen + + def test_compiled_batch_sizes_excludes_prefill(self): + """The prefill spec's input batch (1) must not leak into the compiled decode batch set.""" + base, _ = self._make_dynamic_batch_instance(decode_batches=[2, 4], full_batch_size=4) + assert base._compiled_batch_sizes == [2, 4] + + def test_full_batch_size_reads_kv_capacity_not_input_axis(self): + """B_max is read from the retained KV batch axis (4), not max input batch (2).""" + base, _ = self._make_dynamic_batch_instance(decode_batches=[1, 2], full_batch_size=4) + assert base.full_batch_size == 4 + + def test_resolve_default_uses_largest_compiled_decode_batch(self): + """Default execution batch is the largest compiled decode batch, not B_max.""" + _, gen = self._make_dynamic_batch_instance(decode_batches=[1, 2], full_batch_size=4) + assert gen._resolve_execution_batch_size(None) == 2 + + def test_resolve_explicit_compiled_batch_accepted(self): + _, gen = self._make_dynamic_batch_instance(decode_batches=[1, 2, 4], full_batch_size=4) + assert gen._resolve_execution_batch_size(2) == 2 + + def test_resolve_uncompiled_batch_rejected(self): + """A batch with no decode spec (e.g. the prefill-only 1 when it isn't a decode batch, or 3) + is rejected up front rather than deferred to a runtime spec-match failure.""" + _, gen = self._make_dynamic_batch_instance(decode_batches=[2, 4], full_batch_size=4) + with pytest.raises(ValueError, match="not a compiled decode batch size"): + gen._resolve_execution_batch_size(3) + with pytest.raises(ValueError, match="not a compiled decode batch size"): + gen._resolve_execution_batch_size(1) + + def test_resolve_batch_exceeding_bmax_rejected(self): + _, gen = self._make_dynamic_batch_instance(decode_batches=[2, 4], full_batch_size=4) + with pytest.raises(ValueError, match="not a compiled decode batch size"): + gen._resolve_execution_batch_size(8) + + # --------------------------------------------------------------------------- # Tests: _fetch_next_token_id # --------------------------------------------------------------------------- From 6b7252c72e78da9ac2d4289072fe0b54068f7323 Mon Sep 17 00:00:00 2001 From: eplatero Date: Fri, 31 Jul 2026 14:25:08 -0500 Subject: [PATCH 2/4] Collapse decode_input_batch_size into batch_size in build_decode_specialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit build_decode_specialization carried two params for the same concept — the live decode *input* batch of one specialization: batch_size (read only on the non-CB path) and decode_input_batch_size (the CB override). At the single call site the loop variable was passed under both names, and on the CB path batch_size was dead. Collapse them into a single batch_size that means "decode input batch for this spec" on every path (exec_input_bs = batch_size). The caller now computes the one input-batch value (full_batch_size for plain CB, the per-spec bs for dynamic batching and the non-CB path), so emitted specializations are unchanged. Harden the retained-batch write to fall back to full_batch_size when kv_cache_batch_size is omitted so a direct caller still pins B_max. Also update the dynamic-batch unit test to express the live batch via batch_size. Validation: 219 CPU unit tests pass across test_auto_model_api, test_modeling_auto_cpu, test_speculative_decoding, and test_pld_inference; ruff check/format clean. Co-Authored-By: Claude Opus 4.8 Signed-off-by: eplatero --- .../transformers/models/modeling_auto.py | 44 +++++++++---------- tests/unit_test/utils/test_auto_model_api.py | 7 ++- 2 files changed, 25 insertions(+), 26 deletions(-) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index fef508a7bd..6f726f001e 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -4225,7 +4225,6 @@ def build_decode_specialization( kv_cache_batch_size: Optional[int] = None, full_batch_size: Optional[int] = None, num_speculative_tokens: Optional[int] = None, - decode_input_batch_size: Optional[int] = None, **kwargs, ): """ @@ -4238,19 +4237,17 @@ def build_decode_specialization( ctx_len : int, optional Maximum context length the compiled model can remember. Default is 128. batch_size : int, optional - Batch size for the decode phase. Default is 1. + The decode *input* batch for this specialization, i.e. the size of the input axis + (`input_ids`/`position_ids`/`batch_index`). On every path this is the live decode + batch; the retained KV-cache batch (B_max) is carried separately by + `kv_cache_batch_size`/`full_batch_size`, so a single spec can pin the retained state + while its input batch varies (dynamic batching). Default is 1. kv_cache_batch_size : int, optional Batch size for KV cache. If not provided, it defaults based on `full_batch_size` or `batch_size`. full_batch_size : int, optional Continuous batching batch size. Used if `continuous_batching` is enabled. Default is None. num_speculative_tokens : int, optional Number of speculative tokens for Speculative Decoding Target Language Model. Default is None. - decode_input_batch_size : int, optional - Dynamic-batching only (continuous batching): the live decode *input* batch for this - specialization. When set, the input axis (`input_ids`/`position_ids`/`batch_index`) is - sized to this value while the retained KV cache stays pinned at `full_batch_size` (B_max). - This is what lets one QPC hold several decode input batches without varying retained state. - Default is None (input batch = `full_batch_size`, i.e. plain continuous batching). Returns ------- @@ -4261,13 +4258,10 @@ def build_decode_specialization( decode_seq_len = (num_speculative_tokens + 1) if self.is_tlm else 1 if decode_seq_len == prefill_seq_len and not self.continuous_batching: return None - # Under continuous batching the decode input batch is normally `full_batch_size`. For dynamic - # batching we override it per specialization with `decode_input_batch_size` so the input axis - # varies while the retained KV cache (set via `full_batch_size`/`kv_cache_batch_size`) stays fixed. - if self.continuous_batching: - exec_input_bs = decode_input_batch_size if decode_input_batch_size is not None else full_batch_size - else: - exec_input_bs = batch_size + # `batch_size` is the decode *input* batch for this specialization on every path. The + # retained KV-cache batch (B_max) is carried separately by `kv_cache_batch_size`, so a + # single spec can pin the retained state while its input batch varies (dynamic batching). + exec_input_bs = batch_size if hasattr(self.model, "get_specializations"): spec = self.model.get_specializations( batch_size=exec_input_bs, @@ -4285,10 +4279,13 @@ def build_decode_specialization( spec["num_logits_to_keep"] = (num_speculative_tokens + 1) if self.is_tlm else None + # The retained KV-cache batch (B_max) is `kv_cache_batch_size`; fall back to + # `full_batch_size` so a direct caller that passes only `full_batch_size` still pins it. + retained_batch_size = kv_cache_batch_size if kv_cache_batch_size is not None else full_batch_size if self.continuous_batching: - spec["full_batch_size"] = kv_cache_batch_size + spec["full_batch_size"] = retained_batch_size else: - spec["batch_size"] = kv_cache_batch_size + spec["batch_size"] = retained_batch_size result = {k: v for k, v in spec.items() if v is not None} result["_graph_name"] = "Decode" return result @@ -4640,23 +4637,26 @@ def compile( # (batch_size × comp_ctx_lengths_decode × num_speculative_tokens). [None] sentinels # stand in for "no CCL" / "no SpD" so the single loop covers every combination, # including the legacy single-decode-spec case. For dynamic batching (continuous - # batching), each spec varies only its decode *input* batch (`decode_input_batch_size`) - # while `kv_cache_batch_size` stays pinned at full_batch_size (B_max) so the retained - # KV-cache shape is identical across all specs. + # batching), each spec varies only its decode *input* batch (the `batch_size` passed + # per iteration) while `kv_cache_batch_size` stays pinned at full_batch_size (B_max) + # so the retained KV-cache shape is identical across all specs. _ccl_decode_vals = self.comp_ctx_lengths_decode if self.comp_ctx_lengths_decode is not None else [None] _ks_vals = _decode_ks if (self.is_tlm and _decode_ks is not None) else [None] for bs in _decode_bs: for ccl in _ccl_decode_vals: for k in _ks_vals: + # `batch_size` is the decode input batch for this spec. On plain continuous + # batching the input batch is full_batch_size (B_max); for dynamic batching + # (and the non-CB path) it is the per-spec batch `bs`. + input_bs = full_batch_size if (self.continuous_batching and not _is_dynamic_batch) else bs decode_spec = self.build_decode_specialization( prefill_seq_len=prefill_seq_len, ctx_len=ctx_len, comp_ctx_lengths=ccl, - batch_size=bs, + batch_size=input_bs, kv_cache_batch_size=(kv_cache_batch_size if self.continuous_batching else bs), full_batch_size=full_batch_size, num_speculative_tokens=k, - decode_input_batch_size=(bs if (self.continuous_batching and _is_dynamic_batch) else None), prefill_only=prefill_only, ) if decode_spec is not None: diff --git a/tests/unit_test/utils/test_auto_model_api.py b/tests/unit_test/utils/test_auto_model_api.py index cd9677ccd4..f383adfa80 100644 --- a/tests/unit_test/utils/test_auto_model_api.py +++ b/tests/unit_test/utils/test_auto_model_api.py @@ -246,16 +246,15 @@ def test_build_decode_specialization_ccl_and_num_speculative_tokens_together(sel assert result["comp_ctx_lengths"] == 64 def test_build_decode_specialization_dynamic_batch_pins_kv_at_bmax(self): - """Dynamic batching: decode input batch == decode_input_batch_size, KV batch == full_batch_size.""" + """Dynamic batching: decode input batch == batch_size, KV batch == full_batch_size.""" from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM qeff = QEFFAutoModelForCausalLM(make_tiny_gpt2(), continuous_batching=True) result = qeff.build_decode_specialization( ctx_len=128, - batch_size=2, - kv_cache_batch_size=4, + batch_size=2, # live decode input batch + kv_cache_batch_size=4, # retained KV pinned at B_max full_batch_size=4, - decode_input_batch_size=2, ) assert result is not None # Input axis follows the live batch; retained KV batch stays pinned at B_max. From 73e6cb52adc0b52bf138871cfcb682b1638b4ba0 Mon Sep 17 00:00:00 2001 From: eplatero Date: Tue, 1 Sep 2026 16:40:55 -0500 Subject: [PATCH 3/4] Expose dynamic batching execution batch size Signed-off-by: eplatero --- QEfficient/cloud/execute.py | 295 +++++++++--------- .../generation/text_generation_inference.py | 24 +- tests/cloud/test_export_compile_execute.py | 2 + tests/unit_test/utils/test_auto_model_api.py | 24 ++ tests/unit_test/utils/test_generation.py | 82 +++++ 5 files changed, 279 insertions(+), 148 deletions(-) diff --git a/QEfficient/cloud/execute.py b/QEfficient/cloud/execute.py index 09e989ea06..4e75a8ea35 100644 --- a/QEfficient/cloud/execute.py +++ b/QEfficient/cloud/execute.py @@ -1,142 +1,153 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import argparse -from typing import List, Optional - -from QEfficient.generation.text_generation_inference import cloud_ai_100_exec_kv -from QEfficient.utils import load_hf_tokenizer - - -def main( - model_name: str, - qpc_path: str, - device_group: List[int] = None, - local_model_dir: Optional[str] = None, - prompt: Optional[str] = None, # type: ignore - prompts_txt_file_path: Optional[str] = None, - generation_len: Optional[int] = None, - cache_dir: Optional[str] = None, - hf_token: Optional[str] = None, - full_batch_size: Optional[int] = None, -): - """ - Main function for the QEfficient execution CLI application. - - This function serves as the entry point for running a compiled model - (QPC package) on the Cloud AI 100 Platform. It loads the necessary - tokenizer and then orchestrates the text generation inference. - - Parameters - ---------- - model_name : str - Hugging Face Model Card name (e.g., ``gpt2``) for loading the tokenizer. - qpc_path : str - Path to the generated binary (QPC package) after compilation. - - Other Parameters - ---------------- - device_group : List[int], optional - List of device IDs to be used for inference. If `len(device_group) > 1`, - a multi-card setup is enabled. Default is None. - local_model_dir : str, optional - Path to custom model weights and config files, used if not loading tokenizer - from Hugging Face Hub. Default is None. - prompt : str, optional - Sample prompt(s) for the model text generation. For batch size > 1, - pass multiple prompts separated by a pipe (``|``) symbol. Default is None. - prompts_txt_file_path : str, optional - Path to a text file containing multiple input prompts, one per line. Default is None. - generation_len : int, optional - Maximum number of tokens to be generated during inference. Default is None. - cache_dir : str, optional - Cache directory where downloaded HuggingFace files (like tokenizer) are stored. - Default is None. - hf_token : str, optional - HuggingFace login token to access private repositories. Default is None. - full_batch_size : int, optional - Ignored in this context as continuous batching is managed by the compiled QPC. - However, it might be passed through from CLI arguments. Default is None. - - Example - ------- - To execute a compiled model from the command line: - - .. code-block:: bash - - python -m QEfficient.cloud.execute --model-name gpt2 --qpc-path /path/to/qpc/binaries --prompt "Hello world" - - For multi-device inference: - - .. code-block:: bash - - python -m QEfficient.cloud.execute --model-name gpt2 --qpc-path /path/to/qpc/binaries --device-group "[0,1]" --prompt "Hello | Hi" - - """ - tokenizer = load_hf_tokenizer( - pretrained_model_name_or_path=(local_model_dir if local_model_dir else model_name), - cache_dir=cache_dir, - hf_token=hf_token, - ) - - # Execute - cloud_ai_100_exec_kv( - tokenizer=tokenizer, - qpc_path=qpc_path, - device_id=device_group, - prompt=prompt, - prompts_txt_file_path=prompts_txt_file_path, - generation_len=generation_len, - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Execution script.") - parser.add_argument( - "--model_name", "--model-name", required=False, type=str, help="HF model card name for tokenizing the inputs" - ) - parser.add_argument("--qpc_path", "--qpc-path", required=True, help="Path to generated QPC") - parser.add_argument( - "--device_group", - "--device-group", - type=lambda device_ids: [int(x) for x in device_ids.strip("[]").split(",")], - help="Cloud AI 100 device ids (comma-separated) e.g. [0]", - ) - parser.add_argument( - "--prompt", - type=lambda prompt: prompt.split("|"), - help="Input prompt, if executing for batch size>1, pass input prompts in single string but separate with pipe (|) symbol", - ) - parser.add_argument( - "--prompts_txt_file_path", - "--prompts-txt-file-path", - type=str, - help="File path for taking input prompts from txt file, sample prompts.txt file present in examples/sample_prompts folder", - ) - parser.add_argument("--generation_len", "--generation-len", type=int, help="Number of tokens to generate") - parser.add_argument( - "--local-model-dir", "--local_model_dir", required=False, help="Path to custom model weights and config files" - ) - parser.add_argument( - "--cache-dir", - "--cache_dir", - default=None, - required=False, - help="Cache dir to store HF Downloads", - ) - parser.add_argument( - "--full_batch_size", - "--full-batch-size", - type=int, - default=None, - help="Set full batch size to enable continuous batching mode, default is None", - ) - parser.add_argument( - "--hf-token", "--hf_token", default=None, type=str, required=False, help="HF token id for private HF models" - ) - args = parser.parse_args() - main(**args.__dict__) +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +from typing import List, Optional + +from QEfficient.generation.text_generation_inference import cloud_ai_100_exec_kv +from QEfficient.utils import load_hf_tokenizer + + +def main( + model_name: str, + qpc_path: str, + device_group: List[int] = None, + local_model_dir: Optional[str] = None, + prompt: Optional[str] = None, # type: ignore + prompts_txt_file_path: Optional[str] = None, + generation_len: Optional[int] = None, + cache_dir: Optional[str] = None, + hf_token: Optional[str] = None, + full_batch_size: Optional[int] = None, + execution_batch_size: Optional[int] = None, +): + """ + Main function for the QEfficient execution CLI application. + + This function serves as the entry point for running a compiled model + (QPC package) on the Cloud AI 100 Platform. It loads the necessary + tokenizer and then orchestrates the text generation inference. + + Parameters + ---------- + model_name : str + Hugging Face Model Card name (e.g., ``gpt2``) for loading the tokenizer. + qpc_path : str + Path to the generated binary (QPC package) after compilation. + + Other Parameters + ---------------- + device_group : List[int], optional + List of device IDs to be used for inference. If `len(device_group) > 1`, + a multi-card setup is enabled. Default is None. + local_model_dir : str, optional + Path to custom model weights and config files, used if not loading tokenizer + from Hugging Face Hub. Default is None. + prompt : str, optional + Sample prompt(s) for the model text generation. For batch size > 1, + pass multiple prompts separated by a pipe (``|``) symbol. Default is None. + prompts_txt_file_path : str, optional + Path to a text file containing multiple input prompts, one per line. Default is None. + generation_len : int, optional + Maximum number of tokens to be generated during inference. Default is None. + cache_dir : str, optional + Cache directory where downloaded HuggingFace files (like tokenizer) are stored. + Default is None. + hf_token : str, optional + HuggingFace login token to access private repositories. Default is None. + full_batch_size : int, optional + Ignored in this context as continuous batching is managed by the compiled QPC. + However, it might be passed through from CLI arguments. Default is None. + execution_batch_size : int, optional + Live decode batch for continuous-batching QPCs compiled with dynamic batching. Default is None. + + Example + ------- + To execute a compiled model from the command line: + + .. code-block:: bash + + python -m QEfficient.cloud.execute --model-name gpt2 --qpc-path /path/to/qpc/binaries --prompt "Hello world" + + For multi-device inference: + + .. code-block:: bash + + python -m QEfficient.cloud.execute --model-name gpt2 --qpc-path /path/to/qpc/binaries --device-group "[0,1]" --prompt "Hello | Hi" + + """ + tokenizer = load_hf_tokenizer( + pretrained_model_name_or_path=(local_model_dir if local_model_dir else model_name), + cache_dir=cache_dir, + hf_token=hf_token, + ) + + # Execute + cloud_ai_100_exec_kv( + tokenizer=tokenizer, + qpc_path=qpc_path, + device_id=device_group, + prompt=prompt, + prompts_txt_file_path=prompts_txt_file_path, + generation_len=generation_len, + execution_batch_size=execution_batch_size, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Execution script.") + parser.add_argument( + "--model_name", "--model-name", required=False, type=str, help="HF model card name for tokenizing the inputs" + ) + parser.add_argument("--qpc_path", "--qpc-path", required=True, help="Path to generated QPC") + parser.add_argument( + "--device_group", + "--device-group", + type=lambda device_ids: [int(x) for x in device_ids.strip("[]").split(",")], + help="Cloud AI 100 device ids (comma-separated) e.g. [0]", + ) + parser.add_argument( + "--prompt", + type=lambda prompt: prompt.split("|"), + help="Input prompt, if executing for batch size>1, pass input prompts in single string but separate with pipe (|) symbol", + ) + parser.add_argument( + "--prompts_txt_file_path", + "--prompts-txt-file-path", + type=str, + help="File path for taking input prompts from txt file, sample prompts.txt file present in examples/sample_prompts folder", + ) + parser.add_argument("--generation_len", "--generation-len", type=int, help="Number of tokens to generate") + parser.add_argument( + "--execution_batch_size", + "--execution-batch-size", + type=int, + default=None, + help="Live decode batch for continuous-batching QPCs compiled with dynamic batching", + ) + parser.add_argument( + "--local-model-dir", "--local_model_dir", required=False, help="Path to custom model weights and config files" + ) + parser.add_argument( + "--cache-dir", + "--cache_dir", + default=None, + required=False, + help="Cache dir to store HF Downloads", + ) + parser.add_argument( + "--full_batch_size", + "--full-batch-size", + type=int, + default=None, + help="Set full batch size to enable continuous batching mode, default is None", + ) + parser.add_argument( + "--hf-token", "--hf_token", default=None, type=str, required=False, help="HF token id for private HF models" + ) + args = parser.parse_args() + main(**args.__dict__) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 8e0d3b3dab..b13f43c49d 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -334,6 +334,7 @@ def cloud_ai_100_exec_kv( return_pdfs: bool = False, include_guided_decoding: bool = False, sampling_params: Optional[Dict[str, Any]] = None, + execution_batch_size: Optional[int] = None, ): """ This method generates output until ``eos`` or ``generation_len`` by executing the compiled ``qpc`` on ``Cloud AI 100`` Hardware cards. @@ -355,6 +356,8 @@ def cloud_ai_100_exec_kv( :automation (bool): If true, it prints input, output, and performance stats. ``Defaults to False``. :iteration (int): Number of iterations to run the inference. ``Defaults to 1``. :prompt_to_lora_id_mapping (List[int]): Mapping to associate prompts with their respective LoRA adapter. + :execution_batch_size (int, optional): Live decode batch for continuous-batching QPCs compiled + with dynamic batching. Must match one of the compiled decode batch sizes. :include_sampler (bool, default=False): Enable/Disable sampling of next tokens. :return_pdfs (bool, default=False): Return probability distributions along with sampled next tokens. For Speculative Decoding Target Language Model, @@ -381,12 +384,9 @@ def cloud_ai_100_exec_kv( """ batch_size, ctx_len, full_batch_size = get_compilation_dims(qpc_path) + if full_batch_size is None and execution_batch_size is not None: + raise ValueError("execution_batch_size is only supported for continuous-batching (dynamic-batching) QPCs.") prompt: List[str] = get_input_prompts(prompt, prompts_txt_file_path) - prompt = fix_prompts(prompt, batch_size, full_batch_size) - if prompt_to_lora_id_mapping is not None: - prompt_to_lora_id_mapping = fix_prompt_to_lora_id_mapping( - prompt_to_lora_id_mapping, batch_size, full_batch_size - ) generate_text = TextGeneration( tokenizer=tokenizer, qpc_path=qpc_path, @@ -403,6 +403,15 @@ def cloud_ai_100_exec_kv( include_guided_decoding=include_guided_decoding, sampling_params=sampling_params, ) + if full_batch_size is not None: + resolved_execution_batch_size = generate_text._resolve_execution_batch_size(execution_batch_size) + else: + resolved_execution_batch_size = None + prompt = fix_prompts(prompt, batch_size, resolved_execution_batch_size) + if prompt_to_lora_id_mapping is not None: + prompt_to_lora_id_mapping = fix_prompt_to_lora_id_mapping( + prompt_to_lora_id_mapping, batch_size, resolved_execution_batch_size + ) for _ in range(0, int(iteration)): if full_batch_size is None: @@ -425,7 +434,10 @@ def cloud_ai_100_exec_kv( ) else: exec_info = generate_text.generate( - prompt=prompt, generation_len=generation_len, prompt_to_lora_id_mapping=prompt_to_lora_id_mapping + prompt=prompt, + generation_len=generation_len, + prompt_to_lora_id_mapping=prompt_to_lora_id_mapping, + execution_batch_size=resolved_execution_batch_size, ) print_latency_stats_kv(prompt, exec_info=exec_info, automation=automation) diff --git a/tests/cloud/test_export_compile_execute.py b/tests/cloud/test_export_compile_execute.py index 7d06871541..6805c5835c 100644 --- a/tests/cloud/test_export_compile_execute.py +++ b/tests/cloud/test_export_compile_execute.py @@ -27,6 +27,7 @@ def test_execute_calls_tokenizer_and_runtime(mocker): cache_dir="/tmp/cache", hf_token="token", full_batch_size=3, + execution_batch_size=2, ) load_hf_tokenizer.assert_called_once_with( @@ -41,4 +42,5 @@ def test_execute_calls_tokenizer_and_runtime(mocker): prompt=["My name is"], prompts_txt_file_path="examples/sample_prompts/prompts.txt", generation_len=20, + execution_batch_size=2, ) diff --git a/tests/unit_test/utils/test_auto_model_api.py b/tests/unit_test/utils/test_auto_model_api.py index f383adfa80..7f827201e7 100644 --- a/tests/unit_test/utils/test_auto_model_api.py +++ b/tests/unit_test/utils/test_auto_model_api.py @@ -20,6 +20,9 @@ All tests run on CPU only, using tiny in-memory models. """ +from pathlib import Path +from unittest.mock import patch + import pytest import torch from transformers import GPT2Config, GPT2LMHeadModel @@ -262,6 +265,27 @@ def test_build_decode_specialization_dynamic_batch_pins_kv_at_bmax(self): assert result["full_batch_size"] == 4 +@pytest.mark.cpu_only +class TestQEFFAutoModelForCausalLMGenerate: + def test_generate_forwards_execution_batch_size_to_cloud_runtime(self): + from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM + + qeff = QEFFAutoModelForCausalLM(make_tiny_gpt2()) + qeff.onnx_path = Path("/tmp/qeff/model.onnx") + qeff.qpc_path = Path("/tmp/qeff/qpc") + + with patch("QEfficient.cloud_ai_100_exec_kv", return_value=object()) as cloud_exec: + qeff.generate( + tokenizer=object(), + prompts=["one", "two"], + generation_len=4, + execution_batch_size=2, + ) + + cloud_exec.assert_called_once() + assert cloud_exec.call_args.kwargs["execution_batch_size"] == 2 + + # --------------------------------------------------------------------------- # Tests: QEFFAutoModelForCausalLM prefill toggle # --------------------------------------------------------------------------- diff --git a/tests/unit_test/utils/test_generation.py b/tests/unit_test/utils/test_generation.py index 645fb53659..8a0d898373 100644 --- a/tests/unit_test/utils/test_generation.py +++ b/tests/unit_test/utils/test_generation.py @@ -1072,6 +1072,88 @@ def test_resolve_batch_exceeding_bmax_rejected(self): gen._resolve_execution_batch_size(8) +class TestDynamicBatchPublicRuntime: + def test_cloud_exec_uses_resolved_dynamic_execution_batch_for_prompt_padding(self): + from QEfficient.generation import text_generation_inference as tgi + + exec_info = CloudAI100ExecInfo( + batch_size=2, + generated_texts=["a", "b"], + generated_ids=np.array([[1], [2]]), + perf_metrics=PerfMetrics(0.0, 0.0, 0.0, 0.0), + ) + text_generation = MagicMock() + text_generation._resolve_execution_batch_size.return_value = 2 + text_generation.generate.return_value = exec_info + + with ( + patch.object(tgi, "get_compilation_dims", return_value=(1, CTX_LEN, 4)), + patch.object(tgi, "TextGeneration", return_value=text_generation), + patch.object(tgi, "print_latency_stats_kv"), + ): + result = tgi.cloud_ai_100_exec_kv( + tokenizer=object(), + qpc_path="/fake/path/model.qpc", + prompt=["one"], + generation_len=4, + execution_batch_size=2, + ) + + assert result is exec_info + text_generation._resolve_execution_batch_size.assert_called_once_with(2) + text_generation.generate.assert_called_once() + _, kwargs = text_generation.generate.call_args + assert kwargs["prompt"] == ["one", "one"] + assert kwargs["execution_batch_size"] == 2 + + def test_cloud_exec_default_uses_largest_compiled_dynamic_batch_for_prompt_padding(self): + from QEfficient.generation import text_generation_inference as tgi + + exec_info = CloudAI100ExecInfo( + batch_size=2, + generated_texts=["a", "b"], + generated_ids=np.array([[1], [2]]), + perf_metrics=PerfMetrics(0.0, 0.0, 0.0, 0.0), + ) + text_generation = MagicMock() + text_generation._resolve_execution_batch_size.return_value = 2 + text_generation.generate.return_value = exec_info + + with ( + patch.object(tgi, "get_compilation_dims", return_value=(1, CTX_LEN, 4)), + patch.object(tgi, "TextGeneration", return_value=text_generation), + patch.object(tgi, "print_latency_stats_kv"), + ): + tgi.cloud_ai_100_exec_kv( + tokenizer=object(), + qpc_path="/fake/path/model.qpc", + prompt=["one"], + generation_len=4, + ) + + text_generation._resolve_execution_batch_size.assert_called_once_with(None) + _, kwargs = text_generation.generate.call_args + assert kwargs["prompt"] == ["one", "one"] + assert kwargs["execution_batch_size"] == 2 + + def test_cloud_exec_rejects_execution_batch_size_for_non_cb_qpc(self): + from QEfficient.generation import text_generation_inference as tgi + + with ( + patch.object(tgi, "get_compilation_dims", return_value=(1, CTX_LEN, None)), + patch.object(tgi, "TextGeneration", return_value=MagicMock()) as text_generation_cls, + ): + with pytest.raises(ValueError, match="continuous-batching"): + tgi.cloud_ai_100_exec_kv( + tokenizer=object(), + qpc_path="/fake/path/model.qpc", + prompt=["one"], + generation_len=4, + execution_batch_size=1, + ) + text_generation_cls.assert_not_called() + + # --------------------------------------------------------------------------- # Tests: _fetch_next_token_id # --------------------------------------------------------------------------- From 281eb33a06de867011204a9ab0b97c88113d8fa5 Mon Sep 17 00:00:00 2001 From: eplatero Date: Wed, 2 Sep 2026 14:18:28 -0500 Subject: [PATCH 4/4] Clean up dynamic batching tests Signed-off-by: eplatero --- .../transformers/models/modeling_auto.py | 6 ++-- tests/transformers/spd/test_spd_inference.py | 26 ++------------ .../models/test_modeling_auto_cpu.py | 35 ++++++++++--------- tests/unit_test/utils/test_auto_model_api.py | 17 --------- tests/unit_test/utils/test_generation.py | 15 +++++++- 5 files changed, 39 insertions(+), 60 deletions(-) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 6f726f001e..4d154149d1 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -4465,10 +4465,10 @@ def compile( # Normalize batch_size. A list requests dynamic batching: one decode specialization per # batch size in a single QPC. A plain int N behaves exactly as before (single-element list). _is_dynamic_batch = isinstance(batch_size, (list, tuple)) - _decode_bs = sorted(set(batch_size)) if _is_dynamic_batch else [batch_size] if _is_dynamic_batch: - if any((not isinstance(b, int)) or b < 1 for b in _decode_bs): + if any((not isinstance(b, int)) or b < 1 for b in batch_size): raise ValueError(f"All `batch_size` values must be integers >= 1, got {batch_size}.") + _decode_bs = sorted(set(batch_size)) # Dynamic batching only works on the continuous-batching export path: there the input # batch axis (`input_ids`/`position_ids`/`batch_index`) is a distinct ONNX symbol from the # retained KV-cache batch axis (`full_batch_size`), so decode specializations can vary the @@ -4490,6 +4490,8 @@ def compile( f"Every `batch_size` value must be <= `full_batch_size` (B_max={full_batch_size}); " f"got {batch_size}." ) + else: + _decode_bs = [batch_size] # Infer kv_cache_batch_size if not provided. Under continuous batching the retained KV cache # is always sized at full_batch_size (B_max) regardless of the decode input batch(es). diff --git a/tests/transformers/spd/test_spd_inference.py b/tests/transformers/spd/test_spd_inference.py index 2580f03f4b..a5ac585ca1 100644 --- a/tests/transformers/spd/test_spd_inference.py +++ b/tests/transformers/spd/test_spd_inference.py @@ -600,8 +600,7 @@ def test_multi_spec_qpc_logit_correctness(decode_ks, manual_cleanup): @pytest.mark.on_qaic @pytest.mark.feature -@pytest.mark.parametrize("batch_sizes", [[1, 2, 4]]) -def test_dynamic_batch_qpc_per_batch_execution(batch_sizes, manual_cleanup): +def test_dynamic_batch_qpc_per_batch_execution(manual_cleanup): """ Compile ONE continuous-batching QPC carrying one decode specialization per input batch size while the retained KV cache is pinned at ``full_batch_size`` (B_max), then verify each live @@ -613,6 +612,7 @@ def test_dynamic_batch_qpc_per_batch_execution(batch_sizes, manual_cleanup): that routes those ``b`` sequences into ``b`` of the ``B_max`` KV slots, and assert finite logits with batch dimension ``b``. """ + batch_sizes = [1, 2, 4] b_max = max(batch_sizes) tokenizer = AutoTokenizer.from_pretrained(_DYN_BATCH_MODEL, padding_side="right") if tokenizer.pad_token_id is None: @@ -739,25 +739,3 @@ def test_dynamic_batch_times_spec_len_compiles(manual_cleanup): ) assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) manual_cleanup([tlm.onnx_path]) - - -@pytest.mark.on_qaic -@pytest.mark.feature -def test_dynamic_batch_with_ccl_rejected(): - """Dynamic batching combined with CCL must be rejected at compile time (not deferred to the - QAIC compiler): the two vary different identifying inputs and cannot be disambiguated.""" - model = load_qeff_causal_lm_model( - _DYN_BATCH_MODEL, - num_hidden_layers=_DYN_BATCH_NUM_LAYERS, - continuous_batching=True, - ) - with pytest.raises(ValueError, match="comp_ctx_lengths"): - model.compile( - num_cores=2, - prefill_seq_len=_DYN_BATCH_PREFILL_LEN, - ctx_len=2048, - aic_enable_depth_first=True, - batch_size=[1, 2], - full_batch_size=2, - comp_ctx_lengths_decode=[1024, 2048], - ) diff --git a/tests/unit_test/models/test_modeling_auto_cpu.py b/tests/unit_test/models/test_modeling_auto_cpu.py index ab4e81796a..15cba8deeb 100644 --- a/tests/unit_test/models/test_modeling_auto_cpu.py +++ b/tests/unit_test/models/test_modeling_auto_cpu.py @@ -1333,6 +1333,17 @@ def test_list_batch_size_produces_one_decode_spec_per_batch(self): assert len(decode_specs) == 3 assert {s["batch_size"] for s in decode_specs} == {1, 2, 4} + def test_list_batch_size_normalized_to_unique_sorted_specs(self): + """Unsorted/duplicate dynamic batches emit deterministic decode specs.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + specs = self._capture_specializations( + qeff, prefill_seq_len=32, ctx_len=128, batch_size=[2, 1, 2], full_batch_size=2 + ) + decode_specs = self._decode_specs(specs) + assert {s["batch_size"] for s in decode_specs} == {1, 2} + assert [s["batch_size"] for s in decode_specs] == [1, 2] + def test_all_decode_specs_share_kv_batch_bmax(self): """Retained KV batch (full_batch_size) must be identical (=B_max) across every decode spec. @@ -1446,22 +1457,6 @@ def test_scalar_spec_len_with_ccl_still_allowed(self): assert {s["seq_len"] for s in decode_specs} == {4} assert {s["comp_ctx_lengths"] for s in decode_specs} == {1024, 2048} - def test_batch_ccl_spec_len_combo_rejected(self): - """batch_size list + CCL + num_speculative_tokens → ValueError (CCL cannot combine with a batch list).""" - model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM( - model, continuous_batching=True, qaic_config={"speculative_model_type": "target"} - ) - with pytest.raises(ValueError, match="comp_ctx_lengths"): - qeff.compile( - prefill_seq_len=32, - ctx_len=2048, - batch_size=[1, 2], - full_batch_size=2, - comp_ctx_lengths_decode=[1024, 2048], - num_speculative_tokens=[1, 3], - ) - def test_list_batch_size_rejected_without_continuous_batching(self): """batch_size list without continuous_batching → ValueError (non-CB cannot decouple axes).""" model, _ = make_tiny_llama() @@ -1482,3 +1477,11 @@ def test_batch_size_exceeding_full_batch_size_rejected(self): qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) with pytest.raises(ValueError, match="full_batch_size"): qeff.compile(prefill_seq_len=32, ctx_len=128, batch_size=[1, 2, 8], full_batch_size=4) + + @pytest.mark.parametrize("batch_size", [[0], [-1], [1, "2"]]) + def test_invalid_list_batch_size_rejected(self, batch_size): + """Dynamic batch list entries must be positive integers.""" + model, _ = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + with pytest.raises(ValueError): + qeff.compile(prefill_seq_len=32, ctx_len=128, batch_size=batch_size, full_batch_size=4) diff --git a/tests/unit_test/utils/test_auto_model_api.py b/tests/unit_test/utils/test_auto_model_api.py index 7f827201e7..497ee5599b 100644 --- a/tests/unit_test/utils/test_auto_model_api.py +++ b/tests/unit_test/utils/test_auto_model_api.py @@ -231,23 +231,6 @@ def test_build_decode_specialization_batch_size_from_kv_cache_batch_size(self): result = qeff.build_decode_specialization(ctx_len=32, batch_size=4, kv_cache_batch_size=4, full_batch_size=None) assert result["batch_size"] == 4 - def test_build_decode_specialization_ccl_and_num_speculative_tokens_together(self): - """build_decode_specialization accepts comp_ctx_lengths and num_speculative_tokens simultaneously.""" - qeff = self._make_qeff() - qeff.is_tlm = True - result = qeff.build_decode_specialization( - ctx_len=128, - batch_size=1, - kv_cache_batch_size=1, - full_batch_size=None, - comp_ctx_lengths=64, - num_speculative_tokens=3, - prefill_seq_len=32, - ) - assert result is not None - assert result["seq_len"] == 4 # k+1 - assert result["comp_ctx_lengths"] == 64 - def test_build_decode_specialization_dynamic_batch_pins_kv_at_bmax(self): """Dynamic batching: decode input batch == batch_size, KV batch == full_batch_size.""" from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM diff --git a/tests/unit_test/utils/test_generation.py b/tests/unit_test/utils/test_generation.py index 8a0d898373..df19a51033 100644 --- a/tests/unit_test/utils/test_generation.py +++ b/tests/unit_test/utils/test_generation.py @@ -581,6 +581,18 @@ def test_no_batch_index_without_full_batch_size(self): decode_inputs = obj.prepare_decode_inputs() assert "batch_index" not in decode_inputs + def test_lora_ids_use_live_decode_batch_under_dynamic_batching(self): + obj, _, _ = _make_base_instance(full_batch_size=4) + obj.initialize_decode_inputs(num_prompts=4, execution_batch_size=2, max_gen_length=10) + obj.decode_batch_size = 2 + obj.batch_index = np.arange(2).reshape(-1, 1) + obj._prompt_to_lora_id_mapping_decode = [17, 23, 31, 47] + + decode_inputs = obj.prepare_decode_inputs() + + assert decode_inputs["lora_ids"].shape == (2, 1) + np.testing.assert_array_equal(decode_inputs["lora_ids"], np.array([[17], [23]], dtype=np.int64)) + # --------------------------------------------------------------------------- # Tests: update_decode_input @@ -1068,7 +1080,8 @@ def test_resolve_uncompiled_batch_rejected(self): def test_resolve_batch_exceeding_bmax_rejected(self): _, gen = self._make_dynamic_batch_instance(decode_batches=[2, 4], full_batch_size=4) - with pytest.raises(ValueError, match="not a compiled decode batch size"): + gen._qaic_model._compiled_batch_sizes = [2, 4, 8] + with pytest.raises(ValueError, match="exceeds full_batch_size"): gen._resolve_execution_batch_size(8)