[mlperf][pkgci] Update punet-fp8 with reduction dim as last dim (#19316)

[mlperf][pkgci] Update punet-fp8 with reduction dim as last dim (#19316) #4100