do not pad to nearest pow2 if using simpleGo

Author: RJKeevil

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [0.5.3] - 2025-09-01
+
+### Changed
+
+- Performance improvement: do not pad inputs to POW2 when using simplego backend
+
 ## [0.5.2] - 2025-08-30
 
 ### Changed

pipelineBackends/model_gomlx.go

@@ -193,7 +193,7 @@ func loadInputOutputMetaGoMLX(model *onnx.Model) ([]InputOutputInfo, []InputOutp
 	return inputs, outputs
 }
 
-func createInputTensorsGoMLX(batch *PipelineBatch, model *Model, padBatchDimension bool) error {
+func createInputTensorsGoMLX(batch *PipelineBatch, model *Model, padBatchDimension bool, padSequenceDimension bool) error {
 	leftPad := len(model.EosTokenIDs) > 0
 
 	// TODO: replace this once dynamic input shapes fixed
@@ -204,7 +204,7 @@ func createInputTensorsGoMLX(batch *PipelineBatch, model *Model, padBatchDimensi
 		batchSize = nextPowerOf2(batchSize)
 	}
 	maxSeqLength := batch.MaxSequenceLength
-	if !leftPad {
+	if padSequenceDimension && !leftPad {
 		maxSeqLength = nextPowerOf2(maxSeqLength)
 	}
 	total := batchSize * maxSeqLength

pipelineBackends/pipeline.go

@@ -140,8 +140,10 @@ func RunGenerativeSessionOnBatch(batch *PipelineBatch, p *BasePipeline) error {
 	switch p.Runtime {
 	case "ORT":
 		return runGenerativeORTSessionOnBatch(batch, p)
-	case "GO", "XLA":
-		return errors.New("GO/XLA backend is not yet implemented for generative models")
+	case "GO":
+		return errors.New("GO backend is not yet implemented for generative models")
+	case "XLA":
+		return errors.New("XLA backend is not yet implemented for generative models")
 	default:
 		return errors.New("invalid backend")
 	}
@@ -153,8 +155,10 @@ func CreateInputTensorsTraining(batch *PipelineBatch, model *Model, runtime stri
 	switch runtime {
 	case "ORT":
 		return createInputTensorsORT(batch, model)
-	case "GO", "XLA":
-		return createInputTensorsGoMLX(batch, model, false)
+	case "GO":
+		return createInputTensorsGoMLX(batch, model, false, false)
+	case "XLA":
+		return createInputTensorsGoMLX(batch, model, false, true)
 	}
 	return nil
 }
@@ -163,8 +167,10 @@ func CreateInputTensors(batch *PipelineBatch, model *Model, runtime string) erro
 	switch runtime {
 	case "ORT":
 		return createInputTensorsORT(batch, model)
-	case "GO", "XLA":
-		return createInputTensorsGoMLX(batch, model, true)
+	case "GO":
+		return createInputTensorsGoMLX(batch, model, false, false)
+	case "XLA":
+		return createInputTensorsGoMLX(batch, model, true, true)
 	}
 	return nil
 }