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

[GenTang]

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 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:

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
• norm_fn.py

Minimal reproduction

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 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