onnx-genai-example-gemma4-e2b / inference_metadata.mobius.yaml
justinchuby's picture
Ship faithful onnx-genai v1 inference_metadata.yaml (validated vs #1716 schema); preserve mobius emitter output as .mobius.yaml; add reference + generator
b5cc9f2 verified
Raw
History Blame
72.1 kB
schema_version: '1.0'
pipeline:
workflow:
manifest:
capabilities:
- workflow_ssa
- linear_effects
- nested_control_flow
- typed_emit
- emit_valid_length
- loop_induction_values
- serving_service_contract
- bounded_state_recurrence
inputs:
request.input_ids:
contract:
dtype: int64
rank: 2
shape:
- batch
- sequence_len
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: prompt_tokens
source:
kind: request
required: true
request.max_iterations:
contract:
dtype: int64
rank: 1
shape:
- 1
role:
kind: runtime
version: '1.0'
role: max_output_tokens
source:
kind: request
required: true
package.eos_ids:
contract:
dtype: int64
rank: 1
shape:
- E
role:
kind: opaque
source:
kind: literal
required: true
default: 1
package.one_token:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: literal
required: false
default: 1
package.one_step:
contract:
dtype: int64
rank: 1
shape:
- 1
role:
kind: opaque
source:
kind: literal
required: false
default: 1
package.max_context:
contract:
dtype: int64
rank: 1
shape:
- 1
role:
kind: opaque
source:
kind: literal
required: false
default: 131072
request.prompt_lengths:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: application
name: prompt_lengths
required: false
default: -1
request.eos_ids:
contract:
dtype: int64
rank: 2
shape:
- batch
- num_eos
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: application
name: eos_ids
required: false
default: 1
request.eos_lengths:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: application
name: eos_lengths
required: false
default: 1
request.row_max_iterations:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: application
name: row_max_iterations
required: false
default: -1
request.temperature:
contract:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: sampling_temperature
source:
kind: request
required: false
default: 1.0
request.top_k:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: sampling_top_k
source:
kind: request
required: false
default: 1
request.top_p:
contract:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: sampling_top_p
source:
kind: request
required: false
default: 1.0
request.min_p:
contract:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: sampling_min_p
source:
kind: request
required: false
default: 0.0
request.seed:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: runtime
version: '1.0'
role: seed
source:
kind: request
required: false
default: 0
request.rng_counter:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: application
name: rng_counter
required: false
default: 0
package.active:
contract:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: literal
required: false
default: true
package.not_done:
contract:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: literal
required: false
default: false
package.cache_lengths:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: literal
required: false
default: 0
package.zero_batch:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
role:
kind: opaque
source:
kind: literal
required: false
default: 0
outputs:
tokens:
contract:
dtype: int64
rank: 2
shape:
- batch
- generated_sequence
batch_layout:
kind: request_aligned
axis: 0
role: tokens
stage: pre_adapter
components:
model:
implementation:
kind: onnx
artifact: model.onnx
ports:
roles:
input_ids: token_ids
attention_mask: attention_mask
logits: logits
token_sampler:
implementation:
kind: onnx
artifact: policies/token_sampler.onnx
ports:
inputs:
logits:
dtype: float32
rank: 2
shape:
- batch
- vocabulary
batch_layout:
kind: request_aligned
axis: 0
temperature:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
top_k:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
top_p:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
min_p:
dtype: float32
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
seed:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
counter:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
done:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
token:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
next_counter:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
contract:
id: onnx-genai.token-sampler
version: '2'
bindings:
logits: logits
token: token
temperature: temperature
top_k: top_k
top_p: top_p
min_p: min_p
active: active
done: done
seed: seed
counter: counter
next_counter: next_counter
parameters:
mode: seeded_stochastic
batching: per_row
inactive_rows: preserve
application_overridable: true
termination:
implementation:
kind: onnx
artifact: policies/termination.onnx
ports:
inputs:
tokens:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
eos_ids:
dtype: int64
rank: 2
shape:
- batch
- num_eos
batch_layout:
kind: request_aligned
axis: 0
eos_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
iteration:
dtype: int64
rank: 1
shape:
- 1
max_iterations:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
done:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
next_active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
continue:
dtype: bool
rank: 1
shape:
- 1
contract:
id: onnx-genai.termination-predicate
version: '2'
bindings:
tokens: tokens
eos_ids: eos_ids
iteration: iteration
max_iterations: max_iterations
eos_lengths: eos_lengths
active: active
done: done
next_active: next_active
continue: continue
parameters:
batching: per_row
inactive_rows: preserve
token_state_update:
implementation:
kind: onnx
artifact: policies/token_state_update.onnx
ports:
inputs:
current:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
update:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
done:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
next:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
contract:
id: onnx-genai.state-update
version: '2'
bindings:
current: current
update: update
active: active
done: done
next: next
parameters:
batching: per_row
inactive_rows: preserve
last_token_logits:
implementation:
kind: onnx
artifact: policies/last_token_logits.onnx
ports:
inputs:
logits:
dtype: float16
rank: 3
shape:
- batch
- sequence
- vocabulary
batch_layout:
kind: request_aligned
axis: 0
outputs:
last_logits:
dtype: float32
rank: 2
shape:
- batch
- vocabulary
batch_layout:
kind: request_aligned
axis: 0
decoder_state_initializer:
implementation:
kind: onnx
artifact: policies/decoder_state_initializer.onnx
ports:
inputs:
prompt_tokens:
dtype: int64
rank: 2
shape:
- batch
- prompt_sequence
batch_layout:
kind: request_aligned
axis: 0
prompt_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
max_iterations:
dtype: int64
rank: 1
shape:
- 1
outputs:
attention_mask:
dtype: int64
rank: 2
shape:
- batch
- capacity
batch_layout:
kind: request_aligned
axis: 0
body_attention_mask:
dtype: int64
rank: 2
shape:
- batch
- capacity
batch_layout:
kind: request_aligned
axis: 0
token_slot:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
generated_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
cache_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
past_key_values.0.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.0.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.1.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.1.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.10.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.10.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.11.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.11.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.12.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.12.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.13.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.13.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.14.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
past_key_values.14.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
past_key_values.2.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.2.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.3.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.3.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.4.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
past_key_values.4.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
past_key_values.5.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.5.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.6.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.6.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.7.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.7.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.8.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.8.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 256
batch_layout:
kind: request_aligned
axis: 0
past_key_values.9.key:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
past_key_values.9.value:
dtype: float16
rank: 4
shape:
- batch
- 1
- capacity
- 512
batch_layout:
kind: request_aligned
axis: 0
decoder_step_update:
implementation:
kind: onnx
artifact: policies/decoder_step_update.onnx
ports:
inputs:
attention_mask:
dtype: int64
rank: 2
shape:
- batch
- context
batch_layout:
kind: request_aligned
axis: 0
logical_length:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
next_attention_mask:
dtype: int64
rank: 2
shape:
- batch
- context
batch_layout:
kind: request_aligned
axis: 0
cache_length_update:
implementation:
kind: onnx
artifact: policies/cache_length_update.onnx
ports:
inputs:
left:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
right:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
done:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
total:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
termination_batch_initializer:
implementation:
kind: onnx
artifact: policies/termination_batch_initializer.onnx
ports:
inputs:
input_eos_ids:
dtype: int64
rank: 2
shape:
- batch
- num_eos
batch_layout:
kind: request_aligned
axis: 0
input_eos_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
input_max_iterations:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
fallback_max_iterations:
dtype: int64
rank: 1
shape:
- 1
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
row_eos_ids:
dtype: int64
rank: 2
shape:
- batch
- num_eos
batch_layout:
kind: request_aligned
axis: 0
eos_lengths:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
max_iterations:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
token_to_slot:
implementation:
kind: onnx
artifact: policies/token_to_slot.onnx
ports:
inputs:
token:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
slot:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
generated_length_update:
implementation:
kind: onnx
artifact: policies/generated_length_update.onnx
ports:
inputs:
left:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
right:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
active:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
done:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
outputs:
total:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
state:
token:
contract:
dtype: int64
rank: 2
shape:
- batch
- 1
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: initializer.token_slot
recurrence:
kind: invariant
logits:
contract:
dtype: float32
rank: 2
shape:
- batch
- 262144
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.last_logits
recurrence:
kind: invariant
generated_lengths:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
class: semantic
scope: invocation
initializer: initializer.generated_lengths
recurrence:
kind: invariant
active:
contract:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
class: semantic
scope: invocation
initializer: package.active
recurrence:
kind: invariant
done:
contract:
dtype: bool
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
class: semantic
scope: invocation
initializer: package.not_done
recurrence:
kind: invariant
accepted_len:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
class: semantic
scope: invocation
initializer: package.zero_batch
recurrence:
kind: invariant
cache_lengths:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
class: semantic
scope: invocation
initializer: initializer.cache_lengths
recurrence:
kind: invariant
rng_counter:
contract:
dtype: int64
rank: 1
shape:
- batch
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
class: semantic
initializer: request.rng_counter
recurrence:
kind: invariant
attention_mask:
contract:
dtype: int64
rank: 2
shape:
- batch
- context
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: initializer.attention_mask
recurrence:
kind: invariant
cache_0:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.0.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_1:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.0.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_2:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.1.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_3:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.1.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_4:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.2.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_5:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.2.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_6:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.3.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_7:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.3.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_8:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.4.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
cache_9:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.4.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
cache_10:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.5.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_11:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.5.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_12:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.6.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_13:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.6.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_14:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.7.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_15:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.7.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_16:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.8.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_17:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.8.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_18:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.9.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
cache_19:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.9.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
cache_20:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.10.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_21:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.10.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_22:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.11.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_23:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.11.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_24:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.12.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_25:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.12.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_26:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.13.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_27:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 256
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.13.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_sliding_attention
cache_28:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.14.key
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
cache_29:
contract:
dtype: float16
rank: 4
shape:
- batch
- 1
- past_sequence_len
- 512
batch_layout:
kind: request_aligned
axis: 0
scope: invocation
initializer: decoder.setup.present.14.value
recurrence:
kind: bounded
axis: 2
max: package.max_context
management: runtime
release_boundary: invocation
service_group: decoder_cache_full_attention
serving:
active: active
done: done
accepted_len: accepted_len
state_service:
groups:
decoder_cache_full_attention:
kind: full_attention
sequence_axis: 2
layout: bnsh
logical_lengths: cache_lengths
aliasing: permitted
reuse:
prefix_reusable: true
evictable_prefix: false
ports:
model:
cache_8:
input: past_key_values.4.key
output: present.4.key
role: key
layer: 4
cache_9:
input: past_key_values.4.value
output: present.4.value
role: value
layer: 4
cache_18:
input: past_key_values.9.key
output: present.9.key
role: key
layer: 9
cache_19:
input: past_key_values.9.value
output: present.9.value
role: value
layer: 9
cache_28:
input: past_key_values.14.key
output: present.14.key
role: key
layer: 14
cache_29:
input: past_key_values.14.value
output: present.14.value
role: value
layer: 14
decoder_cache_sliding_attention:
kind: sliding_attention
sequence_axis: 2
layout: bnsh
logical_lengths: cache_lengths
aliasing: permitted
reuse:
prefix_reusable: true
evictable_prefix: true
ports:
model:
cache_0:
input: past_key_values.0.key
output: present.0.key
role: key
layer: 0
cache_1:
input: past_key_values.0.value
output: present.0.value
role: value
layer: 0
cache_2:
input: past_key_values.1.key
output: present.1.key
role: key
layer: 1
cache_3:
input: past_key_values.1.value
output: present.1.value
role: value
layer: 1
cache_4:
input: past_key_values.2.key
output: present.2.key
role: key
layer: 2
cache_5:
input: past_key_values.2.value
output: present.2.value
role: value
layer: 2
cache_6:
input: past_key_values.3.key
output: present.3.key
role: key
layer: 3
cache_7:
input: past_key_values.3.value
output: present.3.value
role: value
layer: 3
cache_10:
input: past_key_values.5.key
output: present.5.key
role: key
layer: 5
cache_11:
input: past_key_values.5.value
output: present.5.value
role: value
layer: 5
cache_12:
input: past_key_values.6.key
output: present.6.key
role: key
layer: 6
cache_13:
input: past_key_values.6.value
output: present.6.value
role: value
layer: 6
cache_14:
input: past_key_values.7.key
output: present.7.key
role: key
layer: 7
cache_15:
input: past_key_values.7.value
output: present.7.value
role: value
layer: 7
cache_16:
input: past_key_values.8.key
output: present.8.key
role: key
layer: 8
cache_17:
input: past_key_values.8.value
output: present.8.value
role: value
layer: 8
cache_20:
input: past_key_values.10.key
output: present.10.key
role: key
layer: 10
cache_21:
input: past_key_values.10.value
output: present.10.value
role: value
layer: 10
cache_22:
input: past_key_values.11.key
output: present.11.key
role: key
layer: 11
cache_23:
input: past_key_values.11.value
output: present.11.value
role: value
layer: 11
cache_24:
input: past_key_values.12.key
output: present.12.key
role: key
layer: 12
cache_25:
input: past_key_values.12.value
output: present.12.value
role: value
layer: 12
cache_26:
input: past_key_values.13.key
output: present.13.key
role: key
layer: 13
cache_27:
input: past_key_values.13.value
output: present.13.value
role: value
layer: 13
steps:
- kind: loop
setup:
- kind: invoke
component: decoder_state_initializer
inputs:
prompt_tokens: request.input_ids
prompt_lengths: request.prompt_lengths
max_iterations: request.max_iterations
outputs:
attention_mask: initializer.attention_mask
body_attention_mask: initializer.body_attention_mask
token_slot: initializer.token_slot
generated_lengths: initializer.generated_lengths
cache_lengths: initializer.cache_lengths
past_key_values.0.key: initializer.past_key_values.0.key
past_key_values.0.value: initializer.past_key_values.0.value
past_key_values.1.key: initializer.past_key_values.1.key
past_key_values.1.value: initializer.past_key_values.1.value
past_key_values.10.key: initializer.past_key_values.10.key
past_key_values.10.value: initializer.past_key_values.10.value
past_key_values.11.key: initializer.past_key_values.11.key
past_key_values.11.value: initializer.past_key_values.11.value
past_key_values.12.key: initializer.past_key_values.12.key
past_key_values.12.value: initializer.past_key_values.12.value
past_key_values.13.key: initializer.past_key_values.13.key
past_key_values.13.value: initializer.past_key_values.13.value
past_key_values.14.key: initializer.past_key_values.14.key
past_key_values.14.value: initializer.past_key_values.14.value
past_key_values.2.key: initializer.past_key_values.2.key
past_key_values.2.value: initializer.past_key_values.2.value
past_key_values.3.key: initializer.past_key_values.3.key
past_key_values.3.value: initializer.past_key_values.3.value
past_key_values.4.key: initializer.past_key_values.4.key
past_key_values.4.value: initializer.past_key_values.4.value
past_key_values.5.key: initializer.past_key_values.5.key
past_key_values.5.value: initializer.past_key_values.5.value
past_key_values.6.key: initializer.past_key_values.6.key
past_key_values.6.value: initializer.past_key_values.6.value
past_key_values.7.key: initializer.past_key_values.7.key
past_key_values.7.value: initializer.past_key_values.7.value
past_key_values.8.key: initializer.past_key_values.8.key
past_key_values.8.value: initializer.past_key_values.8.value
past_key_values.9.key: initializer.past_key_values.9.key
past_key_values.9.value: initializer.past_key_values.9.value
- kind: invoke
component: model
inputs:
input_ids: request.input_ids
past_key_values.0.key: initializer.past_key_values.0.key
past_key_values.0.value: initializer.past_key_values.0.value
past_key_values.1.key: initializer.past_key_values.1.key
past_key_values.1.value: initializer.past_key_values.1.value
past_key_values.2.key: initializer.past_key_values.2.key
past_key_values.2.value: initializer.past_key_values.2.value
past_key_values.3.key: initializer.past_key_values.3.key
past_key_values.3.value: initializer.past_key_values.3.value
past_key_values.4.key: initializer.past_key_values.4.key
past_key_values.4.value: initializer.past_key_values.4.value
past_key_values.5.key: initializer.past_key_values.5.key
past_key_values.5.value: initializer.past_key_values.5.value
past_key_values.6.key: initializer.past_key_values.6.key
past_key_values.6.value: initializer.past_key_values.6.value
past_key_values.7.key: initializer.past_key_values.7.key
past_key_values.7.value: initializer.past_key_values.7.value
past_key_values.8.key: initializer.past_key_values.8.key
past_key_values.8.value: initializer.past_key_values.8.value
past_key_values.9.key: initializer.past_key_values.9.key
past_key_values.9.value: initializer.past_key_values.9.value
past_key_values.10.key: initializer.past_key_values.10.key
past_key_values.10.value: initializer.past_key_values.10.value
past_key_values.11.key: initializer.past_key_values.11.key
past_key_values.11.value: initializer.past_key_values.11.value
past_key_values.12.key: initializer.past_key_values.12.key
past_key_values.12.value: initializer.past_key_values.12.value
past_key_values.13.key: initializer.past_key_values.13.key
past_key_values.13.value: initializer.past_key_values.13.value
past_key_values.14.key: initializer.past_key_values.14.key
past_key_values.14.value: initializer.past_key_values.14.value
attention_mask: initializer.attention_mask
outputs:
logits: decoder.setup.logits
present.0.key: decoder.setup.present.0.key
present.0.value: decoder.setup.present.0.value
present.1.key: decoder.setup.present.1.key
present.1.value: decoder.setup.present.1.value
present.2.key: decoder.setup.present.2.key
present.2.value: decoder.setup.present.2.value
present.3.key: decoder.setup.present.3.key
present.3.value: decoder.setup.present.3.value
present.4.key: decoder.setup.present.4.key
present.4.value: decoder.setup.present.4.value
present.5.key: decoder.setup.present.5.key
present.5.value: decoder.setup.present.5.value
present.6.key: decoder.setup.present.6.key
present.6.value: decoder.setup.present.6.value
present.7.key: decoder.setup.present.7.key
present.7.value: decoder.setup.present.7.value
present.8.key: decoder.setup.present.8.key
present.8.value: decoder.setup.present.8.value
present.9.key: decoder.setup.present.9.key
present.9.value: decoder.setup.present.9.value
present.10.key: decoder.setup.present.10.key
present.10.value: decoder.setup.present.10.value
present.11.key: decoder.setup.present.11.key
present.11.value: decoder.setup.present.11.value
present.12.key: decoder.setup.present.12.key
present.12.value: decoder.setup.present.12.value
present.13.key: decoder.setup.present.13.key
present.13.value: decoder.setup.present.13.value
present.14.key: decoder.setup.present.14.key
present.14.value: decoder.setup.present.14.value
- kind: invoke
component: termination_batch_initializer
inputs:
input_eos_ids: request.eos_ids
input_eos_lengths: request.eos_lengths
input_max_iterations: request.row_max_iterations
fallback_max_iterations: request.max_iterations
active: package.active
outputs:
row_eos_ids: termination.eos_ids
eos_lengths: termination.eos_lengths
max_iterations: termination.max_iterations
- kind: invoke
component: last_token_logits
inputs:
logits: decoder.setup.logits
outputs:
last_logits: decoder.setup.last_logits
steps:
- kind: invoke
component: token_sampler
inputs:
logits: logits
temperature: request.temperature
top_k: request.top_k
top_p: request.top_p
min_p: request.min_p
seed: request.seed
counter: rng_counter
active: active
done: done
outputs:
token: sample.body
next_counter: sample.next_counter
- kind: invoke
component: token_to_slot
inputs:
token: sample.body
outputs:
slot: sample.slot
- kind: invoke
component: generated_length_update
inputs:
left: generated_lengths
right: package.one_token
active: active
done: done
outputs:
total: token.next_lengths
- kind: invoke
component: generated_length_update
inputs:
left: package.zero_batch
right: package.one_token
active: active
done: done
outputs:
total: token.emitted_length
- kind: invoke
component: token_state_update
inputs:
current: token
update: sample.slot
active: active
done: done
outputs:
next: token.body
- kind: invoke
component: termination
inputs:
tokens: sample.body
eos_ids: termination.eos_ids
eos_lengths: termination.eos_lengths
iteration: loop.iteration
max_iterations: termination.max_iterations
active: active
outputs:
done: loop.done
continue: loop.continue
next_active: loop.next_active
- kind: invoke
component: cache_length_update
inputs:
left: cache_lengths
right: package.one_token
active: active
done: done
outputs:
total: cache_lengths.next
- kind: invoke
component: cache_length_update
inputs:
left: package.zero_batch
right: package.one_token
active: active
done: done
outputs:
total: accepted_len.next
- kind: emit
value: token.body
output: tokens
mode: append
valid_length: token.emitted_length
when: active
- kind: invoke
component: decoder_step_update
inputs:
attention_mask: attention_mask
logical_length: cache_lengths
outputs:
next_attention_mask: decoder_step.body_attention_mask
- kind: invoke
component: model
inputs:
input_ids: token.body
past_key_values.0.key: cache_0
past_key_values.0.value: cache_1
past_key_values.1.key: cache_2
past_key_values.1.value: cache_3
past_key_values.2.key: cache_4
past_key_values.2.value: cache_5
past_key_values.3.key: cache_6
past_key_values.3.value: cache_7
past_key_values.4.key: cache_8
past_key_values.4.value: cache_9
past_key_values.5.key: cache_10
past_key_values.5.value: cache_11
past_key_values.6.key: cache_12
past_key_values.6.value: cache_13
past_key_values.7.key: cache_14
past_key_values.7.value: cache_15
past_key_values.8.key: cache_16
past_key_values.8.value: cache_17
past_key_values.9.key: cache_18
past_key_values.9.value: cache_19
past_key_values.10.key: cache_20
past_key_values.10.value: cache_21
past_key_values.11.key: cache_22
past_key_values.11.value: cache_23
past_key_values.12.key: cache_24
past_key_values.12.value: cache_25
past_key_values.13.key: cache_26
past_key_values.13.value: cache_27
past_key_values.14.key: cache_28
past_key_values.14.value: cache_29
attention_mask: decoder_step.body_attention_mask
outputs:
logits: decoder.body.logits
present.0.key: decoder.body.present.0.key
present.0.value: decoder.body.present.0.value
present.1.key: decoder.body.present.1.key
present.1.value: decoder.body.present.1.value
present.2.key: decoder.body.present.2.key
present.2.value: decoder.body.present.2.value
present.3.key: decoder.body.present.3.key
present.3.value: decoder.body.present.3.value
present.4.key: decoder.body.present.4.key
present.4.value: decoder.body.present.4.value
present.5.key: decoder.body.present.5.key
present.5.value: decoder.body.present.5.value
present.6.key: decoder.body.present.6.key
present.6.value: decoder.body.present.6.value
present.7.key: decoder.body.present.7.key
present.7.value: decoder.body.present.7.value
present.8.key: decoder.body.present.8.key
present.8.value: decoder.body.present.8.value
present.9.key: decoder.body.present.9.key
present.9.value: decoder.body.present.9.value
present.10.key: decoder.body.present.10.key
present.10.value: decoder.body.present.10.value
present.11.key: decoder.body.present.11.key
present.11.value: decoder.body.present.11.value
present.12.key: decoder.body.present.12.key
present.12.value: decoder.body.present.12.value
present.13.key: decoder.body.present.13.key
present.13.value: decoder.body.present.13.value
present.14.key: decoder.body.present.14.key
present.14.value: decoder.body.present.14.value
- kind: invoke
component: last_token_logits
inputs:
logits: decoder.body.logits
outputs:
last_logits: decoder.body.last_logits
continue_when: active
max_iterations: request.max_iterations
carried:
- cell: token
next: token.body
- cell: logits
next: decoder.body.last_logits
- cell: generated_lengths
next: token.next_lengths
- cell: active
next: loop.next_active
- cell: done
next: loop.done
- cell: cache_lengths
next: cache_lengths.next
- cell: accepted_len
next: accepted_len.next
- cell: rng_counter
next: sample.next_counter
- cell: attention_mask
next: decoder_step.body_attention_mask
- cell: cache_0
next: decoder.body.present.0.key
- cell: cache_1
next: decoder.body.present.0.value
- cell: cache_2
next: decoder.body.present.1.key
- cell: cache_3
next: decoder.body.present.1.value
- cell: cache_4
next: decoder.body.present.2.key
- cell: cache_5
next: decoder.body.present.2.value
- cell: cache_6
next: decoder.body.present.3.key
- cell: cache_7
next: decoder.body.present.3.value
- cell: cache_8
next: decoder.body.present.4.key
- cell: cache_9
next: decoder.body.present.4.value
- cell: cache_10
next: decoder.body.present.5.key
- cell: cache_11
next: decoder.body.present.5.value
- cell: cache_12
next: decoder.body.present.6.key
- cell: cache_13
next: decoder.body.present.6.value
- cell: cache_14
next: decoder.body.present.7.key
- cell: cache_15
next: decoder.body.present.7.value
- cell: cache_16
next: decoder.body.present.8.key
- cell: cache_17
next: decoder.body.present.8.value
- cell: cache_18
next: decoder.body.present.9.key
- cell: cache_19
next: decoder.body.present.9.value
- cell: cache_20
next: decoder.body.present.10.key
- cell: cache_21
next: decoder.body.present.10.value
- cell: cache_22
next: decoder.body.present.11.key
- cell: cache_23
next: decoder.body.present.11.value
- cell: cache_24
next: decoder.body.present.12.key
- cell: cache_25
next: decoder.body.present.12.value
- cell: cache_26
next: decoder.body.present.13.key
- cell: cache_27
next: decoder.body.present.13.value
- cell: cache_28
next: decoder.body.present.14.key
- cell: cache_29
next: decoder.body.present.14.value
termination: generation_eos
iteration:
value: loop.iteration
contract:
dtype: int64
rank: 1
shape:
- 1