deepseek-ai / deepseek-ai/TileKernels

[Correctness][mHC] mhc_pre ignores norm_weight in the no-grad fused path

Open
#22 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

## Summary

`mhc_pre()` produces different computations depending on whether gradients are enabled when `norm_weight` is not `None`.

In the grad-enabled path, `mhc_pre_norm_fn()` merges the RMSNorm weight into `fn` before the TF32 GEMM. In the no-grad path, `mhc_pre()` dispatches directly to `mhc_pre_big_fuse()` without passing or merging `norm_weight`.

As a result, `norm_weight` is silently ignored during inference/evaluation.

Severity: **High for affected**, because this is a silent numerical correctness issue rather than an explicit failure.

Affected (and):

- `norm_weight is not None`
- `torch.no_grad()` or `torch.inference_mode()`
- calls through the high-level `mhc_pre()` API

Not affected (or):

- `norm_weight is None`
- the grad-enabled path

## Expected behavior according to the paper

Section 4.3.1 of the [mHC paper](https://arxiv.org/pdf/2512.24880#page=10) describes the RMSNorm fusion and states that:

> the RMSNorm weight is also absorbed in φ_l

This is mathematically equivalent to merging the RMSNorm weight into the columns of the projection matrix before the GEMM:

```text
RMSNorm(x, w) @ fn.T
= (x / rms(x) * w) @ fn.T
= (x / rms(x)) @ (fn * w).T
```

The grad-enabled implementation follows this formulation, but the no-grad fused path does not.

## Root cause

The grad-enabled path
1. it passes `norm_weight` to `mhc_pre_norm_fn()`:
2. `mhc_pre_norm_fn()` then performs the fusion.

However, the no-grad path passes the original `fn` directly to `mhc_pre_big_fuse()`. `mhc_pre_big_fuse()` consumes `fn`, but it has no access to `norm_weight`.

Relevant code:

- [`functional.py`](https://github.com/deepseek-ai/TileKernels/blob/main/tile_kernels/modeling/mhc/functional.py)
- [`norm_fn.py`](https://github.com/deepseek-ai/TileKernels/blob/main/tile_kernels/modeling/mhc/ops/norm_fn.py)

## Minimal reproduction

```python
import torch
from tile_kernels.modeling.mhc.functional import mhc_pre

BATCH_SIZE = 1
SEQ_LEN = 1024
MHC_MULT = 4
HIDDEN_SIZE = 256
MIX_SIZE = MHC_MULT * (MHC_MULT + 2)

def make_inputs(device) -> dict[str, torch.Tensor]:

residual = torch.randn(
BATCH_SIZE, SEQ_LEN, MHC_MULT, HIDDEN_SIZE,
device=device, dtype=torch.float32).to(torch.bfloat16)
fn = torch.randn(
MIX_SIZE, MHC_MULT * HIDDEN_SIZE,
device=device, dtype=torch.float32)
scale = torch.randn(3, device=device, dtype=torch.float32)
base = torch.randn(MIX_SIZE, device=device, dtype=torch.float32)
norm_weight = torch.rand(MHC_MULT * HIDDEN_SIZE, device=device)

return {
"residual": residual,
"fn": fn,
"scale": scale,
"base": base,
"norm_weight": norm_weight
}

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
norm_weight_inputs = make_inputs(device)
no_norm_weight_inputs = norm_weight_inputs.copy()
no_norm_weight_inputs.pop("norm_weight")

def test_train_eval_path_output(inputs):
train_path_out, _ = mhc_pre(**inputs)
with torch.no_grad():
eval_path_out, _ = mhc_pre(**inputs)
diff = train_path_out - eval_path_out
rel_err = diff.norm() / train_path_out.norm()
print(f"device : {device}")
print(f"max abs diff : {diff.abs().max().item():.6f}")
print(f"mean abs diff : {diff.abs().mean().item():.6f}")
print(f"train-path output norm : {train_path_out.norm().item():.6f}")
print(f"relative error : {rel_err.item():.4%}")

print("***test with norm weight***")
test_train_eval_path_output(norm_weight_inputs)
print("***test without norm weight***")
test_train_eval_path_output(no_norm_weight_inputs)
```

The result will be as blow:
```
***test with norm weight***
device : cuda
max abs diff : 4.843750
mean abs diff : 0.408203
train-path output norm : 680.000000
relative error : 47.2656%
***test without norm weight***
device : cuda
max abs diff : 0.000000
mean abs diff : 0.000000
train-path output norm : 696.000000
relative error : 0.0000%
```

## Why existing tests did not catch this

The lower-level [`test_norm_fn`](https://github.com/deepseek-ai/TileKernels/blob/main/tests/mhc/test_norm_fn.py) tests cover both `norm_weight=None` and `norm_weight!=None`, so the merge kernel itself is tested.

However:

1. There is no upstream end-to-end test for the high-level `mhc_pre()` API.
2. Therefore, no test compares the grad-enabled and no-grad dispatch paths while passing `None` or `norm_weight`.

This allowed the low-level implementations to appear correct while the high-level goes wrong in inference mode.

## Proposed fix

Before dispatching to `mhc_pre_big_fuse()`, merge `norm_weight` into `fn` using the existing `_mhc_fn_normw_merge_fwd` kernel

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.