I ungdomslandslagen var Stenmalm framgångsrik. Vid U19-VM 2011 var han med när Sveriges U19-landslag tog brons. Vid U20-EM 2012 var han med då U21-landslaget kom på fjärde plats och vid U21-VM 2013 var han högst delaktig då Sverige tog guld. Stenmalm utsågs till VM-turneringens mest värdefulla spelare (MVP). Som senior har han representerat Sveriges landslag bland annat vid OS 2016 i Rio de Janeiro, då Sverige slutade på elfte plats.