-
Notifications
You must be signed in to change notification settings - Fork 2
Expand L4/L5 test coverage: seq2seq generation, golden data, ci_skip_reason #186
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,8 +3,9 @@ | |
|
|
||
| """Autoregressive text generation using ONNX Runtime. | ||
|
|
||
| Provides a self-contained generation loop with KV cache management, | ||
| without depending on onnxruntime-genai. | ||
| Provides self-contained generation loops with KV cache management, | ||
| without depending on onnxruntime-genai. Supports causal-LM (decoder- | ||
| only) and seq2seq (encoder-decoder) architectures. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
@@ -164,6 +165,130 @@ def generate( | |
| return all_ids | ||
|
|
||
|
|
||
| class OnnxSeq2SeqGenerator: | ||
| """Greedy generation for encoder-decoder (seq2seq) ONNX models. | ||
|
|
||
| Runs the encoder once, then autoregressively decodes using the | ||
| decoder with cross-attention to encoder hidden states. Manages | ||
| both self-attention and cross-attention KV caches. | ||
|
|
||
| Example:: | ||
|
|
||
| enc_session = OnnxModelSession(pkg["encoder"]) | ||
| dec_session = OnnxModelSession(pkg["decoder"]) | ||
| gen = OnnxSeq2SeqGenerator(enc_session, dec_session, config) | ||
| output_ids = gen.generate(input_ids, max_new_tokens=20) | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| enc_session: OnnxModelSession, | ||
| dec_session: OnnxModelSession, | ||
| config: ArchitectureConfig, | ||
| ): | ||
| self.enc_session = enc_session | ||
| self.dec_session = dec_session | ||
| self.config = config | ||
|
|
||
| def generate( | ||
| self, | ||
| input_ids: np.ndarray, | ||
| max_new_tokens: int = 20, | ||
| eos_token_id: int | None = None, | ||
| decoder_start_token_id: int = 0, | ||
| ) -> np.ndarray: | ||
| """Generate tokens from encoder input using greedy decoding. | ||
|
|
||
| Args: | ||
| input_ids: [batch, src_seq_len] int64 encoder input tokens. | ||
| max_new_tokens: Maximum number of tokens to generate. | ||
| eos_token_id: If set, stop when this token is produced. | ||
| decoder_start_token_id: Token to seed the decoder. | ||
|
|
||
| Returns: | ||
| [batch, generated_len] int64 array of generated token IDs | ||
| (decoder start token + generated, no encoder input). | ||
| """ | ||
| batch_size = input_ids.shape[0] | ||
| src_seq_len = input_ids.shape[1] | ||
|
|
||
| # Step 1: Run encoder once | ||
| enc_feeds = { | ||
| "input_ids": input_ids, | ||
| "attention_mask": np.ones_like(input_ids), | ||
| } | ||
| enc_outputs = self.enc_session.run(enc_feeds) | ||
|
|
||
| # Extract encoder hidden states | ||
| enc_hidden = None | ||
| for key in ("encoder_hidden_states", "last_hidden_state"): | ||
| if key in enc_outputs: | ||
| enc_hidden = enc_outputs[key] | ||
| break | ||
| if enc_hidden is None: | ||
| raise KeyError( | ||
| f"Encoder output missing hidden states. Keys: {sorted(enc_outputs.keys())}" | ||
| ) | ||
|
|
||
| # Step 2: Initialize decoder KV caches | ||
| num_kv_heads = self.config.num_key_value_heads | ||
| head_dim = self.config.head_dim | ||
| past_kv: dict[str, np.ndarray] = {} | ||
|
|
||
| for name in self.dec_session.input_names: | ||
| if not name.startswith("past_key_values."): | ||
| continue | ||
| if ".cross." in name: | ||
| # Cross-attention cache: starts empty, populated on first step | ||
| past_kv[name] = np.zeros( | ||
| (batch_size, num_kv_heads, 0, head_dim), | ||
| dtype=np.float32, | ||
| ) | ||
| else: | ||
| # Self-attention cache: grows each step | ||
| past_kv[name] = np.zeros( | ||
| (batch_size, num_kv_heads, 0, head_dim), | ||
| dtype=np.float32, | ||
| ) | ||
|
|
||
| # Step 3: Autoregressive decode loop | ||
| cur_dec_ids = np.full((batch_size, 1), decoder_start_token_id, dtype=np.int64) | ||
| all_ids = cur_dec_ids.copy() | ||
|
|
||
| for _step in range(max_new_tokens): | ||
| dec_feeds: dict[str, np.ndarray] = { | ||
| "input_ids": cur_dec_ids, | ||
| "encoder_hidden_states": enc_hidden, | ||
| "attention_mask": np.ones((batch_size, src_seq_len), dtype=np.int64), | ||
| **past_kv, | ||
| } | ||
|
Comment on lines
+258
to
+264
|
||
|
|
||
| outputs = self.dec_session.run(dec_feeds) | ||
|
|
||
| # Extract logits and take argmax of last token | ||
| logits = outputs["logits"] # [batch, dec_seq_len, vocab] | ||
| next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True) | ||
|
|
||
| all_ids = np.concatenate([all_ids, next_token], axis=1) | ||
|
|
||
| # Check EOS | ||
| if eos_token_id is not None and np.all(next_token == eos_token_id): | ||
| break | ||
|
|
||
| # Update KV caches from present outputs | ||
| for name in list(past_kv.keys()): | ||
| # past_key_values.N.self.key → present.N.self.key | ||
| layer_suffix = name.replace("past_key_values.", "") | ||
| present_name = f"present.{layer_suffix}" | ||
| if present_name in outputs: | ||
| past_kv[name] = outputs[present_name] | ||
|
|
||
| # Next step: only the new token | ||
| cur_dec_ids = next_token.astype(np.int64) | ||
|
|
||
| return all_ids | ||
|
|
||
|
|
||
| def torch_generate_greedy( | ||
| model, | ||
| input_ids: np.ndarray, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
KV cache initialization in
OnnxSeq2SeqGeneratorhard-codes(batch_size, config.num_key_value_heads, 0, config.head_dim)for everypast_key_values.*input. HoweverSeq2SeqTaskcreates caches usingnum_attention_heads(and dtype from the model), and naming is split into.self.and.cross.. To avoid head-count mismatches and future shape variations, initialize caches fromdec_session.get_input_shape(name)(like_make_empty_kv_cache()does intests/e2e_golden_test.py) and only override the sequence dimension (0 for self; encoder seq len for cross if needed).