Gemma4 sampling optimizations. #5621

Job	Run time
Build tunix package / Build Tunix wheel	14s
tunix_cpu_unit_tests / run	3m 0s
tunix_tpu_unit_tests / run_prod	11m 24s
tunix_tpu_unit_tests / run_dev	28m 51s
Notify failed build	4s
	43m 33s

Provide feedback