Discuss the backward impl of mHC

I have long been wondering if the backward of sinkhorn step can be computed like this:

triton code

```py
"""
Corrected Sinkhorn Backward Pass (fix.md Scheme 2)

When R satisfies R @ 1 = r and R^T @ 1 = c (r, c != 1),
the corrected linear system for beta is:
(diag(c) - R^T diag(r)^{-1} R) beta = s_c - R^T diag(r)^{-1} s_r

where s_r = (R ⊙ dR) @ 1, s_c = (R ⊙ dR)^T @ 1.

This file implements the corrected kernel and compares with:
1. PyTorch autograd
2. The original triton_reduce.py (assumes r=1, c=1)
"""

from icecream import ic
import torch
import triton
import triton.language as tl
from tqdm import trange

# TMA descriptors require a global memory allocation
def alloc_fn(size: int, alignment: int, stream: int | None):
return torch.empty(size, device="cuda", dtype=torch.int8)

triton.set_allocator(alloc_fn)

dtype = torch.float64
EPS = tl.constexpr(1e-10)
eps = 1e-10

# ─────────────────────────────────────────────────────────────────────────────
# Forward: standard Sinkhorn (same as triton_reduce.py)
# ─────────────────────────────────────────────────────────────────────────────
def sinkhorn_forward(M, iters=20):
P = torch.exp(M)
R = P
for _ in range(iters):
R = R / R.sum(-2, keepdim=True).clamp(min=eps)
R = R / R.sum(-1, keepdim=True).clamp(min=eps)
return R, P

# ─────────────────────────────────────────────────────────────────────────────
# Original matvec_S: S = I - R^T R (assumes r=1, c=1)
# ─────────────────────────────────────────────────────────────────────────────
@triton.jit
def matvec_S(R, x):
"""
S = I - R^T R, perform S @ x WITHOUT materializing RTR.
Computes: x - R^T (R x) using two matvecs.
R: (tilesize, n, n)
x: (tilesize, n, 1)
returns: (tilesize, n, 1)
"""
Rx = tl.sum(R * x.permute(0, 2, 1), axis=-1).expand_dims(-1)
RT = R.permute(0, 2, 1)
RTRx = tl.sum(RT * Rx.permute(0, 2, 1), axis=-1).expand_dims(-1)
return x - RTRx

# ─────────────────────────────────────────────────────────────────────────────
# Corrected matvec_S: S = diag(c) - R^T diag(r)^{-1} R
# ─────────────────────────────────────────────────────────────────────────────
@triton.jit
def matvec_S_corrected(R, x, r, c):
"""
Corrected S = diag(c) - R^T diag(r)^{-1} R.
Computes: diag(c) x - R^T (diag(r)^{-1} (R x))

R: (tilesize, n, n)
x: (tilesize, n, 1)
r: (tilesize, n, 1) row sums of R
c: (tilesize, n, 1) col sums of R
returns: (tilesize, n, 1)
"""
# R @ x -> (tilesize, n, 1)
Rx = tl.sum(R * x.permute(0, 2, 1), axis=-1).expand_dims(-1)
# diag(r)^{-1} @ Rx -> element-wise divide
Rx_scaled = Rx / r
# R^T @ Rx_scaled -> (tilesize, n, 1)
RT = R.permute(0, 2, 1)
RTRx = tl.sum(RT * Rx_scaled.permute(0, 2, 1), axis=-1).expand_dims(-1)
# diag(c) @ x - R^T diag(r)^{-1} R @ x
return c * x - RTRx

# ─────────────────────────────────────────────────────────────────────────────
# Corrected backward kernel
# ─────────────────────────────────────────────────────────────────────────────
@triton.autotune(
configs=[
triton.Config({"tilesize": tilesize}, num_stages=1, num_warps=num_warps)
for tilesize in [1, 2, 4, 8, 16, 32, 64]
for num_warps in [1, 2, 4, 8]
],
key=[],
)
@triton.jit
def sinkhorn_bwd_corrected_kernel(
seqlen,
out,
dout,
res,
out_stride_0,
out_stride_1,
out_stride_2,
dout_stride_0,
dout_stride_1,
dout_stride_2,
res_stride_0,
res_stride_1,
res_stride_2,
n_stream: tl.constexpr,
tilesize: tl.constexpr,
):
out_desc = tl.make_tensor_descriptor(
out,
shape=[seqlen, n_stream, n_stream],
strides=[out_stride_0, out_stride_1, out_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)
dout_desc = tl.make_tensor_descriptor(
dout,
shape=[seqlen, n_stream, n_stream],
strides=[dout_stride_0, dout_stride_1, dout_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)
res_desc = tl.make_tensor_descriptor(
res,
shape=[seqlen, n_stream, n_stream],
strides=[res_stride_0, res_stride_1, res_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)

seq_off = tl.program_id(0) * tilesize

R = out_desc.load([seq_off, 0, 0]) # (tilesize, n, n)
RT = R.permute(0, 2, 1)
dR = dout_desc.load([seq_off, 0, 0])

# ── Compute row sums r and col sums c of R ──────────────────────────────
# r: (tilesize, n, 1), c: (tilesize, n, 1)
r = tl.sum(R, axis=-1).expand_dims(-1) # row sums
c = tl.sum(R, axis=-2).expand_dims(-1) # col sums

# ── Step 1: s_r = (R ⊙ dR) @ 1, s_c = (R ⊙ dR)^T @ 1 ────────────────
RdR = R * dR
s_r = tl.sum(RdR, axis=-1).expand_dims(-1) # (tilesize, n, 1)
s_c = tl.sum(RdR, axis=-2).expand_dims(-1) # (tilesize, n, 1)

# ── Step 2: b = s_c - R^T diag(r)^{-1} s_r ─────────────────────────────
s_r_scaled = s_r / (r + EPS) # diag(r)^{-1} s_r
RT_sr = tl.sum(RT * s_r_scaled.permute(0, 2, 1), axis=-1).expand_dims(-1)
b = s_c - RT_sr # (tilesize, n, 1)

# ── Step 3: CG to solve (diag(c) - R^T diag(r)^{-1} R) beta = b ────────
x = tl.zeros((tilesize, n_stream, 1), dtype=R.dtype)
r_cg = b # residual (using r_cg to avoid name clash with row sums r)
p = r_cg
r_normsq = tl.sum(r_cg * r_cg, axis=1, keep_dims=True)

for _ in range(n_stream):
Sp = matvec_S_corrected(R, p, r, c)
pSp = tl.sum(p * Sp, axis=1, keep_dims=True)
alpha = r_normsq / (pSp + EPS)

x += alpha * p
r_cg -= alpha * Sp

r_new_normsq = tl.sum(r_cg * r_cg, axis=1, keep_dims=True)
beta_cg = r_new_normsq / (r_normsq + EPS)

p = r_cg + beta_cg * p
r_normsq = r_new_normsq

# ── Step 4: u = diag(r)^{-1} (s_r - R beta), v = beta ─────────────────
# R @ x: (tilesize, n, n) x (tilesize, n, 1) -> (tilesize, n, 1)
Rx = tl.sum(R * x.permute(0, 2, 1), axis=-1).expand_dims(-1)
u = (s_r - Rx) / r # (tilesize, n, 1)
v = x # (tilesize, n, 1)

# ── Step 5: M_ij = u_i + v_j ────────────────────────────────────────────
vt = v.reshape(tilesize, 1, n_stream)
M_mat = u + vt # broadcast -> (tilesize, n, n)

# ── Step 6: grad = (dR - M) ⊙ R ─────────────────────────────────────────
res_tile = (dR - M_mat) * R

res_desc.store([seq_off, 0, 0], res_tile)

def sinkhorn_bwd_corrected(
out: torch.Tensor,
dout: torch.Tensor,
repeat: int = 1,
warmup: bool = True,
):
seqlen = out.size(0)
n_stream = out.size(1)

res = torch.empty_like(out)

def grid(META):
return (triton.cdiv(seqlen, META["tilesize"]), 1, 1)

def _run():
sinkhorn_bwd_corrected_kernel[grid](
seqlen,
out,
dout,
res,
out.stride(0),
out.stride(1),
out.stride(2),
dout.stride(0),
dout.stride(1),
dout.stride(2),
res.stride(0),
res.stride(1),
res.stride(2),
n_stream,
)

if warmup:
for _ in trange(4, desc="warmup"):
_run()
torch.cuda.synchronize()

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start_event.record()
for _ in range(repeat):
_run()
end_event.record()
torch.cuda.synchronize()

elapsed_ms = start_event.elapsed_time(end_event)
if repeat > 1:
print(f"[corrected] Kernel execution time ({repeat=}): {elapsed_ms:.3f} ms")
print(f"[corrected] Average time per iteration: {elapsed_ms / repeat:.3f} ms")

return res

# ─────────────────────────────────────────────────────────────────────────────
# Original backward kernel (from triton_reduce.py, verbatim)
# ─────────────────────────────────────────────────────────────────────────────
@triton.autotune(
configs=[
triton.Config({"tilesize": tilesize}, num_stages=1, num_warps=num_warps)
for tilesize in [1, 2, 4, 8, 16, 32, 64]
for num_warps in [1, 2, 4, 8]
],
key=[],
)
@triton.jit
def sinkhorn_bwd_original_kernel(
seqlen,
out,
dout,
res,
out_stride_0,
out_stride_1,
out_stride_2,
dout_stride_0,
dout_stride_1,
dout_stride_2,
res_stride_0,
res_stride_1,
res_stride_2,
n_stream: tl.constexpr,
tilesize: tl.constexpr,
):
out_desc = tl.make_tensor_descriptor(
out,
shape=[seqlen, n_stream, n_stream],
strides=[out_stride_0, out_stride_1, out_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)
dout_desc = tl.make_tensor_descriptor(
dout,
shape=[seqlen, n_stream, n_stream],
strides=[dout_stride_0, dout_stride_1, dout_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)
res_desc = tl.make_tensor_descriptor(
res,
shape=[seqlen, n_stream, n_stream],
strides=[res_stride_0, res_stride_1, res_stride_2],
block_shape=[tilesize, n_stream, n_stream],
)

seq_off = tl.program_id(0) * tilesize

R = out_desc.load([seq_off, 0, 0])
RT = R.permute(0, 2, 1)
dR = dout_desc.load([seq_off, 0, 0])

RdR = R * dR
s_r = tl.sum(RdR, axis=-1).expand_dims(-1)
s_c = tl.sum(RdR, axis=-2).expand_dims(-1)

b = s_c - tl.sum(RT * s_r.permute(0, 2, 1), axis=-1).expand_dims(-1)

x = tl.zeros((tilesize, n_stream, 1), dtype=R.dtype)
r = b
p = r
r_normsq = tl.sum(r * r, axis=1, keep_dims=True)

for _ in range(n_stream):
Sp = matvec_S(R, p)
pSp = tl.sum(p * Sp, axis=1, keep_dims=True)
alpha = r_normsq / (pSp + EPS)

x += alpha * p
r -= alpha * Sp

r_new_normsq = tl.sum(r * r, axis=1, keep_dims=True)
beta = r_new_normsq / (r_normsq + EPS)

p = r + beta * p
r_normsq = r_new_normsq

u = s_r - tl.sum(R * x.permute(0, 2, 1), axis=-1).expand_dims(-1)
v = x

vt = v.reshape(tilesize, 1, n_stream)
M_mat = u + vt

res_tile = (dR - M_mat) * R
res_desc.store([seq_off, 0, 0], res_tile)

def sinkhorn_bwd_original(
out: torch.Tensor,
dout: torch.Tensor,
repeat: int = 1,
warmup: bool = True,
):
seqlen = out.size(0)
n_stream = out.size(1)

res = torch.empty_like(out)

def grid(META):
return (triton.cdiv(seqlen, META["tilesize"]), 1, 1)

def _run():
sinkhorn_bwd_original_kernel[grid](
seqlen,
out,
dout,
res,
out.stride(0),
out.stride(1),
out.stride(2),
dout.stride(0),
dout.stride(1),
dout.stride(2),
res.stride(0),
res.stride(1),
res.stride(2),
n_stream,
)

if warmup:
for _ in trange(4, desc="warmup (original)"):
_run()
torch.cuda.synchronize()

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start_event.record()
for _ in range(repeat):
_run()
end_event.record()
torch.cuda.synchronize()

elapsed_ms = start_event.elapsed_time(end_event)
if repeat > 1:
print(f"[original] Kernel execution time ({repeat=}): {elapsed_ms:.3f} ms")
print(f"[original] Average time per iteration: {elapsed_ms / repeat:.3f} ms")

return res

# ─────────────────────────────────────────────────────────────────────────────
# Main: accuracy comparison
# ─────────────────────────────────────────────────────────────────────────────
def main():
seqlen = 65536
n_stream = 8
iters = 10
repeat = 512

print(f"\nRunning {iters} iters...")

device = torch.device("cuda")
dist = torch.distributions.uniform.Uniform(0.0, 4.0)
M = dist.sample((seqlen, n_stream, n_stream)).to(device)
M.requires_grad_(True)

# ── Forward ──────────────────────────────────────────────────────────────
R, P = sinkhorn_forward(M, iters)
loss_weight = torch.randn_like(R)

# ── Method A: Autograd (ground truth) ────────────────────────────────────
loss_a = (R * loss_weight).sum()
loss_a.backward()
grad_autograd = M.grad.detach().clone()

grad_R = loss_weight # dL/dR

# ── Method B: Original triton_reduce.py (assumes r=1, c=1) ───────────────
grad_original = sinkhorn_bwd_original(R.detach(), grad_R, repeat=repeat, warmup=True)

# ── Method C: Corrected (fix.md scheme 2, uses actual r, c) ─────────────
grad_corrected = sinkhorn_bwd_corrected(R.detach(), grad_R, repeat=repeat, warmup=True)

# ── Accuracy comparison ───────────────────────────────────────────────────
def format_list(ls):
return [f"{x:.2e}" for x in ls]

def print_metrics(name, g1, g2):
abs_diff = (g1 - g2).abs()
rel_diff = abs_diff / (g1.abs() + 1e-12)

MAE = abs_diff.mean(dim=(-1, -2)).tolist()
max_abs_diff = abs_diff.reshape(seqlen, -1).max(-1).values.tolist()
mean_rel_diff = rel_diff.mean(dim=(-1, -2)).tolist()
max_rel_diff = rel_diff.reshape(seqlen, -1).max(-1).values.tolist()

print(f"\n[{name}]")
print(f"Max MAE = {max(MAE)}")
print(f"Max max_abs_diff = {max(max_abs_diff)}")
print(f"Max mean_rel_diff = {max(mean_rel_diff)}")
print(f"Max max_rel_diff = {max(max_rel_diff)}")
return max(MAE)

print("\nComparison of gradients dL/dM")
print("--------------------------------")
mae_orig = print_metrics("Original (r=1,c=1 assumption)", grad_autograd, grad_original)
mae_gen = print_metrics("Corrected (fix.md scheme 2)", grad_autograd, grad_corrected)

print("\nGrad (autograd) sample:\n", grad_autograd[0])
print("\nGrad (original) sample:\n", grad_original[0])
print("\nGrad (corrected) sample:\n", grad_corrected[0])
print(grad_autograd.dtype)
print(grad_corrected.dtype)

# ── Sanity check: row/col sums of R ──────────────────────────────────────
r_sums = R.sum(-1) # (seqlen, n)
c_sums = R.sum(-2) # (seqlen, n)
print(f"\nR row sums - mean={r_sums.mean():.4f}, std={r_sums.std():.4f}")
print(f"R col sums - mean={c_sums.mean():.4f}, std={c_sums.std():.4f}")
print("(If sums ≈ 1.0, original and generalized should agree)")

assert mae_gen < 5e-4, f"Corrected MAE too large: {mae_gen:.3e}"
print("\n✓ Accuracy check passed.")

if __name__ == "__main__":
main()

```

driver code

```python
"""
Compare the backward gradients of two Sinkhorn backward implementations:
1. tl_impl.py - TileLang implementation (layer-by-layer backprop)
2. tt_impl.py - Triton implementation (CG solver for corrected linear system)

Ground truth: PyTorch autograd
"""

import argparse
import torch
import sys

from tl_impl import _mhc_sinkhorn_bwd, _mhc_sinkhorn_fwd
from tt_impl import sinkhorn_bwd_corrected, sinkhorn_forward

def sinkhorn_forward_pytorch(M, iters=10, eps=1e-10):
P = torch.softmax(M, dim=-1) + eps
R = P / (P.sum(dim=-2, keepdim=True) + eps)
for _ in range(iters - 1):
R = R / (R.sum(dim=-1, keepdim=True) + eps)
R = R / (R.sum(dim=-2, keepdim=True) + eps)
return R

def main():
parser = argparse.ArgumentParser()
parser.add_argument("--iter", type=int, default=10, help="Number of Sinkhorn iterations")
args = parser.parse_args()

hidden_size = 8
num_tokens = 4096
token_block_size = 4
repeat = args.iter
eps = 1e-10

device = torch.device("cuda")
torch.manual_seed(42)

M = torch.randn(num_tokens, hidden_size, hidden_size, dtype=torch.float32, device=device)
grad_output = torch.randn(num_tokens, hidden_size, hidden_size, dtype=torch.float32, device=device)

print(f"Configuration: num_tokens={num_tokens}, hidden_size={hidden_size}, repeat={repeat}, eps={eps}")
print(f"token_block_size={token_block_size}")
print()

# PyTorch autograd (softmax forward)
M_ref = M.detach().clone().requires_grad_(True)
R_ref = sinkhorn_forward_pytorch(M_ref, iters=repeat, eps=eps)
(R_ref * grad_output).sum().backward()
grad_autograd = M_ref.grad.detach().clone()

# TileLang
fwd_kernel = _mhc_sinkhorn_fwd(hidden_size, token_block_size, repeat, eps)
R_tl = torch.empty(num_tokens, hidden_size, hidden_size, dtype=torch.float32, device=device)
fwd_kernel(M.detach(), R_tl)
torch.cuda.synchronize()

bwd_kernel = _mhc_sinkhorn_bwd(hidden_size, token_block_size, repeat, eps)
grad_tl = torch.empty(num_tokens, hidden_size, hidden_size, dtype=torch.float32, device=device)
bwd_kernel(grad_output, M.detach(), grad_tl)
torch.cuda.synchronize()

# Triton
M_tt = M.detach().clone().requires_grad_(True)
R_tt, _ = sinkhorn_forward(M_tt, iters=repeat)
(R_tt * grad_output).sum().backward()
grad_autograd_tt = M_tt.grad.detach().clone()

grad_tt = sinkhorn_bwd_corrected(R_tt.detach(), grad_output, repeat=1, warmup=False)
torch.cuda.synchronize()

# Comparison 1: TileLang vs PyTorch softmax
print("=== TileLang vs PyTorch (softmax forward) ===")
abs_diff = (grad_tl - grad_autograd).abs()
rel_diff = abs_diff / (grad_autograd.abs() + 1e-12)
print(f" MAE = {abs_diff.mean().item():.6e}")
print(f" MaxAE= {abs_diff.max().item():.6e}")
print(f" MRE = {rel_diff.mean().item():.6e}")
print(f" MaxRE= {rel_diff.max().item():.6e}")
print()

# Comparison 2: Triton vs PyTorch exp
print("=== Triton vs PyTorch (exp forward) ===")
abs_diff = (grad_tt - grad_autograd_tt).abs()
rel_diff = abs_diff / (grad_autograd_tt.abs() + 1e-12)
print(f" MAE = {abs_diff.mean().item():.6e}")
print(f" MaxAE= {abs_diff.max().item():.6e}")
print(f" MRE = {rel_diff.mean().item():.6e}")
print(f" MaxRE= {rel_diff.max().item():.6e}")
print()

print("Sample gradient (token 0):")
print(f" Autograd (softmax):\n{grad_autograd[0]}")
print(f" TileLang: \n{grad_tl[0]}")
print(f" Autograd (exp): \n{grad_autograd_tt[0]}")
print(f" Triton: \n{grad_tt[0]}")

if __name__ == "__main__":
main()
```

However,this heavily rely on the convergence of forward pass, which may not be guaranteed during training

Start by reading the existing triton_reduce.py backward implementation referenced in the issue, then compare it with the proposed corrected Triton kernel and its PyTorch autograd comparison in the issue body. Run the comparison from main with the stated CUDA configuration. Done means reaching agreement on whether the corrected backward approach is valid and what implementation change, if any, should be adopted.

Written by the indexing model from the issue text.