[Model builder] Add option to exclude cache in inputs and outputs

[xenova]

In certain cases (e.g., single-round conversations), it is not necessary to require past_key_values as inputs and present outputs, like with https://huggingface.co/livekit/turn-detector (and its usage here).

So, this PR adds an option to exclude these inputs and outputs from the graph.

Example usage:

python builder.py -m livekit/turn-detector -o converted -e cpu -p fp32 --extra_options exclude_cache=True

Output graph signature:

For comparison, graph signature w/ cache IO:

[xenova]

@microsoft-github-policy-service agree [company="{Hugging Face}"]

[xenova]

@microsoft-github-policy-service agree company="Hugging Face"

[xenova]

This modification only works for models w/ GQA. Maybe someone with a bit more experience with the model builder could help get it working for models w/ MHA? 😇

[ambroser53]

Does this give a memory overhead improvement?

[xenova]

Does this give a memory overhead improvement?

Indeed - not having to store and pass these values back has improved execution time in my tests.

[kunal-vaishnavi]

Suggested change

	exclude_cache = Remove cache inputs and outputs from your ONNX model.
	exclude_kv_cache = Remove KV cache inputs and outputs from your ONNX model.

[kunal-vaishnavi]

Suggested change

	Use this option when you want to remove the `past_key_values` inputs and `present` outputs from within your ONNX model.
	Use this option when you want to remove the `past_key_values` inputs and `present` outputs from within your ONNX model.
	Note that this should be used when you want to run ONNX models with ONNX Runtime only. ONNX Runtime GenAI requires the KV cache inputs and outputs for inference.

[kunal-vaishnavi]

Suggested change

	self.exclude_cache = "exclude_cache" in extra_options
	self.exclude_cache = extra_options.get("exclude_cache", False)

[kunal-vaishnavi]

I think present_k and present_v should be empty strings since the KV cache inputs and outputs are not in the model.

[kunal-vaishnavi]

Thanks for the contribution! Can you also update the following places?

1. Add the option in the validation of the boolean extra options

onnxruntime-genai/src/python/py/models/builder.py

Lines 3133 to 3141 in 44e541e

	bools = ["int4_is_symmetric", "exclude_embeds", "exclude_lm_head", "include_hidden_states", "enable_cuda_graph", "use_8bits_moe", "use_qdq", "include_prompt_templates"]
	for key in bools:
	if key in kv_pairs:
	if kv_pairs[key] in {"false", "False", "0"}:
	kv_pairs[key] = False
	elif kv_pairs[key] in {"true", "True", "1"}:
	kv_pairs[key] = True
	else:
	raise ValueError(f"{key} must be false/False/0 or true/True/1.")

2. Add the option in the README and provide a usage example

For MultiHeadAttention where the repeat KV operation occurs, you may need to remove the Concat node that combines past_kv and curr_kv to produce present_kv.

onnxruntime-genai/src/python/py/models/builder.py

Lines 1287 to 1304 in 44e541e

	# Make the initial subgraph
	#
	# +------> Gather --> Unsqueeze -----+
	# | |
	# past_kv +------> Gather --> Unsqueeze -----+---> Mul --> Concat (4D)
	# | | |
	# root_input --> Reshape --> Transpose --> Concat --> Shape ---> Gather --> Unsqueeze -----+---> Concat (5D)
	# | | |
	# present_kv +------> Gather --> Unsqueeze -----+
	reshape_1_name = f"{basename}/Reshape_1"
	reshape_1_inputs = [root_input, f"/model/constants/TensorProto.INT64/1D/0, 0, {self.num_kv_heads}, -1"]
	self.make_reshape(reshape_1_name, reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length', self.num_kv_heads, self.head_size])
	transpose_1_name = f"{basename}/Transpose_1"
	transpose_1_input = f"{reshape_1_name}/output_0"
	self.make_transpose(transpose_1_name, transpose_1_input, dtype=self.io_dtype, shape=['batch_size', self.num_kv_heads, 'sequence_length', self.head_size], perm=[0,2,1,3])
	concat_1_name = f"{basename}/Concat_1"
	concat_1_inputs = [past_kv, f"{transpose_1_name}/output_0"]
	self.make_node("Concat", inputs=concat_1_inputs, outputs=[present_kv], name=concat_1_name, axis=2)