Add ShardedRMSNorm for Q-K normalization under tensor parallelism#47

aws-satyajith merged 2 commits intoaws-neuron:mainaws-neuron/neuronx-distributed-inference:mainfrom

sdeeptan-aws:olmo2sdeeptan-aws/neuronx-distributed-inference:olmo2Copy head branch name to clipboard