[BUG] FlexMDM predict() includes suffix decode loop in timed path

[dhruvdcoder]

Describe the bug
FlexMDMPredictor.predict() calls decode_generation_suffix() at the end of generation and folds that work into time_taken. That helper walks each batch row in Python and runs per-sample tokenizer.decode over the non-fixed region.

Suffix text extraction is not part of generation. It exists so TinyGSM / post-hoc evaluators can read generated_text. Putting it inside predict() makes FlexMDM look slower than peers when comparing predict() latency, and mixes logging/eval work into the timed generation path.

to_dict() already has a fallback that can rebuild generated_text from preds["ids"] + batch["fixed"] when it is missing from preds.

Relevant code:

• decode_generation_suffix and the call at the end of predict() in xlm-models/flexmdm/predictor_flexmdm.py
• Fallback in FlexMDMPredictor.to_dict() in the same file

To Reproduce
Run any FlexMDM seq2seq prediction path that logs time_taken (e.g. TinyGSM / GSM8K eval, or IWSLT prediction once wired), and compare wall time attributed to predict() vs models that only batch_decode the full sequence once.

# Example: any FlexMDM experiment with a prediction dataloader + LogPredictions
xlm job_type=eval experiment=tinygsm_flexmdm_gsm8k_eval +eval.ckpt_path=...

Inspect prediction JSONL / profiler: time_taken includes the Python per-row suffix decode after sampling finishes.

Expected behavior

• predict() returns generation tensors plus whatever full-sequence text decode is already required for history/debug.
• Suffix-only string materialization (generated_text) happens in to_dict() (logging) and/or post-hoc evaluators.
• time_taken reflects generation (and intentional full-sequence decode), not post-prediction string surgery.
• TinyGSM / GSM8K continue to get generated_text via to_dict → JSONL without needing it inside predict().

Actual behavior
predict() calls decode_generation_suffix(xt, fixed_gaps, ...) before returning, and time_taken is measured around that call. FlexMDM therefore pays an extra Python decode loop on every prediction batch inside the timed path.

Relevant config
Any FlexMDM prediction config that uses seq2seq / TinyGSM-style fixed masks, e.g.:

experiment=tinygsm_flexmdm_gsm8k_eval
# or other FlexMDM seq2seq prediction runs

Environment

• xlm-core / xlm-models: current main (FlexMDM predictor as above)
• Not environment-specific

Additional context
STAR-style seq2seq metrics already score token ids via slice metadata (output_start_idx / fixed / token_type_ids) and do not need this decode in predict(). Post-hoc text metrics (TinyGSM code-exec, future MT BLEU) should derive suffix strings outside the timed generation path.

Suggested fix:

1. Remove decode_generation_suffix from predict() return path (stop putting generated_text on the predict dict, or build it only in to_dict).
2. Rely on / harden the existing to_dict fallback so JSONL still gets generated_text.
3. Confirm TinyGSM Gsm8kCodeEval still sees generated_text from logged predictions.
4. Optionally add a short note in predictor docs: suffix decode is a logging concern, not part of timed predict().