Repository: pytorch
Issue: TorchDispatchMode cann't capture the operator which name is aten::index_put_ impl_
Labels: triaged,module: __torch_dispatch__

### 🐛 Describe the bug

🐛 Describe the bug
In my understanding, **TorchDispatchMode** should capture atens which call actual kernel. When I run the code below, I found it missed **aten::index_put_ impl_**.

And I try to print torch ops dispatch stack, and try using profiler. They both can see **aten::index_put_ impl_**.

What's the reason for this? 

# 0x01 Using TorchDispatchMode

```python
import torch
from tests.unit_tests.test_utilities import Utils
import numpy as np
import torch.distributed as dist

from torch.utils._python_dispatch import TorchDispatchMode
from megatron.core.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy

class PrintingMode(TorchDispatchMode):
      def __torch_dispatch__(self, func, types, args=(), kwargs=None):
          print(f"{func.__module__}.{func.__name__}")
          return func(*args, **kwargs)

      def __enter__(self):
          # 进入with块时执行的代码
          print("Entering PrintingMode")
          return super().__enter__()

      def __exit__(self, exc_type, exc_value, traceback):
          # 退出with块时执行的代码
          return super().__exit__(exc_type, exc_value, traceback)

def test_vocab_parallel_cross_entropy():
    Utils.initialize_model_parallel(1,1)
    # vocab_parallel_logits = torch.range(0,7).repeat(16,4).cuda()
    # target = torch.arange(0,32,2).cuda()
    vocab_parallel_logits = torch.empty((4096, 1, 32000), dtype=torch.float16, device='cuda:0')
    # 设置步长
    vocab_parallel_logits = vocab_parallel_logits.as_strided(
        (4096, 1, 32000), (32000, 32000, 1))

    # 创建 target 张量
    target = torch.empty((4096, 1), dtype=torch.int64, device='cuda:0')
    # 设置步长
    target = target.as_strided((4096, 1), (1, 4096))

    print(vocab_parallel_logits.shape)
    print(target.shape)


    output = vocab_parallel_cross_entropy(vocab_parallel_logits, target)

    Utils.destroy_model_parallel()

 # 初始化分布式环境
#dist.init_process_group(backend='nccl', init_method='env://', world_size=1, rank=0, device_ids=[0])

with PrintingMode():
    test_vocab_parallel_cross_entropy()

# 销毁进程组
dist.destroy_process_group()


```

Outputs is bellow: 

> No aten::index_put_ impl_

```bash
> python test_cross_entropy.py
Entering PrintingMode
Initializing torch.distributed with rank: 0, world_size: 1
torch._ops.aten.empty.memory_format
torch._ops.c10d.barrier.default
[rank0]:[W117 10:31:55.661054174 ProcessGroupNCCL.cpp:4457] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in barrier() to force use of a particular device, or call init_process_group() with a device_id.
torch._ops.aten.empty.memory_format
torch._ops.aten.as_strided.default
torch._ops.aten.empty.memory_format
torch._ops.aten.as_strided.default
torch.Size([4096, 1, 32000])
torch.Size([4096, 1])
torch._ops.aten._to_copy.default
torch._ops.aten.max.dim
torch._ops.c10d.allreduce_.default
torch._ops.aten.unsqueeze.default
torch._ops.aten.sub_.Tensor
torch._ops.aten.lt.Scalar
torch._ops.aten.ge.Scalar
torch._ops.aten.bitwise_or.Tensor
torch._ops.aten.clone.default
torch._ops.aten.sub.Tensor
torch._ops.aten.lift_fresh.default
torch._ops.aten.index_put_.default
torch._ops.aten.view.default
torch._ops.aten.view.default
torch._ops.aten.arange.start
torch._ops.aten.index.Tensor
torch._ops.aten.clone.default
torch._ops.aten.view.default
torch._ops.aten.lift_fresh.default
torch._ops.aten.index_put_.default
torch._ops.aten.exp.out
torch._ops.aten.sum.dim_IntList
torch._ops.c10d.allreduce_.default
torch._ops.c10d.allreduce_.default
torch._ops.aten.log.default
torch._ops.aten.sub.Tensor
torch._ops.aten.unsqueeze.default
torch._ops.aten.div_.Tensor
torch._ops.aten.empty.memory_format
torch._ops.c10d.barrier.default

```

# 0x02 export TORCH_SHOW_DISPATCH_TRACE=1

![Image](https://github.com/user-attachments/assets/c8b4f583-d655-46ff-8857-f2125d881890)

We can see the line 476, it called  **aten::index_put_ impl_**, and seems like using actual kernel.

> Has aten::index_put_ impl_

![Image](https://github.com/user-attachments/assets/093adec1-ee7c-4132-80e4-1a49cd43f3f7)

# 0x03 Profiler

While I using profiler, I can also see the  **aten::index_put_ impl_**

![Image](https://github.com/user-attachments/assets/d8b280a1-c3b4-4298-91fc-269cc1017a29)

> Also has aten::index_put_ impl_

![Image](https://github.com/user-attachments/assets/73c5ea3c-9f8e-4244-b673-2810e18d5b7f)



### Versions

# Version

```bash
> python collect_env.py                                                                                                                                                                                                                                 
Collecting environment information...                                                                                                                                                                                                                   
PyTorch version: 2.6.0a0+git78543e6                                                                                                                                                                                                                     
Is debug build: True                                                                                                                                                                                                                                    
CUDA used to build PyTorch: 12.2                                                                                                                                                                                                                        
ROCM used to build PyTorch: N/A                                                                                                                                                                                                                         
                                                                                                                                                                                                                                                        
OS: Ubuntu 22.04.5 LTS (x86_64)                                                                                                                                                                                                                         
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0                                                                                                                                                                                                      
Clang version: Could not collect                                                                                                                                                                                                                        
CMake version: version 3.31.1                                                                                                                                                                                                                           
Libc version: glibc-2.35                                                                                                                                                                                                                                
                                                                                                                                                                                                                                                        
Python version: 3.9.20 (main, Oct  3 2024, 07:27:41)  [GCC 11.2.0] (64-bit runtime)                                                                                                                                                                     
Python platform: Linux-6.8.0-51-generic-x86_64-with-glibc2.35                                                                                                                                                                                           
Is CUDA available: True                                                                                                                                                                                                                                 
CUDA runtime version: Could not collect                                                                                                                                                                                                                 
CUDA_MODULE_LOADING set to: LAZY                                                                                                                                                                                                                        
GPU models and configuration: GPU 0: NVIDIA T1000                                                                                                                                                                                                       
Nvidia driver version: 535.183.01                                                                                                                                                                                                                       
cuDNN version: Could not collect                                                                                                                                                                                                                        
HIP runtime version: N/A                                                                                                                                                                                                                                
MIOpen runtime version: N/A                                                                                                                                                                                                                             
Is XNNPACK available: True                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                        
CPU:                                                                                                                                                                                                                                                    
架构：                                x86_64                                                                                                                                                                                                            
CPU 运行模式：                        32-bit, 64-bit                                                                                                                                                                                                    
Address sizes:                        39 bits physical, 48 bits virtual                                                                                                                                                                                 
字节序：                              Little Endian
CPU:                                  16
在线 CPU 列表：                       0-15
厂商 ID：                             GenuineIntel
型号名称：                            11th Gen Intel(R) Core(TM) i7-11700 @ 2.50GHz
CPU 系列：                            6
型号：                                167
每个核的线程数：                      2                                                                                                                                                                                                                 
每个座的核数：                        8                                                                                                                                                                                                                 
座：                                  1                                                                                                                                                                                                                 
步进：                                1                                                                                                                                                                                                                 
CPU 最大 MHz：                        4900.0000                                                                                                                                                                                                         
CPU 最小 MHz：                        800.0000                                                                                                                                                                                                          
BogoMIPS：                            4992.00                                                                                                                                                                                                           
标记：                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap avx512ifma clflushopt intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm md_clear flush_l1d arch_capabilities
虚拟化：                              VT-x                                                                                                                                                                                                              
L1d 缓存：                            384 KiB (8 instances)                                                                                                                                                                                             
L1i 缓存：                            256 KiB (8 instances)                                                                                                                                                                                             
L2 缓存：                             4 MiB (8 instances)                                                                                                                                                                                               
L3 缓存：                             16 MiB (1 instance)                                                                                                                                                                                               
NUMA 节点：                           1                                                                                                                                                                                                                 
NUMA 节点0 CPU：                      0-15                                                                                                                                                                                                              
Vulnerability Gather data sampling:   Mitigation; Microcode                                                                                                                                                                                             
Vulnerability Itlb multihit:          Not affected                                                                                                                                                                                                      
Vulnerability L1tf:                   Not affected                                                                                                                                                                                                      
Vulnerability Mds:                    Not affected                                                                                                                                                                                                      
Vulnerability Meltdown:               Not affected                                                                                                                                                                                                      
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT vulnerable                                                                                                                                                                     
Vulnerability Reg file data sampling: Not affected                                                                                                                                                                                                      
Vulnerability Retbleed:               Mitigation; Enhanced IBRS                                                                                                                                                                                         
Vulnerability Spec rstack overflow:   Not affected                                                                                                                                                                                                      
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl                                                                                                                                                           
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization                                                                                                                                              
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop                                                                                           
Vulnerability Srbds:                  Not affected                                                                                                                                                                                                      
Vulnerability Tsx async abort:        Not affected                                                                                                                                                                                                      
                                                                                                                                                                                                                                                        
Versions of relevant libraries:                                                                                                                                                                                                                         
[pip3] numpy==2.0.2                                                                                                                                                                                                                                     
[pip3] optree==0.13.1                                                                                                                                                                                                                                   
[pip3] torch==2.6.0a0+git78543e6                                                                                                                                                                                                                        
[conda] magma-cuda121             2.6.1                         1    pytorch                                                                                                                                                                            
[conda] mkl-include               2025.0.0           hf2ce2f3_941    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge                                                                                                                    
[conda] mkl-static                2025.0.0           ha770c72_941    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge                                                                                                                    
[conda] numpy                     2.0.2                    pypi_0    pypi                                                                                                                                                                               
[conda] optree                    0.13.1                   pypi_0    pypi                                                                                                                                                                               
[conda] torch                     2.6.0a0+git78543e6           dev_0    <develop>                
```

cc @Chillee @ezyang @zou3519 @albanD @samdow