deepseek-ai / deepseek-ai/TileKernels
Engram TileLang gate kernel does not match original Engram Pytorch Demo
- Dominant language
- Python
- Stars
- 1.8k
- Forks
- 162
- PR merge metrics
- No merged PRs in 30d
Description
In PyTorch Engram, the module returns only the Engram contribution:
```python
value = gates * self.value_proj(embeddings).unsqueeze(2)
output = value + self.short_conv(value)
return output
```
Reference:
https://github.com/deepseek-ai/Engram/blob/fb7f84a21f91223715394a33a1dc24bbfb7f788e/engram_demo_v1.py#L376-L378
The residual add with `hidden_states` happens outside the Engram module:
```python
hidden_states = self.engram(hidden_states=hidden_states,input_ids=input_ids) + hidden_states
```
Reference:
https://github.com/deepseek-ai/Engram/blob/fb7f84a21f91223715394a33a1dc24bbfb7f788e/engram_demo_v1.py#L389-L391
But the TileLang gate kernel adds `hidden_states` inside the kernel,even before short conv:
```python
output[i_s, pid_h, sub_base + thread_idx * vec_size + i_k] = x_local[i_k] + gate_score_reducer[0] * v_local[i_k]
```
Reference:
https://github.com/deepseek-ai/TileKernels/blob/36d9e45d38e204ebb87e6f6e833821eee0482fe5/tile_kernels/engram/engram_gate_kernel.py#L176-L177
So the two implementations differ:
```python
# PyTorch Engram
engram_output = gate * value + short_conv(gate * value)
hidden_states = hidden_states + engram_output
```
```python
# TileLang gate kernel
gate_output = hidden_states + gate * value
```
This makes the TileLang kernel inconsistent with the PyTorch Engram op. The gate kernel should output only:
```python
output = gate * value
```
if it is intended to match the PyTorch Engram module semantics.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.