deepseek-ai / deepseek-ai/TileKernels

Engram TileLang gate kernel does not match original Engram Pytorch Demo

Open
#19 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.