Optimize torch.nn.functional.one_hot

[xinyu-intel]

🚀 The feature, motivation and pitch

one_hot is used in vLLM Gemma 4 26B workload https://github.com/vllm-project/vllm/blob/308cec5864890f5c0724e1d4531d9fe2ee0a8209/vllm/model_executor/models/gemma4.py#L214. We found 120 D2H copy were invoked per step and most likely caused by the boundary check in https://github.com/pytorch/pytorch/blob/beae96dfc1a2880f80a62f196306188a8d6dfdd9/aten/src/ATen/native/Onehot.cpp#L62. Pls consider to follow cuda to skip the check.

Alternatives

No response

Additional context

No response

[Silv3S]

Thanks @xinyu-intel for detailed description of potential root cause. I adjusted the boundary checks and prepared small benchmark. After patching, there are no more d2h calls

Before - Self CPU time total: 494.844ms

                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
            aten::one_hot         1.59%       7.868ms       100.00%     494.844ms     494.844us          1000  
               aten::item         0.37%       1.827ms        79.27%     392.264ms     196.132us          2000  
aten::_local_scalar_dense        78.90%     390.438ms        78.90%     390.438ms     195.219us          2000  
                aten::max         5.06%      25.023ms         5.48%      27.106ms      27.106us          1000  
           aten::scatter_         5.13%      25.370ms         5.21%      25.802ms      25.802us          1000  
                aten::min         4.38%      21.696ms         4.73%      23.424ms      23.424us          1000  
              aten::zeros         0.45%       2.247ms         3.13%      15.465ms      15.465us          1000  
              aten::zero_         0.33%       1.638ms         2.32%      11.497ms      11.497us          1000  
              aten::fill_         1.99%       9.860ms         1.99%       9.860ms       9.860us          1000  
              aten::empty         0.95%       4.706ms         0.95%       4.706ms       1.569us          3000  
          aten::unsqueeze         0.48%       2.379ms         0.59%       2.915ms       2.915us          1000  
         aten::as_strided         0.36%       1.795ms         0.36%       1.795ms       0.449us          4000  

After - Self CPU time total: 15.692ms

            Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
   aten::one_hot         5.71%     895.750us       100.00%      15.692ms      15.692us          1000  
  aten::scatter_        57.82%       9.073ms        59.01%       9.260ms       9.260us          1000  
     aten::zeros         4.64%     728.849us        30.98%       4.861ms       4.861us          1000  
     aten::zero_         2.70%     423.864us        23.17%       3.636ms       3.636us          1000  
     aten::fill_        20.47%       3.212ms        20.47%       3.212ms       3.212us          1000  
 aten::unsqueeze         3.41%     535.263us         4.30%     675.194us       0.675us          1000  
     aten::empty         3.16%     496.520us         3.16%     496.520us       0.497us          1000  
aten::as_strided         2.09%     327.425us         2.09%     327.425us       0.164us          2000 

Could you please check if this patch improves real model performance pytorch/pytorch#179831?

[xinyu-intel]

Thanks @xinyu-intel for detailed description of potential root cause. I adjusted the boundary checks and prepared small benchmark. After patching, there are no more d2h calls

Before - Self CPU time total: 494.844ms

                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
            aten::one_hot         1.59%       7.868ms       100.00%     494.844ms     494.844us          1000  
               aten::item         0.37%       1.827ms        79.27%     392.264ms     196.132us          2000  
aten::_local_scalar_dense        78.90%     390.438ms        78.90%     390.438ms     195.219us          2000  
                aten::max         5.06%      25.023ms         5.48%      27.106ms      27.106us          1000  
           aten::scatter_         5.13%      25.370ms         5.21%      25.802ms      25.802us          1000  
                aten::min         4.38%      21.696ms         4.73%      23.424ms      23.424us          1000  
              aten::zeros         0.45%       2.247ms         3.13%      15.465ms      15.465us          1000  
              aten::zero_         0.33%       1.638ms         2.32%      11.497ms      11.497us          1000  
              aten::fill_         1.99%       9.860ms         1.99%       9.860ms       9.860us          1000  
              aten::empty         0.95%       4.706ms         0.95%       4.706ms       1.569us          3000  
          aten::unsqueeze         0.48%       2.379ms         0.59%       2.915ms       2.915us          1000  
         aten::as_strided         0.36%       1.795ms         0.36%       1.795ms       0.449us          4000  

After - Self CPU time total: 15.692ms

            Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
   aten::one_hot         5.71%     895.750us       100.00%      15.692ms      15.692us          1000  
  aten::scatter_        57.82%       9.073ms        59.01%       9.260ms       9.260us          1000  
     aten::zeros         4.64%     728.849us        30.98%       4.861ms       4.861us          1000  
     aten::zero_         2.70%     423.864us        23.17%       3.636ms       3.636us          1000  
     aten::fill_        20.47%       3.212ms        20.47%       3.212ms       3.212us          1000  
 aten::unsqueeze         3.41%     535.263us         4.30%     675.194us       0.675us          1000  
     aten::empty         3.16%     496.520us         3.16%     496.520us       0.497us          1000  
aten::as_strided         2.09%     327.425us         2.09%     327.425us       0.164us          2000 

Could you please check if this patch improves real model performance pytorch/pytorch#179831?

thanks for the quick fix, it definitely should work. we're using zeros+scatter_ for a temp WA, I think the perf should be similar.

[yintong-lu]

Thanks @Silv3S for the quick fix. I have verified your fix on Gemma4-26B-A4B model and compared it with a scatter_add_ workaround.
I profiled under torch.compile with concurrency=1,4,16. Both approaches deliver essentially the same speedup. At c=16, TTFT dropped from 6.0s to ~1.6s; throughput rose from 1285 to ~4700 tok/s.

[chuanqi129]

@Silv3S @xinyu-intel @yintong-lu Do we need this fix for PT 2.12? or can we target it for PT 2.13?

[xinyu-intel]

2.12 should be better

[Silv3S]

I think the sooner the better, one_hot is very common op so reducing D2H calls could also improve performance in other models

[chuanqi129]

Please @Silv3S drive the PR review and merge ASAP