Repository: pytorch
Issue: Embedding forward performance analysis
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

Embedding operation maps indices into a dense vector space, typically used to look up embeddings for discrete input tokens in tasks like natural language processing.

Inductor compiled kernel is slower than liger version for this op. For input 11, liger is 1.26X faster than inductor.

## **Reproduce tests**

```
% python run.py --op embedding  --mode fwd --num-inputs 1 --input-id 11 --precision fp32 --metrics latency,speedup --cudagraph
  0%|                                                                                                                                                         | 0/1 [00:00<?, ?it/s]INFO:tritonbench.utils.triton_op:Took 178.71ms to get benchmark function for torch_embedding
INFO:tritonbench.utils.triton_op:Took 178.65ms to get benchmark function for liger_embedding
INFO:tritonbench.utils.triton_op:Took 822.35ms to get benchmark function for inductor_embedding
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.54s/it]
         (B, T, D, V)    torch_embedding-latency    liger_embedding-speedup    liger_embedding-latency    inductor_embedding-speedup    inductor_embedding-latency
---------------------  -------------------------  -------------------------  -------------------------  ----------------------------  ----------------------------
(8, 2048, 4096, 8192)                   0.472864                    2.50288                   0.188928                        1.9795                       0.23888
```

## Report Analysis

For input 11, the input dimensions are \``(B, T, D, V) = (8, 2048, 4096, 8192)`\`, so 

```
embeddings.shape = (V, D) = (8192, 4096)
indices.shape = (B, T) = (8, 2048)
output.shape = (B, T, D) = (8, 2048, 4096)
```

For FP32, the theoretical optimal write bytes should be 

sizeof(output) \= 8\*2048\*4096\*4/1024/1024 \= 256MB

The theoretical optimal read bytes should be 

sizeof(indices \+ embeddings) \= (8\*2048  \+  \*8192\*4096) \* 4/1024/1024 MB= 128.0625MB 

### Inductor compiled triton kernel

The following is the compiled triton kernel and its wrapper code.

In the `grid=grid(67108864)` function, it will compute the final grid size.   
[https://github.com/pytorch/pytorch/blob/a84779040049377aec7b62e37becb7327950541e/torch/\_inductor/runtime/triton\_heuristics.py\#L1962-L2004](https://github.com/pytorch/pytorch/blob/a84779040049377aec7b62e37becb7327950541e/torch/_inductor/runtime/triton_heuristics.py#L1962-L2004)

Kernel triton\_poi\_fused\_embedding\_0’s final grid size \= (65536, 1,1) while its XBLOCK= 1024\.  
So, theoretical memory write request per warp is

65536\*1024/32/4 \=524,288 

32 threads per warp, 4 FP32 per STG.E.128

It matches the memory write request in ncu reports. 
<img width="1009" alt="image" src="https://github.com/user-attachments/assets/c3c589c5-9c82-4d06-a939-541492916604">

However, for memory read, the indices load is compiled to the following triton line.
```
x1 = (xindex // 4096)
tmp0 = tl.load(in_ptr0 + (x1), None, eviction_policy='evict_last') 
```

Because each thread will execute this load, the total read requests for indices are 524,288 too. So the total number of load requests is 1048576. This matches the memory table in ncu report. 
<img width="919" alt="image" src="https://github.com/user-attachments/assets/32af253e-23c0-44b4-8aa1-035edb681cfd">

I also confirmed with a modified [mem\_trace tool](https://github.com/FindHao/nvbit_test2) in nvbit. 

`MEMTRACE: Instruction execution counts for kernel triton_poi_fused_embedding_0 in context 0x55ed1dd95330:`  
 `Instruction: STG.E.128            Count: 524288`  
 `Instruction: LDG.E.128            Count: 524288`  
 `Instruction: LDG.E.EL.64          Count: 524288`

LDG.E.128 for embedding loads, LDG.E.EL.64 for indices loads in this case.

### Liger implementation

[https://github.com/linkedin/Liger-Kernel/blob/main/src/liger\_kernel/ops/experimental/embedding.py\#L81-L95](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/experimental/embedding.py#L81-L95)  
For liger version, it converts shapes to make it compatible with two-dimension blocks (128, 128), and use gridsize \= (128,32), which is much smaller than inductor compiled version. Triton also automatically place the smaller indices to shared memory. Thus, the liger version saves a lot of global memory loads. 

<img width="1001" alt="image" src="https://github.com/user-attachments/assets/a01d1de0-c9f1-4d1b-bd1d-297a839664a3">


### Potential solution:

fix its tiling


### Error logs

_No response_

### Versions

main branch

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov