From a89288128f74d452512d83b8901029d4bb3481ac Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 30 Apr 2026 10:12:31 +0800 Subject: [PATCH 01/34] feat: unify mHC fused kernels --- megatron/core/fusions/fused_mhc_kernels.py | 2003 +++++++++++++---- megatron/core/fusions/triton_mhc_kernels.py | 395 ++++ megatron/core/transformer/hyper_connection.py | 107 +- .../core/transformer/transformer_config.py | 27 +- .../core/transformer/transformer_layer.py | 12 +- .../fusions/test_fused_mhc_kernels.py | 767 ++++++- .../test_fused_mhc_kernels_autotune.py | 237 ++ .../fusions/verify_compute_h_bwd.py | 198 ++ .../test_hyper_connection_recompute.py | 453 +++- 9 files changed, 3712 insertions(+), 487 deletions(-) create mode 100644 megatron/core/fusions/triton_mhc_kernels.py create mode 100644 tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py create mode 100644 tests/unit_tests/fusions/verify_compute_h_bwd.py diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 6a19255196a..c18e38e8a6a 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -24,10 +24,17 @@ # Check cuTile availability # --------------------------------------------------------------------------- _CUTILE_AVAILABLE = False +_CUTILE_EXPERIMENTAL_AVAILABLE = False try: import cuda.tile as ct _CUTILE_AVAILABLE = True + try: + import cuda.tile_experimental as ct_experimental + + _CUTILE_EXPERIMENTAL_AVAILABLE = True + except ImportError: + pass except ImportError: pass @@ -37,6 +44,24 @@ def is_cutile_available() -> bool: return _CUTILE_AVAILABLE +# --------------------------------------------------------------------------- +# Check Triton availability +# --------------------------------------------------------------------------- +_TRITON_AVAILABLE = False +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + pass + + +def is_triton_available() -> bool: + """Return True if Triton is available for supported mHC kernels.""" + return _TRITON_AVAILABLE + + # ============================================================================ # CuTile implementations (only defined when cuda.tile is available) # ============================================================================ @@ -112,22 +137,63 @@ def _ct_sinkhorn_bwd_kernel( grad = grad * M ct.store(grad_inp, index=(pid, 0, 0), tile=grad.astype(grad_inp.dtype)) + def _sinkhorn_autotune_tile_sizes(N_batch): + """Generate autotune search space for sinkhorn kernels.""" + for ts in (1, 2, 4, 8, 16, 32, 64, 128): + if ts <= N_batch: + yield ts + + _sinkhorn_fwd_best_cfg: dict = {} + _sinkhorn_bwd_best_cfg: dict = {} + def _cutile_sinkhorn_fwd( input_logits: Tensor, num_iterations: int, eps: float = 1e-8 ) -> Tuple[Tensor, Tensor]: original_shape = input_logits.shape hc = original_shape[-1] N_batch = input_logits.numel() // (hc * hc) - TILE_SIZE = math.gcd(N_batch, 128) dev = input_logits.device + stream = torch.cuda.current_stream() out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) - ct.launch( - torch.cuda.current_stream(), - (math.ceil(N_batch / TILE_SIZE), 1, 1), - _ct_sinkhorn_fwd_kernel, - (input_logits.view(N_batch, hc, hc), out, M_init, eps, hc, num_iterations, TILE_SIZE), - ) + inp = input_logits.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, ts), + ) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) + for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_fwd_kernel, + args_fn=lambda cfg: ( + inp, out, M_init, eps, hc, num_iterations, cfg.TILE_SIZE, + ), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_fwd_best_cfg[cache_key] = best_ts + ct.launch( + stream, + (math.ceil(N_batch / best_ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, best_ts), + ) + return out.view(original_shape), M_init.view(original_shape) def _cutile_sinkhorn_bwd( @@ -136,29 +202,55 @@ def _cutile_sinkhorn_bwd( original_shape = grad_output.shape hc = original_shape[-1] N_batch = grad_output.numel() // (hc * hc) - TILE_SIZE = math.gcd(N_batch, 128) dev = grad_output.device - ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) - ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) - ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + stream = torch.cuda.current_stream() grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) - ct.launch( - torch.cuda.current_stream(), - (math.ceil(N_batch / TILE_SIZE), 1, 1), - _ct_sinkhorn_bwd_kernel, - ( - grad_output.view(N_batch, hc, hc), - M_init.view(N_batch, hc, hc), - grad_input, - ws_M, - ws_rs, - ws_cs, - eps, - hc, - num_iterations, - TILE_SIZE, - ), - ) + go = grad_output.view(N_batch, hc, hc) + mi = M_init.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_bwd_best_cfg.get(cache_key) + + def _alloc_and_launch(ts): + ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_bwd_kernel, + (go, mi, grad_input, ws_M, ws_rs, ws_cs, eps, hc, num_iterations, ts), + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + _alloc_and_launch(ts) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) + for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + # Allocate workspace for largest tile size (all configs share same workspace shape). + ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_bwd_kernel, + args_fn=lambda cfg: ( + go, mi, grad_input, ws_M, ws_rs, ws_cs, eps, hc, num_iterations, + cfg.TILE_SIZE, + ), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_bwd_best_cfg[cache_key] = best_ts + # Re-launch with best config. + _alloc_and_launch(best_ts) + return grad_input.view(original_shape) # -- H_aggregate kernels ------------------------------------------------- @@ -194,18 +286,35 @@ def _ct_h_agg_bwd_kernel(go, x, h_pre, gx, gh, N: ConstInt, TILE_M: ConstInt, TI gh_acc += ct.sum(go_expanded * x_tile, axis=2) ct.store(gh, index=(pid, 0), tile=gh_acc.astype(gh.dtype)) + def _h_agg_autotune_configs(sb, C): + """Generate autotune search space for h_aggregate kernels.""" + for tile_m in (1, 2, 4, 8): + for tile_c in (32, 64, 128, 256, 512, 1024): + if tile_m <= sb and tile_c <= C: + yield {"TILE_M": tile_m, "TILE_C": tile_c} + + _h_agg_fwd_best_cfg: dict = {} + _h_agg_bwd_best_cfg: dict = {} + def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: s, b, n, C = x.shape sb = s * b - TILE_SIZE = math.gcd(sb, 4) - TILE_C = math.gcd(C, 1024) + stream = torch.cuda.current_stream() out = torch.empty(sb, C, dtype=x.dtype, device=x.device) - ct.launch( - torch.cuda.current_stream(), - (math.ceil(sb / TILE_SIZE),), - _ct_h_agg_fwd_kernel, - (x.view(sb, n, C), h_pre.view(sb, n), out, n, TILE_SIZE, TILE_C), - ) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + cache_key = (sb, n, C) + cached = _h_agg_fwd_best_cfg.get(cache_key) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + if cached is not None: + tm, tc = cached + else: + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch(stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, + (x_flat, h_flat, out, n, tm, tc)) + return out.view(s, b, C) def _cutile_h_aggregate_bwd( @@ -213,25 +322,24 @@ def _cutile_h_aggregate_bwd( ) -> Tuple[Tensor, Tensor]: s, b, n, C = x.shape sb = s * b - TILE_C = math.gcd(C, 1024) - TILE_M = math.gcd(sb, 4) + stream = torch.cuda.current_stream() gx = torch.empty(sb, n, C, dtype=x.dtype, device=x.device) - gh = torch.empty(sb, n, dtype=x.dtype, device=x.device) - ct.launch( - torch.cuda.current_stream(), - (math.ceil(sb / TILE_M),), - _ct_h_agg_bwd_kernel, - ( - grad_output.view(sb, C), - x.view(sb, n, C), - h_pre.view(sb, n), - gx, - gh, - n, - TILE_M, - TILE_C, - ), - ) + gh = torch.empty(sb, n, dtype=h_pre.dtype, device=x.device) + go_flat = grad_output.view(sb, C) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + cache_key = (sb, n, C) + cached = _h_agg_bwd_best_cfg.get(cache_key) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + if cached is not None: + tm, tc = cached + else: + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch(stream, (math.ceil(sb / tm),), _ct_h_agg_bwd_kernel, + (go_flat, x_flat, h_flat, gx, gh, n, tm, tc)) + return gx.view(s, b, n, C), gh.view(s, b, n) # -- H_post BDA kernels -------------------------------------------------- @@ -243,28 +351,22 @@ def _ct_hpb_fwd_kernel( pid = ct.bid(0) num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) - hp_2d = ct.reshape(hp_tile, (N, 1)) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) - hr_2d = ct.reshape(hr_tile, (N, N)) for ct_idx in range(num_c_tiles): orig_tile = ct.load( orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - orig_2d = ct.reshape(orig_tile, (N, TILE_C)) x_tile = ct.load( x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO ) - x_2d = ct.reshape(x_tile, (1, TILE_C)) - out_2d = hp_2d * x_2d + x_exp = ct.expand_dims(x_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * x_exp # (TILE_SIZE, N, TILE_C) for j in range(N): - out_2d += ct.extract(hr_2d, (0, j), shape=(N, 1)) * ct.extract( - orig_2d, (j, 0), shape=(1, TILE_C) - ) - ct.store( - out, - index=(pid, 0, ct_idx), - tile=ct.reshape(out_2d, (TILE_SIZE, N, TILE_C)).astype(out.dtype), - ) + hr_col = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) @ct.kernel def _ct_hpb_fwd_bias_kernel( @@ -273,210 +375,164 @@ def _ct_hpb_fwd_bias_kernel( pid = ct.bid(0) num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) - hp_2d = ct.reshape(hp_tile, (N, 1)) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) - hr_2d = ct.reshape(hr_tile, (N, N)) for ct_idx in range(num_c_tiles): orig_tile = ct.load( orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - orig_2d = ct.reshape(orig_tile, (N, TILE_C)) x_tile = ct.load( x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO ) bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) - xb_2d = ct.reshape(x_tile, (1, TILE_C)) + ct.reshape(bias_tile, (1, TILE_C)) - out_2d = hp_2d * xb_2d + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * xb_exp # (TILE_SIZE, N, TILE_C) for j in range(N): - out_2d += ct.extract(hr_2d, (0, j), shape=(N, 1)) * ct.extract( - orig_2d, (j, 0), shape=(1, TILE_C) - ) - ct.store( - out, - index=(pid, 0, ct_idx), - tile=ct.reshape(out_2d, (TILE_SIZE, N, TILE_C)).astype(out.dtype), - ) + hr_col = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) @ct.kernel - def _ct_hpb_bwd_kernel( - go, - hr, - orig, - hp, - x, - g_hr, - g_orig, - g_hp, - g_x, - N: ConstInt, - TILE_C: ConstInt, - TILE_SIZE: ConstInt, + def _ct_hpb_bwd_g_x_orig_kernel( + go, hr, hp, + g_orig, g_x, + N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, ): + """Compute g_x = hp @ go and g_orig = hr.T @ go. + + Grid: (ceil(sb / TILE_SIZE), ceil(C / TILE_C)). + 2D grid — no loop, no accumulators. + """ pid = ct.bid(0) - num_c_tiles = ct.cdiv(go.shape[2], TILE_C) - hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N)) - hp_2d = ct.reshape(hp_tile, (1, N)) + ct_idx = ct.bid(1) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) - hr_2d = ct.reshape(hr_tile, (N, N)) - acc_g_hp_2d = ct.full((N, 1), 0, dtype=ct.float32) - acc_g_hr_2d = ct.full((N, N), 0, dtype=ct.float32) + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + g_x_tile = ct.full((TILE_SIZE, 1, TILE_C), 0, dtype=hp.dtype) + g_orig_tile = ct.full((TILE_SIZE, N, TILE_C), 0, dtype=hp.dtype) + for j in range(N): + hp_j = ct.extract(hp_tile, (0, j), shape=(TILE_SIZE, 1)) + hp_j_exp = ct.expand_dims(hp_j, axis=2) # [TS, 1, 1] + go_j = ct.extract(go_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + g_x_tile = g_x_tile + hp_j_exp * go_j + hr_row_j = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + g_orig_tile = g_orig_tile + ct.reshape(hr_row_j, (TILE_SIZE, N, 1)) * go_j + ct.store( + g_x, index=(pid, ct_idx), + tile=ct.reshape(g_x_tile, (TILE_SIZE, TILE_C)).astype(g_x.dtype), + ) + ct.store(g_orig, index=(pid, 0, ct_idx), tile=g_orig_tile.astype(g_orig.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_hp_hr_kernel( + go, orig, x, + g_hr, g_hp, + N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, + ): + """Compute g_hp = sum(go * x) and g_hr = go @ orig.T (no bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) for ct_idx in range(num_c_tiles): - x_tile = ct.load( - x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO - ) - x_2d = ct.reshape(x_tile, (1, TILE_C)) + x_tile = ct.load(x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO) + x_exp = ct.expand_dims(x_tile, axis=1) # [TS, 1, TC] go_tile = ct.load( go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - go_2d = ct.reshape(go_tile, (N, TILE_C)) orig_tile = ct.load( orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - orig_2d = ct.reshape(orig_tile, (N, TILE_C)) - g_x_2d = ct.full((1, TILE_C), 0, dtype=hp.dtype) - g_orig_2d = ct.full((N, TILE_C), 0, dtype=hp.dtype) - for j in range(N): - g_x_2d += ct.extract(hp_2d, (0, j), shape=(1, 1)).item() * ct.extract( - go_2d, (j, 0), shape=(1, TILE_C) - ) - g_orig_2d += ct.extract(hr_2d, (j, 0), shape=(1, N)).reshape((N, 1)) * ct.extract( - go_2d, (j, 0), shape=(1, TILE_C) - ) - acc_g_hp_2d += ct.sum(go_2d * x_2d, axis=1, keepdims=True) - acc_g_hr_2d += ct.sum( - ct.expand_dims(go_2d, axis=1) * ct.expand_dims(orig_2d, axis=0), axis=2 - ) - ct.store( - g_x, - index=(pid, ct_idx), - tile=ct.reshape(g_x_2d, (TILE_SIZE, TILE_C)).astype(g_x.dtype), - ) - ct.store( - g_orig, - index=(pid, 0, ct_idx), - tile=ct.reshape(g_orig_2d, (TILE_SIZE, N, TILE_C)).astype(g_orig.dtype), + acc_g_hp = acc_g_hp + ct.sum(go_tile * x_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 ) ct.store( - g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp_2d, (TILE_SIZE, N)).astype(g_hp.dtype) - ) - ct.store( - g_hr, - index=(pid, 0, 0), - tile=ct.reshape(acc_g_hr_2d, (TILE_SIZE, N, N)).astype(g_hr.dtype), + g_hp, index=(pid, 0), + tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype), ) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) @ct.kernel - def _ct_hpb_bwd_bias_kernel( - go, - hr, - orig, - hp, - x, - bias, - g_hr, - g_orig, - g_hp, - g_x, - N: ConstInt, - TILE_C: ConstInt, - TILE_SIZE: ConstInt, + def _ct_hpb_bwd_g_hp_hr_bias_kernel( + go, orig, x, bias, + g_hr, g_hp, + N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, ): + """Compute g_hp = sum(go * (x+bias)) and g_hr = go @ orig.T (with bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ pid = ct.bid(0) num_c_tiles = ct.cdiv(go.shape[2], TILE_C) - hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N)) - hp_2d = ct.reshape(hp_tile, (1, N)) - hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) - hr_2d = ct.reshape(hr_tile, (N, N)) - acc_g_hp_2d = ct.full((N, 1), 0, dtype=ct.float32) - acc_g_hr_2d = ct.full((N, N), 0, dtype=ct.float32) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) for ct_idx in range(num_c_tiles): - x_tile = ct.load( - x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO - ) + x_tile = ct.load(x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO) bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) - xb_2d = ct.reshape(x_tile, (1, TILE_C)) + ct.reshape(bias_tile, (1, TILE_C)) + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # [TS, 1, TC] go_tile = ct.load( go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - go_2d = ct.reshape(go_tile, (N, TILE_C)) orig_tile = ct.load( orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - orig_2d = ct.reshape(orig_tile, (N, TILE_C)) - g_x_2d = ct.full((1, TILE_C), 0, dtype=hp.dtype) - g_orig_2d = ct.full((N, TILE_C), 0, dtype=hp.dtype) - for j in range(N): - g_x_2d += ct.extract(hp_2d, (0, j), shape=(1, 1)).item() * ct.extract( - go_2d, (j, 0), shape=(1, TILE_C) - ) - g_orig_2d += ct.extract(hr_2d, (j, 0), shape=(1, N)).reshape((N, 1)) * ct.extract( - go_2d, (j, 0), shape=(1, TILE_C) - ) - acc_g_hp_2d += ct.sum(go_2d * xb_2d, axis=1, keepdims=True) - acc_g_hr_2d += ct.sum( - ct.expand_dims(go_2d, axis=1) * ct.expand_dims(orig_2d, axis=0), axis=2 - ) - ct.store( - g_x, - index=(pid, ct_idx), - tile=ct.reshape(g_x_2d, (TILE_SIZE, TILE_C)).astype(g_x.dtype), - ) - ct.store( - g_orig, - index=(pid, 0, ct_idx), - tile=ct.reshape(g_orig_2d, (TILE_SIZE, N, TILE_C)).astype(g_orig.dtype), + acc_g_hp = acc_g_hp + ct.sum(go_tile * xb_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 ) ct.store( - g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp_2d, (TILE_SIZE, N)).astype(g_hp.dtype) - ) - ct.store( - g_hr, - index=(pid, 0, 0), - tile=ct.reshape(acc_g_hr_2d, (TILE_SIZE, N, N)).astype(g_hr.dtype), + g_hp, index=(pid, 0), + tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype), ) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) + + # -- H_post BDA autotune configs & caches -------------------------------- + + def _hpb_autotune_configs(sb, C): + """Generate TILE_SIZE × TILE_C search space for h_post_bda kernels.""" + for tile_size in (1, 2, 4, 8): + for tile_c in (32, 64, 128, 256, 512, 1024): + if tile_c <= C and tile_size <= sb: + yield {"TILE_SIZE": tile_size, "TILE_C": tile_c} + + _hpb_fwd_best_cfg: dict = {} + _hpb_bwd_g_x_orig_best_cfg: dict = {} + _hpb_bwd_g_hp_hr_best_cfg: dict = {} def _cutile_h_post_bda_fwd( h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: s, b, n, C = original_residual.shape sb = s * b - TILE_C = math.gcd(C, 1024) - TILE_SIZE = math.gcd(sb, 1) + stream = torch.cuda.current_stream() out = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) - grid = (math.ceil(sb / TILE_SIZE),) - if bias is not None: - ct.launch( - torch.cuda.current_stream(), - grid, - _ct_hpb_fwd_bias_kernel, - ( - h_res.view(sb, n, n), - original_residual.view(sb, n, C), - h_post.view(sb, n), - x.view(sb, C), - bias, - out, - n, - TILE_C, - TILE_SIZE, - ), - ) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + cache_key = (sb, n, C, bias is not None) + cached = _hpb_fwd_best_cfg.get(cache_key) + kernel = _ct_hpb_fwd_bias_kernel if bias is not None else _ct_hpb_fwd_kernel + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + if cached is not None: + ts, tc = cached else: - ct.launch( - torch.cuda.current_stream(), - grid, - _ct_hpb_fwd_kernel, - ( - h_res.view(sb, n, n), - original_residual.view(sb, n, C), - h_post.view(sb, n), - x.view(sb, C), - out, - n, - TILE_C, - TILE_SIZE, - ), - ) + ts, tc = 1, math.gcd(C, 1024) + args = (hr_flat, orig_flat, hp_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (out, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), kernel, args) + return out.view(s, b, n, C) def _cutile_h_post_bda_bwd( @@ -489,55 +545,94 @@ def _cutile_h_post_bda_bwd( ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: s, b, n, C = original_residual.shape sb = s * b - TILE_C = math.gcd(C, 1024) - TILE_SIZE = math.gcd(sb, 1) + stream = torch.cuda.current_stream() g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=h_res.device) - g_res = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) - g_hp = torch.empty(sb, n, dtype=h_res.dtype, device=h_res.device) - g_x = torch.empty(sb, C, dtype=h_res.dtype, device=h_res.device) - grid = (sb,) - if bias is not None: + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=h_res.device) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=h_res.device) + g_x = torch.empty(sb, C, dtype=x.dtype, device=h_res.device) + go_flat = grad_output.view(sb, n, C) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + # --- Kernel A: g_x, g_orig (2D grid, no loop) --- + cache_key_a = ('hpb_bwd_g_x_orig', sb, n, C) + cached_a = _hpb_bwd_g_x_orig_best_cfg.get(cache_key_a) + + if cached_a is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_a is not None: + ts, tc = cached_a + else: + ts, tc = 1, math.gcd(C, 1024) ct.launch( - torch.cuda.current_stream(), - grid, - _ct_hpb_bwd_bias_kernel, - ( - grad_output.view(sb, n, C), - h_res.view(sb, n, n), - original_residual.view(sb, n, C), - h_post.view(sb, n), - x.view(sb, C), - bias, - g_hr, - g_res, - g_hp, - g_x, - n, - TILE_C, - TILE_SIZE, - ), + stream, (math.ceil(sb / ts), math.ceil(C / tc)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, tc, ts), ) else: - ct.launch( - torch.cuda.current_stream(), - grid, - _ct_hpb_bwd_kernel, - ( - grad_output.view(sb, n, C), - h_res.view(sb, n, n), - original_residual.view(sb, n, C), - h_post.view(sb, n), - x.view(sb, C), - g_hr, - g_res, - g_hp, - g_x, - n, - TILE_C, - TILE_SIZE, + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE), math.ceil(C / cfg.TILE_C)), + kernel=_ct_hpb_bwd_g_x_orig_kernel, + args_fn=lambda cfg: ( + go_flat, hr_flat, hp_flat, g_res, g_x, n, cfg.TILE_C, cfg.TILE_SIZE, ), + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_x_orig_best_cfg[cache_key_a] = (best.TILE_SIZE, best.TILE_C) + ct.launch( + stream, (math.ceil(sb / best.TILE_SIZE), math.ceil(C / best.TILE_C)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, best.TILE_C, best.TILE_SIZE), ) - g_bias = g_x.sum(dim=0) if bias is not None else None + + # --- Kernel B: g_hp, g_hr (1D grid, loops C-tiles) --- + cache_key_b = ('hpb_bwd_g_hp_hr', sb, n, C, bias is not None) + cached_b = _hpb_bwd_g_hp_hr_best_cfg.get(cache_key_b) + hp_hr_kernel = _ct_hpb_bwd_g_hp_hr_bias_kernel if bias is not None else _ct_hpb_bwd_g_hp_hr_kernel + + if cached_b is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_b is not None: + ts, tc = cached_b + else: + ts, tc = 1, math.gcd(C, 1024) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), hp_hr_kernel, args) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + + def _hp_hr_args_fn(cfg): + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + return args + (g_hr, g_hp, n, cfg.TILE_C, cfg.TILE_SIZE) + + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE),), + kernel=hp_hr_kernel, + args_fn=_hp_hr_args_fn, + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_hp_hr_best_cfg[cache_key_b] = (best.TILE_SIZE, best.TILE_C) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, best.TILE_C, best.TILE_SIZE) + ct.launch(stream, (math.ceil(sb / best.TILE_SIZE),), hp_hr_kernel, args) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None return ( g_hr.view(s, b, n, n), g_res.view(s, b, n, C), @@ -549,10 +644,9 @@ def _cutile_h_post_bda_bwd( # -- Proj RMS kernels ---------------------------------------------------- @ct.function - def _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K): + def _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K, eps=1e-6): inv_norm = ct.where(norm_tile > 0, 1.0 / norm_tile, 0.0) inv_sqrt_k = 1.0 / ct.sqrt(K) - eps = 1e-8 u = norm_tile * inv_sqrt_k + eps coeff = -(1.0 / (u * u)) * inv_sqrt_k return dr_tile * coeff * a_tile * inv_norm @@ -571,12 +665,21 @@ def _ct_proj_rms_fwd_kernel( TILE_M: ConstInt, TILE_N: ConstInt, TILE_K: ConstInt, + SPLIT_K: ConstInt, ): + ''' + Grid: (num_tiles_m, num_tiles_k). Fused matmul + norm + r: proj, norm, r in one pass over K. + ''' tile_m_id = ct.bid(0) + split_k_id = ct.bid(1) + num_m_tiles = ct.cdiv(M, TILE_M) num_k_tiles = ct.cdiv(K, TILE_K) + num_k_tiles_per_split = ct.cdiv(num_k_tiles, SPLIT_K) + tile_k_id_start = split_k_id * num_k_tiles_per_split + tile_k_id_end = ct.minimum(tile_k_id_start + num_k_tiles_per_split, num_k_tiles) acc = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32) sum_sq = ct.full((TILE_M, 1), 0.0, dtype=ct.float32) - for tile_k_id in range(num_k_tiles): + for tile_k_id in range(tile_k_id_start, tile_k_id_end): a_tile = ct.load( A, index=(tile_m_id, tile_k_id), shape=(TILE_M, TILE_K), padding_mode=PAD_ZERO ) @@ -585,12 +688,119 @@ def _ct_proj_rms_fwd_kernel( a_tile.astype(ct.tfloat32), b_tile.transpose().astype(ct.tfloat32), acc=acc ) sum_sq += ct.sum(a_tile * a_tile, axis=1, keepdims=True) - norm_tile = ct.sqrt(sum_sq) - v = norm_tile / ct.sqrt(K) + eps - r_tile = 1.0 / v - ct.store(PROJ, index=(tile_m_id, 0), tile=acc.astype(PROJ.dtype)) - ct.store(NORM, index=(tile_m_id, 0), tile=norm_tile.astype(NORM.dtype)) - ct.store(R, index=(tile_m_id, 0), tile=r_tile.astype(R.dtype)) + # norm_tile = ct.sqrt(sum_sq) + + bid_m_k = tile_m_id + split_k_id * num_m_tiles + ct.store(PROJ, index=(bid_m_k, 0), tile=acc.astype(PROJ.dtype)) + ct.store(NORM, index=(bid_m_k, 0), tile=sum_sq.astype(NORM.dtype)) + # v = norm_tile / ct.sqrt(K) + eps + # r_tile = 1.0 / v + # ct.atomic_add(PROJ, proj_indices, acc.astype(PROJ.dtype)) + # ct.atomic_add(NORM, norm_indices, norm_tile.astype(NORM.dtype)) + # ct.store(R, index=(tile_m_id, 0), tile=r_tile.astype(R.dtype)) + + # -- Sigmoid helper for cuTile kernels ------------------------------------ + + @ct.function + def _ct_sigmoid(x): + """Sigmoid via exp2: σ(x) = 1 / (1 + 2^(-x * log2(e))).""" + return 1.0 / (1.0 + ct.exp2(-x * LOG2E)) + + # -- Reduce split-K + compute_h kernel ------------------------------------ + + @ct.kernel + def _ct_reduce_compute_h_kernel( + Y_acc, + R_acc, + Bias, + Alpha_pre, + Alpha_post, + Alpha_res, + Y, + R, + PROJ_OUT, + M: int, + N: int, + K: int, + n: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + SPLIT_K: ConstInt, + ): + """Reduce split-K partial proj/norm, compute r, and apply compute_h activations. + + Grid: (ceil(M / TILE_SIZE_M),). + TILE_SIZE_N = next_power_of_2(N) so one tile covers the full N dimension. + Alpha_{pre,post,res} are [1] tensors (scalar parameters). + """ + bid_m = ct.bid(0) + num_bid_m = ct.cdiv(M, TILE_SIZE_M) + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + # 1. Reduce split-K partials + y_accum = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0, dtype=ct.float32) + r_accum = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + y_tile = ct.load( + Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + y_accum = y_accum + ct.astype(y_tile, ct.float32) + + r_tile = ct.load( + R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_accum = r_accum + ct.astype(r_tile, ct.float32) + + # Store reduced proj for backward + ct.store(PROJ_OUT, index=(bid_m, 0), tile=y_accum.astype(PROJ_OUT.dtype)) + + # 2. Compute r = norm / sqrt(K) + denom = ct.full((TILE_SIZE_M, 1), K * 1.0, dtype=ct.float32) + mean = ct.truediv(r_accum, denom) + rstd = ct.rsqrt(mean) + ones = ct.full((TILE_SIZE_M, 1), 1.0, dtype=ct.float32) + r_val = ct.truediv(ones, rstd) # norm / sqrt(K) + + ct.store(R, index=(bid_m, 0), tile=r_val.astype(R.dtype)) + + # 3. Build per-column scale from 3 scalar alphas via masks + offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) + + one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) + zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) + + mask_pre = ct.where(ct.less(offsets, n), one, zero) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre + mask_res = one - mask_pre - mask_post + + scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res + scale = ct.reshape(scale, (1, TILE_SIZE_N)) + + mask_pre = ct.reshape(mask_pre, (1, TILE_SIZE_N)) + mask_post = ct.reshape(mask_post, (1, TILE_SIZE_N)) + mask_res = ct.reshape(mask_res, (1, TILE_SIZE_N)) + + # 4. Apply compute_h: linear = proj * scale / (r + eps) + bias + bias_tile = ct.load( + Bias, index=(0, 0), shape=(1, TILE_SIZE_N), padding_mode=PAD_ZERO, + ) + bias_tile = ct.astype(bias_tile, ct.float32) + + linear = y_accum * scale / (r_val + eps) + bias_tile + + # 5. Mask-based activations + sig = _ct_sigmoid(linear) + out = sig * mask_pre + 2.0 * sig * mask_post + linear * mask_res + + ct.store(Y, index=(bid_m, 0), tile=out.astype(Y.dtype)) @ct.kernel def _ct_proj_rms_bwd_kernel( @@ -604,6 +814,7 @@ def _ct_proj_rms_bwd_kernel( M: int, N: int, K: int, + eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, TILE_SIZE_K: ConstInt, @@ -626,7 +837,7 @@ def _ct_proj_rms_bwd_kernel( dr_tile = ct.load( DR, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=zero_pad ) - accumulator_da = accumulator_da + _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K) + accumulator_da = accumulator_da + _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K, eps) b_tile = ct.load( B, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=zero_pad ) @@ -643,7 +854,7 @@ def _ct_proj_rms_bwd_kernel( @ct.kernel def _ct_proj_rms_bwd_small_k_kernel( - A, B, NORM, DD, DR, DA, DB, M: int, N: int, K: int, TILE_N_SIZE: ConstInt + A, B, NORM, DD, DR, DA, DB, M: int, N: int, K: int, eps: float, TILE_N_SIZE: ConstInt ): zero_pad = ct.PaddingMode.ZERO TILE_DB_SIZE_M = 128 @@ -699,7 +910,7 @@ def _ct_proj_rms_bwd_small_k_kernel( DR, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad ) accumulator_da = accumulator_da + _ct_rms_dnorm( - a_tile.astype(ct.float32), norm_tile, dr_tile, K + a_tile.astype(ct.float32), norm_tile, dr_tile, K, eps ) b_tile = ct.load( B, @@ -729,26 +940,282 @@ def _next_power_of_2(n: int) -> int: n += 1 return n + def _proj_rms_fwd_autotune_configs(N): + """Generate autotune search space for proj_rms forward kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128) + split_ks = (1, 2, 4, 8, 16) + for tile_m in tile_ms: + for tile_k in tile_ks: + for split_k in split_ks: + yield {"TILE_M": tile_m, "TILE_N": TILE_N, "TILE_K": tile_k, 'SPLIT_K': split_k} + + # Cache the best config across calls (keyed by M, N, K). + _proj_rms_fwd_best_cfg: dict = {} + def _cutile_proj_rms_fwd( x: Tensor, weight: Tensor, eps: float = 1e-8 ) -> Tuple[Tensor, Tensor, Tensor]: M, K = x.shape N = weight.shape[0] - TILE_M = 128 TILE_N = _next_power_of_2(N) - TILE_K = 128 - num_tiles_m = math.ceil(M / TILE_M) - proj = torch.empty(M, N, dtype=x.dtype, device=x.device) - norm = torch.empty(M, 1, dtype=x.dtype, device=x.device) - r = torch.empty(M, 1, dtype=x.dtype, device=x.device) - ct.launch( - torch.cuda.current_stream(), - (num_tiles_m,), - _ct_proj_rms_fwd_kernel, - (x, weight, proj, norm, r, M, N, K, eps, TILE_M, TILE_N, TILE_K), - ) + dev = x.device + stream = torch.cuda.current_stream() + + cache_key = (M, N, K) + cached = _proj_rms_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + # Use cached best config, or fall back to default if no experimental. + if cached is not None: + tm, tn, tk, split_k = cached + else: + tm, tn, tk, split_k = 128, TILE_N, 128, 1 + + # print(cached) + # print('hahahahahah') + # split_k = 1 + proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), + ) + proj = proj.view(split_k, M, N).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(split_k, M, 1).sum(dim=0).to(dtype=x.dtype) + # r = r.view(split_k, M, 1).sum(dim=0) + else: + # Autotune on first call for this shape. + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] + # filter out configs with TILE_K > K or TILE_M > M + configs = [cfg for cfg in configs if cfg.TILE_K <= K and cfg.TILE_M <= M] + # if no valid config, set a default config with TILE_M = M and TILE_K = K and SPLIT_K = 1 + # if len(configs) == 0: + # configs = [SimpleNamespace(TILE_M=M, TILE_N=TILE_N, TILE_K=K, SPLIT_K=1)] + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, weight, proj, norm, r, M, N, K, eps, + cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + # print(best) + # print('hahahahahah') + # best.SPLIT_K = 1 + _proj_rms_fwd_best_cfg[cache_key] = (best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K) + proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + # Re-launch with best config for correct output. + ct.launch( + stream, + (math.ceil(M / best.TILE_M), best.SPLIT_K), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, + best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K), + ) + + proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) + norm = norm.view(best.SPLIT_K, M, 1).sum(dim=0) + norm = torch.sqrt(norm) + # r = r.view(split_k, M, 1).sum(dim=0) + r = 1.0 / (norm / math.sqrt(K) + eps) return proj, norm, r + # -- Reduce + compute_h launcher ------------------------------------------ + + def _reduce_compute_h_autotune_configs(M): + """Generate autotune search space for reduce_compute_h kernel.""" + for tile_m in (1, 2, 4, 8, 16, 32, 64, 128): + if tile_m <= M: + yield tile_m + + _reduce_compute_h_best_cfg: dict = {} + + def _cutile_reduce_compute_h( + proj_acc: Tensor, + norm_acc: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + M: int, + N: int, + K: int, + eps: float, + tile_n: int, + split_k: int, + ) -> Tuple[Tensor, Tensor, Tensor]: + """Launch reduce split-K + compute_h kernel. + + Returns: + y_activated: [M, N] with h_pre | h_post | h_res activated + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + dev = proj_acc.device + stream = torch.cuda.current_stream() + + bias_2d = bias.unsqueeze(0).contiguous() # [1, N] + + y_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) + r_out = torch.empty(M, 1, dtype=proj_acc.dtype, device=dev) + proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) + + cache_key = (M, N, K, n, split_k) + cached = _reduce_compute_h_best_cfg.get(cache_key) + + def _make_args(tm): + return ( + proj_acc, norm_acc, bias_2d, + alpha_pre, alpha_post, alpha_res, + y_out, r_out, proj_out, + M, N, K, n, eps, + tm, tile_n, split_k, + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + tm = cached if cached is not None else min(128, M) + ct.launch(stream, (math.ceil(M / tm),), _ct_reduce_compute_h_kernel, _make_args(tm)) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_M=tm) + for tm in _reduce_compute_h_autotune_configs(M) + ] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M),), + kernel=_ct_reduce_compute_h_kernel, + args_fn=lambda cfg: _make_args(cfg.TILE_M), + search_space=configs, + ) + best_tm = tuned.tuned_config.TILE_M + _reduce_compute_h_best_cfg[cache_key] = best_tm + ct.launch( + stream, (math.ceil(M / best_tm),), + _ct_reduce_compute_h_kernel, _make_args(best_tm), + ) + + return y_out, r_out, proj_out + + # -- Combined proj_rms + compute_h forward -------------------------------- + + def _cutile_proj_rms_compute_h_fwd( + x: Tensor, + weight: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + eps: float, + ) -> Tuple[Tensor, Tensor, Tensor]: + """Fused proj_rms + compute_h forward. + + Launches the existing _ct_proj_rms_fwd_kernel (split-K matmul + partial norm), + then _ct_reduce_compute_h_kernel (reduce + r + activations). + + Returns: + y_activated: [M, N] activated h values + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + dev = x.device + stream = torch.cuda.current_stream() + + cache_key = (M, N, K) + cached = _proj_rms_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk, split_k = cached + else: + tm, tn, tk, split_k = 128, TILE_N, 128, 1 + + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, tm, tn, tk, split_k), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] + configs = [cfg for cfg in configs if cfg.TILE_K <= K and cfg.TILE_M <= M] + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj_acc = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, + cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_fwd_best_cfg[cache_key] = (best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K) + tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, tm, tn, tk, split_k), + ) + + # Launch reduce + compute_h kernel + y_activated, r, proj_reduced = _cutile_reduce_compute_h( + proj_acc, norm_acc, bias, + alpha_pre, alpha_post, alpha_res, + n, M, N, K, eps, + TILE_N, split_k, + ) + return y_activated, r, proj_reduced + + def _proj_rms_bwd_autotune_configs(N): + """Generate autotune search space for proj_rms backward kernel (K >= 8192 path).""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128, 256) + for tile_m in tile_ms: + for tile_k in tile_ks: + yield {"TILE_SIZE_M": tile_m, "TILE_SIZE_N": TILE_N, "TILE_SIZE_K": tile_k} + + _proj_rms_bwd_best_cfg: dict = {} + def _cutile_proj_rms_bwd( grad_proj: Tensor, grad_r: Tensor, @@ -763,67 +1230,728 @@ def _cutile_proj_rms_bwd( db = torch.empty_like(weight) TILE_SIZE_N = _next_power_of_2(N) assert TILE_SIZE_N <= 256, f"TILE_SIZE_N too large: {TILE_SIZE_N}" - num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + stream = torch.cuda.current_stream() + if K >= 8192: - TILE_SIZE_M, TILE_SIZE_K = 128, 128 - grid = (math.ceil(K / TILE_SIZE_K), 1) - ct.launch( - torch.cuda.current_stream(), - grid, - _ct_proj_rms_bwd_kernel, - ( - x, - weight, - norm, - grad_proj, - grad_r, - da, - db, - M, - N, - K, - TILE_SIZE_M, - TILE_SIZE_N, - TILE_SIZE_K, - ), - ) + cache_key = (M, N, K) + cached = _proj_rms_bwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk = cached + else: + tm, tn, tk = 128, TILE_SIZE_N, 128 + ct.launch( + stream, + (math.ceil(K / tk), 1), + _ct_proj_rms_bwd_kernel, + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, tm, tn, tk), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_bwd_autotune_configs(N)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K), 1), + kernel=_ct_proj_rms_bwd_kernel, + args_fn=lambda cfg: ( + x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, + cfg.TILE_SIZE_M, cfg.TILE_SIZE_N, cfg.TILE_SIZE_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_bwd_best_cfg[cache_key] = ( + best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + ) + ct.launch( + stream, + (math.ceil(K / best.TILE_SIZE_K), 1), + _ct_proj_rms_bwd_kernel, + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, + best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K), + ) else: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count grid = (num_sms, 2, 1) ct.launch( - torch.cuda.current_stream(), + stream, grid, _ct_proj_rms_bwd_small_k_kernel, - (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, TILE_SIZE_N), + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, TILE_SIZE_N), ) return da, db + # -- Fused compute_h + proj_rms backward kernels ---------------------------- -# ============================================================================ -# Autograd Functions (cuTile only – guarded by _CUTILE_AVAILABLE) -# ============================================================================ + @ct.kernel + def _ct_fused_grad_h_proj_kernel( + GRAD_Y, # [M, N] + Y_ACT, # [M, N] + PROJ, # [M, N] + R, # [M, 1] + GRAD_R_EXT, # [M, 1] + Alpha_pre, # [1] + Alpha_post, # [1] + Alpha_res, # [1] + GRAD_H, # [M, TILE_SIZE_N] output + GRAD_PROJ, # [M, TILE_SIZE_N] output + GRAD_R_TOTAL, # [M, 1] output + M: int, + N: int, + n: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + ): + """Precompute grad_h, grad_proj, and grad_r_total for downstream backward kernels. -if not _CUTILE_AVAILABLE: + Grid: (ceil(M / TILE_SIZE_M),). + """ + tile_m_id = ct.bid(0) - def _no_cutile_error(*_args, **_kwargs): - raise RuntimeError( - "Fused mHC kernels require cuda.tile (cuTile) which is not installed. " - "Either install cuTile or set use_fused_mhc=False to use reference " - "implementations." + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) + one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) + zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) + mask_pre = ct.where(ct.less(offsets, n), one, zero) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre + mask_res = one - mask_pre - mask_post + + scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res + scale_2d = ct.reshape(scale, (1, TILE_SIZE_N)) + mask_pre_2d = ct.reshape(mask_pre, (1, TILE_SIZE_N)) + mask_post_2d = ct.reshape(mask_post, (1, TILE_SIZE_N)) + mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) + + gy_tile = ct.load( + GRAD_Y, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + gy_tile = ct.astype(gy_tile, ct.float32) + y_tile = ct.load( + Y_ACT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + y_tile = ct.astype(y_tile, ct.float32) + proj_tile = ct.load( + PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, ) + proj_tile = ct.astype(proj_tile, ct.float32) + r_tile = ct.load( + R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_tile = ct.astype(r_tile, ct.float32) + + # Activation backward → grad_h + gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d + half_y = y_tile * 0.5 + gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d + gh_res = gy_tile * mask_res_2d + grad_h = gh_pre + gh_post + gh_res + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_proj_tile = grad_h * scale_2d * inv_r_eps + + # grad_r_total: reduction over N → [TILE_M, 1] + grad_r_from_h = ct.sum( + grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), + axis=1, keepdims=True, + ) + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + grad_r_total = grad_r_from_h + grad_r_ext_tile - fused_sinkhorn = _no_cutile_error - fused_h_aggregate = _no_cutile_error - fused_h_post_bda = _no_cutile_error - fused_proj_rms = _no_cutile_error + ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_tile.astype(GRAD_PROJ.dtype)) + ct.store(GRAD_R_TOTAL, index=(tile_m_id, 0), tile=grad_r_total.astype(GRAD_R_TOTAL.dtype)) -else: + @ct.kernel + def _ct_fused_grad_x_weight_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_SIZE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + TILE_SIZE_K: ConstInt, + ): + """Compute grad_x and grad_weight simultaneously. - class FusedSinkhornKnopp(torch.autograd.Function): - """Fused Sinkhorn-Knopp projection to doubly stochastic matrix (cuTile).""" + Grid: (ceil(K / TILE_SIZE_K),). + Each block handles one K-tile and loops over all M-tiles. + Per M-tile: computes and stores grad_x, accumulates grad_weight. + """ + tile_k_id = ct.bid(0) + NUM_M_TILES = ct.cdiv(M, TILE_SIZE_M) + + # Load weight tile once — only depends on K-tile + weight_tile = ct.load( + WEIGHT, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), + padding_mode=PAD_ZERO, + ) + + acc_grad_weight = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for tile_m_id in range(NUM_M_TILES): + grad_proj_tile = ct.load( + GRAD_PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + x_tile = ct.load( + X, index=(tile_m_id, tile_k_id), shape=(TILE_SIZE_M, TILE_SIZE_K), + padding_mode=PAD_ZERO, + ) + grad_r_total = ct.load( + GRAD_R_TOTAL, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_tile = ct.load( + R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_tile = ct.astype(r_tile, ct.float32) + + # grad_x = grad_proj @ weight + grad_r_total * x / (r * K) + inv_rK = 1.0 / (r_tile * K) + acc_grad_x = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + acc_grad_x = ct.mma( + grad_proj_tile.astype(ct.tfloat32), + weight_tile.astype(ct.tfloat32), + acc=acc_grad_x, + ) + ct.store(GRAD_X, index=(tile_m_id, tile_k_id), tile=acc_grad_x.astype(GRAD_X.dtype)) + + # Accumulate grad_weight += x.T @ grad_proj + acc_grad_weight = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=acc_grad_weight, + ) + + ct.store(GRAD_WEIGHT, index=(0, tile_k_id), tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype)) + + @ct.kernel + def _ct_scalar_grads_kernel( + GRAD_H, # [M, TILE_SIZE_N] precomputed + PROJ, # [M, N] + R, # [M, 1] + GRAD_ALPHA_PRE, # [1] output (atomic) + GRAD_ALPHA_POST, # [1] output (atomic) + GRAD_ALPHA_RES, # [1] output (atomic) + GRAD_BIAS, # [TILE_SIZE_N] output (atomic) + M: int, + N: int, + n: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + ): + """Compute scalar gradients (grad_alpha_pre/post/res, grad_bias). + + Grid: (ceil(M / TILE_SIZE_M),). Each block processes one M-tile. + """ + bid_m = ct.bid(0) + + offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) + one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) + zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) + mask_pre = ct.where(ct.less(offsets, n), one, zero) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre + mask_res = one - mask_pre - mask_post + + mask_pre_2d = ct.reshape(mask_pre, (1, TILE_SIZE_N)) + mask_post_2d = ct.reshape(mask_post, (1, TILE_SIZE_N)) + mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) + + grad_h = ct.load( + GRAD_H, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + proj_tile = ct.load( + PROJ, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + proj_tile = ct.astype(proj_tile, ct.float32) + r_tile = ct.load( + R, index=(bid_m, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_tile = ct.astype(r_tile, ct.float32) + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + + ga_all = grad_h * proj_tile * inv_r_eps + ct.atomic_add(GRAD_ALPHA_PRE, 0, ct.sum(ga_all * mask_pre_2d).astype(GRAD_ALPHA_PRE.dtype)) + ct.atomic_add(GRAD_ALPHA_POST, 0, ct.sum(ga_all * mask_post_2d).astype(GRAD_ALPHA_POST.dtype)) + ct.atomic_add(GRAD_ALPHA_RES, 0, ct.sum(ga_all * mask_res_2d).astype(GRAD_ALPHA_RES.dtype)) + partial_gb = ct.sum(grad_h, axis=0, keepdims=False) + ct.atomic_add(GRAD_BIAS, offsets, partial_gb.astype(GRAD_BIAS.dtype)) + + @ct.kernel + def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_Y, # [M, N] + Y_ACT, # [M, N] + PROJ, # [M, N] + R, # [M, 1] + GRAD_R_EXT, # [M, 1] + Alpha_pre, # [1] + Alpha_post, # [1] + Alpha_res, # [1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + n: int, + eps: float, + TILE_N_SIZE: ConstInt, + ): + """Fused backward (small K path) with work-stealing. + + Grid: (num_sms, 2). + bid(1)==0: grad_weight via work-stealing over K-tiles, loops M. + bid(1)==1: grad_x via work-stealing over (M×K) tiles. + Scalar gradients computed by separate _ct_scalar_grads_kernel. + """ + zero_pad = ct.PaddingMode.ZERO + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + # Build masks and scale + offsets = ct.arange(TILE_N_SIZE, dtype=ct.int32) + one = ct.full((TILE_N_SIZE,), 1.0, dtype=ct.float32) + zero_v = ct.full((TILE_N_SIZE,), 0.0, dtype=ct.float32) + mask_pre = ct.where(ct.less(offsets, n), one, zero_v) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero_v) - mask_pre + mask_res = one - mask_pre - mask_post + + scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res + scale_2d = ct.reshape(scale, (1, TILE_N_SIZE)) + mask_pre_2d = ct.reshape(mask_pre, (1, TILE_N_SIZE)) + mask_post_2d = ct.reshape(mask_post, (1, TILE_N_SIZE)) + mask_res_2d = ct.reshape(mask_res, (1, TILE_N_SIZE)) + + TILE_DB_SIZE_M = 128 + TILE_DB_SIZE_K = 64 + NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) + NUM_K_TILES = ct.cdiv(K, TILE_DB_SIZE_K) + + if ct.bid(1) == 0: + # --- grad_weight path --- + for tile_id in range(ct.bid(0), NUM_K_TILES, ct.num_blocks(0)): + accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) + for m_tile in range(NUM_M_TILES): + x_tile = ct.load( + X, index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), padding_mode=zero_pad, + ) + # Compute grad_proj on-the-fly + gy_tile = ct.load( + GRAD_Y, index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + ) + gy_tile = ct.astype(gy_tile, ct.float32) + y_tile = ct.load( + Y_ACT, index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + ) + y_tile = ct.astype(y_tile, ct.float32) + r_tile = ct.load( + R, index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, 1), padding_mode=zero_pad, + ) + r_tile = ct.astype(r_tile, ct.float32) + + gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d + half_y = y_tile * 0.5 + gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d + gh_res = gy_tile * mask_res_2d + grad_h = gh_pre + gh_post + gh_res + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_proj_tile = grad_h * scale_2d * inv_r_eps + + accumulator_db = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=accumulator_db, + ) + + ct.store( + GRAD_WEIGHT, index=(0, tile_id), + tile=accumulator_db.transpose().astype(GRAD_WEIGHT.dtype), + allow_tma=False, + ) + + TILE_DA_SIZE_M = 128 + TILE_DA_SIZE_K = 256 + NUM_DA_TILES = ct.cdiv(M, TILE_DA_SIZE_M) * ct.cdiv(K, TILE_DA_SIZE_K) + NUM_DA_K_TILES = ct.cdiv(K, TILE_DA_SIZE_K) + + if ct.bid(1) == 1: + # --- grad_x path --- + for tile_id in range(ct.bid(0), NUM_DA_TILES, ct.num_blocks(0)): + b_tile_idx = tile_id % NUM_DA_K_TILES + dd_tile_idx = tile_id // NUM_DA_K_TILES + + # Compute grad_proj on-the-fly + gy_tile = ct.load( + GRAD_Y, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + ) + gy_tile = ct.astype(gy_tile, ct.float32) + y_tile = ct.load( + Y_ACT, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + ) + y_tile = ct.astype(y_tile, ct.float32) + proj_tile = ct.load( + PROJ, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + ) + proj_tile = ct.astype(proj_tile, ct.float32) + r_tile = ct.load( + R, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, + ) + r_tile = ct.astype(r_tile, ct.float32) + + gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d + half_y = y_tile * 0.5 + gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d + gh_res = gy_tile * mask_res_2d + grad_h = gh_pre + gh_post + gh_res + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_proj_tile = grad_h * scale_2d * inv_r_eps + + # Simplified rms_dnorm: grad_r_total * x / (r * K) + grad_r_from_h = ct.sum( + grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), + axis=1, keepdims=True, + ) + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + grad_r_total = grad_r_from_h + grad_r_ext_tile + + x_tile = ct.load( + X, index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), padding_mode=zero_pad, + ) + inv_rK = 1.0 / (r_tile * K) + accumulator_da = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + + weight_tile = ct.load( + WEIGHT, index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), padding_mode=zero_pad, + ) + accumulator_da = ct.mma( + grad_proj_tile.astype(ct.tfloat32), + weight_tile.astype(ct.tfloat32), + acc=accumulator_da, + ) + ct.store(GRAD_X, index=(dd_tile_idx, b_tile_idx), tile=accumulator_da.astype(GRAD_X.dtype)) + + def _fused_grad_x_weight_autotune_configs(N): + """Autotune search space for fused grad_x + grad_weight kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128, 256) + for tile_m in tile_ms: + for tile_k in tile_ks: + yield {"TILE_SIZE_M": tile_m, "TILE_SIZE_N": TILE_N, "TILE_SIZE_K": tile_k} + + _fused_grad_x_weight_best_cfg: dict = {} + + def _cutile_fused_compute_h_proj_rms_bwd( + x: Tensor, + weight: Tensor, + grad_y: Tensor, + y_activated: Tensor, + proj: Tensor, + r: Tensor, + grad_r_ext: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Fused compute_h + proj_rms backward. + + Returns: + grad_x: [M, K] + grad_weight: [N, K] + grad_alpha_pre: [1] + grad_alpha_post: [1] + grad_alpha_res: [1] + grad_bias: [N] + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + assert TILE_N <= 256, f"TILE_SIZE_N too large: {TILE_N}" + dev = x.device + stream = torch.cuda.current_stream() + + grad_x = torch.empty_like(x) + grad_weight = torch.empty_like(weight) + grad_alpha_pre_acc = torch.zeros(1, dtype=torch.float32, device=dev) + grad_alpha_post_acc = torch.zeros(1, dtype=torch.float32, device=dev) + grad_alpha_res_acc = torch.zeros(1, dtype=torch.float32, device=dev) + grad_bias_acc = torch.zeros(TILE_N, dtype=torch.float32, device=dev) + + if grad_r_ext is None: + grad_r_ext = torch.zeros(M, 1, dtype=x.dtype, device=dev) + + # 0. Precompute grad_h, grad_proj, grad_r_total + grad_h_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_proj_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_r_total_buf = torch.empty(M, 1, dtype=torch.float32, device=dev) + + tile_m_precomp = min(128, M) + ct.launch( + stream, + (math.ceil(M / tile_m_precomp),), + _ct_fused_grad_h_proj_kernel, + ( + grad_y, y_activated, proj, r, grad_r_ext, + alpha_pre, alpha_post, alpha_res, + grad_h_buf, grad_proj_buf, grad_r_total_buf, + M, N, n, eps, + tile_m_precomp, TILE_N, + ), + ) + + if K >= 8192: + # 1. Fused grad_x + grad_weight kernel — 1D grid (K-tiles), loops M + cache_key = ('grad_x_weight', M, N, K) + cached = _fused_grad_x_weight_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk = cached + else: + tm, tn, tk = 128, TILE_N, 128 + ct.launch( + stream, + (math.ceil(K / tk),), + _ct_fused_grad_x_weight_kernel, + ( + x, weight, grad_proj_buf, grad_r_total_buf, r, + grad_x, grad_weight, + M, N, K, + tm, tn, tk, + ), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _fused_grad_x_weight_autotune_configs(N)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K),), + kernel=_ct_fused_grad_x_weight_kernel, + args_fn=lambda cfg: ( + x, weight, grad_proj_buf, grad_r_total_buf, r, + grad_x, grad_weight, + M, N, K, + cfg.TILE_SIZE_M, cfg.TILE_SIZE_N, cfg.TILE_SIZE_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _fused_grad_x_weight_best_cfg[cache_key] = ( + best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + ) + ct.launch( + stream, + (math.ceil(K / best.TILE_SIZE_K),), + _ct_fused_grad_x_weight_kernel, + ( + x, weight, grad_proj_buf, grad_r_total_buf, r, + grad_x, grad_weight, + M, N, K, + best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + ), + ) + else: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + ct.launch( + stream, + (num_sms, 2, 1), + _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, + ( + x, weight, grad_y, y_activated, proj, r, grad_r_ext, + alpha_pre, alpha_post, alpha_res, + grad_x, grad_weight, + M, N, K, n, eps, + TILE_N, + ), + ) + + # 2. Separate lightweight kernel for scalar gradients (grad_alpha, grad_bias) + tile_m_scalar = min(128, M) + num_m_blocks = math.ceil(M / tile_m_scalar) + ct.launch( + stream, + (num_m_blocks,), + _ct_scalar_grads_kernel, + ( + grad_h_buf, proj, r, + grad_alpha_pre_acc, grad_alpha_post_acc, grad_alpha_res_acc, + grad_bias_acc, + M, N, n, eps, + tile_m_scalar, TILE_N, + ), + ) + + return ( + grad_x, + grad_weight, + grad_alpha_pre_acc.to(alpha_pre.dtype), + grad_alpha_post_acc.to(alpha_post.dtype), + grad_alpha_res_acc.to(alpha_res.dtype), + grad_bias_acc[:N].to(bias.dtype), + ) + + +# ============================================================================ +# Unified public dispatch +# ============================================================================ +# The public fused API chooses the fastest validated backend per operation: +# +# sinkhorn fwd/bwd: Triton -> cuTile -> torch +# h_post_bda fwd: cuTile -> torch +# h_post_bda bwd: Triton -> cuTile -> torch +# h_aggregate/proj_rms: cuTile -> torch +# +# Runtime CUDA launch failures are intentionally not swallowed; after such an +# error the CUDA context may not be safely reusable for fallback work. +# ============================================================================ + +from megatron.core.transformer.hyper_connection import ( + native_fused_add_3 as fused_add_3, + native_h_aggregate, + native_h_post_bda, + native_proj_rms, + native_sinkhorn, +) + + +def _get_triton_sinkhorn(): + if not _TRITON_AVAILABLE: + return None + try: + from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn + + return triton_fused_sinkhorn + except ImportError: + return None + + +def _get_triton_h_post_bda_bwd(): + if not _TRITON_AVAILABLE: + return None + try: + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + + return _triton_h_post_bda_bwd + except ImportError: + return None + + +def _torch_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + go = grad_output.reshape(sb, n, C) + hr = h_res.reshape(sb, n, n) + orig = original_residual.reshape(sb, n, C) + hp = h_post.reshape(sb, n) + x_flat = x.reshape(sb, C) + + g_hr = torch.bmm(go, orig.transpose(1, 2)).view(s, b, n, n) + g_res = torch.bmm(hr.transpose(1, 2), go).view(s, b, n, C) + g_x = torch.sum(go * hp.unsqueeze(-1), dim=1).view(s, b, C) + xb = x_flat if bias is None else x_flat + bias.view(1, C) + g_hp = torch.sum(go * xb.unsqueeze(1), dim=2).view(s, b, n) + g_bias = g_x.reshape(sb, C).sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.to(dtype=h_res.dtype), + g_res.to(dtype=original_residual.dtype), + g_hp.to(dtype=h_post.dtype), + g_x.to(dtype=x.dtype), + g_bias, + ) + + +def _torch_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, +) -> Tuple[Tensor, Tensor]: + proj = torch.matmul(x, weight.t()) + r = x.float().norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) + alpha = torch.cat( + [ + alpha_pre.expand(n), + alpha_post.expand(n), + alpha_res.expand(weight.shape[0] - 2 * n), + ], + dim=-1, + ) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + y_activated = torch.cat([h_pre, h_post, h[..., 2 * n :]], dim=-1) + return y_activated, r.to(dtype=x.dtype) + + +if _CUTILE_AVAILABLE: + + class CutileSinkhornKnopp(torch.autograd.Function): + """cuTile Sinkhorn-Knopp projection fallback.""" @staticmethod def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): - """cuTile fused Sinkhorn forward.""" output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) ctx.save_for_backward(M_init) ctx.num_iterations = num_iterations @@ -832,65 +1960,29 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output): - """cuTile fused Sinkhorn backward.""" (M_init,) = ctx.saved_tensors grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None - class FusedHAggregate(torch.autograd.Function): - """Fused n-stream weighted aggregation (cuTile).""" + class CutileHAggregate(torch.autograd.Function): + """cuTile n-stream weighted aggregation.""" @staticmethod def forward(ctx, x: Tensor, h_pre: Tensor): - """cuTile fused h_aggregate forward.""" output = _cutile_h_aggregate_fwd(x, h_pre) ctx.save_for_backward(x, h_pre) return output @staticmethod def backward(ctx, grad_output): - """cuTile fused h_aggregate backward.""" x, h_pre = ctx.saved_tensors return _cutile_h_aggregate_bwd(grad_output, x, h_pre) - class FusedHPostBDA(torch.autograd.Function): - """Fused: output = H_res @ orig_res + H_post * (x [+ bias]) (cuTile).""" - - @staticmethod - def forward( - ctx, - h_res: Tensor, - original_residual: Tensor, - h_post: Tensor, - x: Tensor, - bias: Optional[Tensor], - ): - """cuTile fused h_post_bda forward.""" - output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) - if bias is not None: - ctx.save_for_backward(h_res, original_residual, h_post, x, bias) - ctx.has_bias = True - else: - ctx.save_for_backward(h_res, original_residual, h_post, x) - ctx.has_bias = False - return output - - @staticmethod - def backward(ctx, grad_output): - """cuTile fused h_post_bda backward.""" - if ctx.has_bias: - h_res, orig_res, h_post, x, bias = ctx.saved_tensors - else: - h_res, orig_res, h_post, x = ctx.saved_tensors - bias = None - return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) - - class FusedProjRms(torch.autograd.Function): - """Fused projection + RMS normalization (cuTile).""" + class CutileProjRms(torch.autograd.Function): + """cuTile projection + RMS normalization.""" @staticmethod def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): - """cuTile fused proj_rms forward.""" proj, norm, r = _cutile_proj_rms_fwd(x, weight, eps) ctx.save_for_backward(x, weight, norm) ctx.eps = eps @@ -898,67 +1990,140 @@ def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): @staticmethod def backward(ctx, grad_proj, grad_r): - """cuTile fused proj_rms backward.""" x, weight, norm = ctx.saved_tensors grad_x, grad_weight = _cutile_proj_rms_bwd(grad_proj, grad_r, x, weight, norm, ctx.eps) return grad_x, grad_weight, None - # ======================================================================== - # Public API (only available when cuTile is installed) - # ======================================================================== + class CutileProjRmsComputeH(torch.autograd.Function): + """cuTile projection + RMS norm + compute_h activations.""" + + @staticmethod + def forward( + ctx, + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + ): + y_activated, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( + x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, + ) + ctx.save_for_backward( + x, weight, y_activated, proj_reduced, r, + alpha_pre, alpha_post, alpha_res, bias, + ) + ctx.n = n + ctx.eps = eps + return y_activated, r - def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: - """Project logits to doubly stochastic matrix via Sinkhorn-Knopp. + @staticmethod + def backward(ctx, grad_y, grad_r_ext): + ( + x, weight, y_activated, proj, r, + alpha_pre, alpha_post, alpha_res, bias_param, + ) = ctx.saved_tensors + + grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias = ( + _cutile_fused_compute_h_proj_rms_bwd( + x, weight, grad_y, y_activated, proj, r, grad_r_ext, + alpha_pre, alpha_post, alpha_res, + bias_param, ctx.n, ctx.eps, + ) + ) - Args: - input_logits: [..., n, n] raw logits - num_iterations: Sinkhorn iterations - eps: numerical stability + return grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None - Returns: - [..., n, n] doubly stochastic matrix - """ - return FusedSinkhornKnopp.apply(input_logits, num_iterations, eps) - def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: - """Weighted n-stream to 1-stream aggregation. +class FusedHPostBDA(torch.autograd.Function): + """H_post_bda with cuTile/torch forward and Triton/cuTile/torch backward.""" - Args: - x: [s, b, n, C] n-stream hidden states - h_pre: [s, b, n] aggregation weights + @staticmethod + def forward( + ctx, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ): + if _CUTILE_AVAILABLE: + output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) + else: + output = native_h_post_bda(h_res, original_residual, h_post, x, bias) + if bias is not None: + ctx.save_for_backward(h_res, original_residual, h_post, x, bias) + ctx.has_bias = True + else: + ctx.save_for_backward(h_res, original_residual, h_post, x) + ctx.has_bias = False + return output + + @staticmethod + def backward(ctx, grad_output): + if ctx.has_bias: + h_res, orig_res, h_post, x, bias = ctx.saved_tensors + else: + h_res, orig_res, h_post, x = ctx.saved_tensors + bias = None - Returns: - [s, b, C] aggregated hidden states - """ - return FusedHAggregate.apply(x, h_pre) + triton_bwd = _get_triton_h_post_bda_bwd() + if triton_bwd is not None: + return triton_bwd(grad_output, h_res, orig_res, h_post, x, bias) + if _CUTILE_AVAILABLE: + return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + return _torch_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) - def fused_h_post_bda( - h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] - ) -> Tensor: - """Fused H_res @ residual + H_post * (x + bias). - Args: - h_res: [s, b, n, n] residual mixing matrix - original_residual: [s, b, n, C] n-stream residual - h_post: [s, b, n] expansion weights - x: [s, b, C] layer output - bias: [C] or None +def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Project logits to a doubly stochastic matrix using Triton, cuTile, then torch.""" + triton_sinkhorn = _get_triton_sinkhorn() + if triton_sinkhorn is not None: + return triton_sinkhorn(input_logits, num_iterations, eps) + if _CUTILE_AVAILABLE: + return CutileSinkhornKnopp.apply(input_logits, num_iterations, eps) + return native_sinkhorn(input_logits, num_iterations, eps) - Returns: - [s, b, n, C] fused output - """ - return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) - def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: - """Fused projection + RMS normalization. +def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Weighted n-stream to 1-stream aggregation using cuTile, then torch.""" + if _CUTILE_AVAILABLE: + return CutileHAggregate.apply(x, h_pre) + return native_h_aggregate(x, h_pre) - Args: - x: [M, K] input - weight: [N, K] projection weight - eps: stability epsilon - Returns: - proj: [M, N] = x @ weight^T - r: [M, 1] = 1 / (||x|| / sqrt(K) + eps) - """ - return FusedProjRms.apply(x, weight, eps) +def fused_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """Fused H_res @ residual + H_post * (x + bias).""" + if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: + return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) + return native_h_post_bda(h_res, original_residual, h_post, x, bias) + + +def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Projection + RMS normalization using cuTile, then torch.""" + if _CUTILE_AVAILABLE: + return CutileProjRms.apply(x, weight, eps) + return native_proj_rms(x, weight, eps) + + +def fused_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, +) -> Tuple[Tensor, Tensor]: + """Projection + RMS norm + compute_h activations using cuTile, then torch.""" + if _CUTILE_AVAILABLE: + return CutileProjRmsComputeH.apply( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps + ) + return _torch_proj_rms_compute_h(x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps) diff --git a/megatron/core/fusions/triton_mhc_kernels.py b/megatron/core/fusions/triton_mhc_kernels.py new file mode 100644 index 00000000000..d21ea7905d0 --- /dev/null +++ b/megatron/core/fusions/triton_mhc_kernels.py @@ -0,0 +1,395 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Triton kernels used by the unified mHC fused path. + +Only the Triton kernels that are faster than the cuTile implementations are kept +here: + - Sinkhorn forward/backward + - H_post_bda backward + +The public mHC API lives in ``fused_mhc_kernels.py`` and handles fallback. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from torch import Tensor + +LOG2E = 1.4426950408889634 + + +# ============================================================================ +# Sinkhorn-Knopp +# ============================================================================ + + +@triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], + key=["HC", "NUM_ITERS"], +) +@triton.jit +def _triton_sinkhorn_fwd_kernel( + inp_ptr, out_ptr, M_init_ptr, + N_batch, eps, + HC: tl.constexpr, NUM_ITERS: tl.constexpr, +): + """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) + row_max = tl.max(logits, axis=1) + M = tl.exp2((logits - row_max[:, None]) * LOG2E) + tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) + + for _ in range(NUM_ITERS): + row_sum = tl.sum(M, axis=1) + M = M / (row_sum[:, None] + eps) + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + + tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], + key=["HC", "NUM_ITERS"], +) +@triton.jit +def _triton_sinkhorn_bwd_kernel( + grad_out_ptr, M_init_ptr, grad_inp_ptr, + ws_M_ptr, ws_rs_ptr, ws_cs_ptr, + N_batch, eps, + HC: tl.constexpr, NUM_ITERS: tl.constexpr, +): + """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + M_ws_base = pid * 2 * NUM_ITERS * HC * HC + v_ws_base = pid * NUM_ITERS + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + M = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + for t in range(NUM_ITERS): + ws_off = M_ws_base + (2 * t) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + row_sum = tl.sum(M, axis=1) + tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) + M = M / (row_sum[:, None] + eps) + + ws_off = M_ws_base + (2 * t + 1) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + col_sum = tl.sum(M, axis=0) + tl.store(ws_cs_ptr + (v_ws_base + t) * HC + offs_c, col_sum) + M = M / (col_sum[None, :] + eps) + + # M is the final forward output. It is the right value for the first VJP + # through the last column-normalization step. + grad = tl.load(grad_out_ptr + mat_ptrs).to(tl.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + + col_s = tl.load(ws_cs_ptr + (v_ws_base + t) * HC + offs_c).to(tl.float32) + grad = grad / (col_s[None, :] + eps) + col_corr = tl.sum(grad * M, axis=0) + grad = grad - col_corr[None, :] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t + 1) * HC * HC + + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) + grad = grad / (row_s[:, None] + eps) + row_corr = tl.sum(grad * M, axis=1) + grad = grad - row_corr[:, None] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t) * HC * HC + + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + grad = grad * M_init + tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) + + +def _triton_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-8 +) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.contiguous().view(N_batch, hc, hc) + _triton_sinkhorn_fwd_kernel[(N_batch,)]( + inp, out, M_init, N_batch, eps, hc, num_iterations, + ) + return out.view(original_shape), M_init.view(original_shape) + + +def _triton_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 +) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.contiguous().view(N_batch, hc, hc) + mi = M_init.contiguous().view(N_batch, hc, hc) + ws_M = torch.empty(N_batch * 2 * num_iterations * hc * hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + _triton_sinkhorn_bwd_kernel[(N_batch,)]( + go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations, + ) + return grad_input.view(original_shape) + + +class TritonFusedSinkhorn(torch.autograd.Function): + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, grad_output: Tensor): + (M_init,) = ctx.saved_tensors + grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + +def triton_fused_sinkhorn( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 +) -> Tensor: + return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) + + +# ============================================================================ +# H_post BDA backward +# ============================================================================ + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], +) +@triton.jit +def _triton_hpb_bwd_g_x_orig_kernel( + go_ptr, hr_ptr, hp_ptr, g_orig_ptr, g_x_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_go_s, stride_go_n, stride_go_c, + stride_hr_s, stride_hr_i, stride_hr_j, + stride_orig_s, stride_orig_n, stride_orig_c, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, +): + """g_x = hp @ go, g_orig = hr.T @ go.""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + g_x_acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, other=0.0, + ).to(tl.float32) + hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) + g_x_acc += hp_j[:, None] * go_j + tl.store( + g_x_ptr + offs_s[:, None] * C + offs_c[None, :], + g_x_acc.to(g_x_ptr.dtype.element_ty), + mask=mask_2d, + ) + + for i in tl.static_range(N): + g_orig_i = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, other=0.0, + ).to(tl.float32) + hr_ji = tl.load( + hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, + mask=mask_s, other=0.0, + ).to(tl.float32) + g_orig_i += hr_ji[:, None] * go_j + tl.store( + g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], + g_orig_i.to(g_orig_ptr.dtype.element_ty), + mask=mask_2d, + ) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], +) +@triton.jit +def _triton_hpb_bwd_g_hp_hr_kernel( + go_ptr, orig_ptr, x_ptr, bias_ptr, + g_hr_ptr, g_hp_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_go_s, stride_go_n, stride_go_c, + stride_orig_s, stride_orig_n, stride_orig_c, + stride_hr_s, stride_hr_i, stride_hr_j, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, +): + """g_hp = sum_c go*(x+bias), g_hr = go @ orig.T.""" + pid_s = tl.program_id(0) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + mask_s = offs_s < sb + + g_hp_acc = tl.zeros((BLOCK_S, N), dtype=tl.float32) + g_hr_acc = tl.zeros((BLOCK_S, N * N), dtype=tl.float32) + + for c_start in range(0, C, BLOCK_C): + offs_c = c_start + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + go_i = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + i * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hp = tl.sum(go_i * x_tile, axis=1) + g_hp_acc += tl.where( + tl.arange(0, N)[None, :] == i, + dot_hp[:, None], + tl.zeros((BLOCK_S, N), dtype=tl.float32), + ) + for j in tl.static_range(N): + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hr = tl.sum(go_i * orig_j, axis=1) + g_hr_acc += tl.where( + tl.arange(0, N * N)[None, :] == i * N + j, + dot_hr[:, None], + tl.zeros((BLOCK_S, N * N), dtype=tl.float32), + ) + + offs_n = tl.arange(0, N) + tl.store( + g_hp_ptr + offs_s[:, None] * N + offs_n[None, :], + g_hp_acc.to(g_hp_ptr.dtype.element_ty), + mask=mask_s[:, None], + ) + + nn_offs = tl.arange(0, N * N) + for i in tl.static_range(N): + for j in tl.static_range(N): + col_mask = (nn_offs == (i * N + j)).to(tl.float32) + val = tl.sum(g_hr_acc * col_mask[None, :], axis=1) + tl.store( + g_hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + val.to(g_hr_ptr.dtype.element_ty), + mask=mask_s, + ) + + +def _triton_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=dev) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=dev) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=dev) + g_x = torch.empty(sb, C, dtype=x.dtype, device=dev) + + go_flat = grad_output.contiguous().view(sb, n, C) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_bwd_g_x_orig_kernel[grid_a]( + go_flat, hr_flat, hp_flat, g_res, g_x, sb, C, n, + go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), + hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), + g_res.stride(0), g_res.stride(1), g_res.stride(2), + ) + + grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) + _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( + go_flat, orig_flat, x_flat, bias if bias is not None else x_flat, + g_hr, g_hp, sb, C, n, + go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), + orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), + g_hr.stride(0), g_hr.stride(1), g_hr.stride(2), + HAS_BIAS=(bias is not None), + ) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + +__all__ = [ + "_triton_h_post_bda_bwd", + "triton_fused_sinkhorn", +] diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index aa6a4bf7fbb..b5ee204d890 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -92,6 +92,37 @@ def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tenso return proj, r +@torch.compile +def native_fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: + """Native 3-way elementwise add (torch.compile fuses into single kernel).""" + return a + b + c + + +class BroadcastTensorFused(torch.autograd.Function): + """Split one tensor into 3 autograd-graph children sharing the same storage. + + During backward the three incoming gradients are summed with a caller- + supplied fused-add function (cuTile or torch.compile fallback) instead of + PyTorch's default sequential accumulation. + """ + + @staticmethod + def forward(ctx, x, fused_add_3_fn): + ctx.fused_add_3_fn = fused_add_3_fn + return x.view_as(x), x.view_as(x), x.view_as(x) + + @staticmethod + def backward(ctx, grad1, grad2, grad3): + grads = [g for g in (grad1, grad2, grad3) if g is not None] + if len(grads) == 0: + return None, None + if len(grads) == 1: + return grads[0], None + if len(grads) == 2: + return grads[0] + grads[1], None + return ctx.fused_add_3_fn(grad1, grad2, grad3), None + + @torch.compile def learned_output_contract( hidden_states: Tensor, head_fn: Tensor, base: Tensor, scale: Tensor, n: int, eps: float @@ -161,14 +192,19 @@ def __init__(self, config: TransformerConfig, layer_number: int): self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n)) self.norm_eps = 1e-6 - # Choose implementation: fused cuTile kernels vs reference modules. - # Both paths expose the same call signatures so the rest of the code - # is implementation-agnostic. + # Choose implementation: unified fused kernels vs reference modules. + # The fused public API selects the backend per operation internally. + # fused_add_3 always uses torch.compile (native_fused_add_3) regardless + # of the kernel backend. cuTile's register overhead (56 regs/thread for + # a trivial a+b+c) is not worth it for a pure memory-bound elementwise op. + self._fused_add_3_op = native_fused_add_3 + if config.use_fused_mhc: from megatron.core.fusions.fused_mhc_kernels import ( fused_h_aggregate, fused_h_post_bda, fused_proj_rms, + fused_proj_rms_compute_h, fused_sinkhorn, ) @@ -176,11 +212,13 @@ def __init__(self, config: TransformerConfig, layer_number: int): self._h_aggregate_op = fused_h_aggregate self._h_post_bda_op = fused_h_post_bda self._proj_rms_op = fused_proj_rms + self._proj_rms_compute_h_op = fused_proj_rms_compute_h else: self._sinkhorn_op = native_sinkhorn self._h_aggregate_op = native_h_aggregate self._h_post_bda_op = native_h_post_bda self._proj_rms_op = native_proj_rms + self._proj_rms_compute_h_op = None self._init_weights() @@ -258,12 +296,29 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: h_res: [s, b, n, n] - residual mixing matrix (doubly stochastic) """ s, b, _ = x.shape - with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"): - proj, r = self._projection_and_get_norm(x) - with torch.cuda.nvtx.range("HyperConnection::compute_h"): - h_pre, h_post, h_res = self._compute_h(proj, r) + + if self._proj_rms_compute_h_op is not None: + # Fused path: proj_rms + compute_h in one kernel launch sequence + x_2d = x.reshape(s * b, self.n * self.hidden_size) + with torch.cuda.nvtx.range("HyperConnection::fused_proj_rms_compute_h"): + y_activated, r = self._proj_rms_compute_h_op( + x_2d, self.mapping_proj.weight, + self.alpha_pre, self.alpha_post, self.alpha_res, + self.bias, self.n, self.norm_eps, + ) + h_pre = y_activated[:, : self.n].view(s, b, self.n) + h_post = y_activated[:, self.n : 2 * self.n].view(s, b, self.n) + h_res = y_activated[:, 2 * self.n :].view(s, b, self.n, self.n) + else: + # Native path: separate proj_rms + _compute_h + with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"): + proj, r = self._projection_and_get_norm(x) + with torch.cuda.nvtx.range("HyperConnection::compute_h"): + h_pre, h_post, h_res = self._compute_h(proj, r) + h_res = h_res.view(s, b, self.n, self.n) + h_res = self._sinkhorn_op( - h_res.view(s, b, self.n, self.n), self.sinkhorn_iterations, self.norm_eps + h_res, self.sinkhorn_iterations, self.norm_eps ) # [s, b, n, n] return h_pre, h_post, h_res @@ -394,10 +449,15 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: def forward( self, hidden_states: Tensor, mhc_recompute_manager: Optional['CheckpointManager'] = None - ) -> Tuple[Tensor, Tensor, Tensor]: + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """ Full mHC forward pass. + Uses BroadcastTensorFused to split hidden_states into 3 autograd-graph + children so that gradient accumulation from the 3 consumers + (compute_mappings, aggregate, fused_h_res_h_post_bda) is handled by a + single fused add instead of PyTorch's default sequential accumulation. + Args: hidden_states: [s, b, n*C] - n-stream hidden states mhc_recompute_manager: Optional CheckpointManager for checkpoint management. @@ -407,13 +467,14 @@ def forward( aggregated: [s, b, C] - aggregated input for layer computation h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda """ if mhc_recompute_manager is not None: return self._forward_with_checkpoint(hidden_states, mhc_recompute_manager) else: return self._forward_normal(hidden_states) - def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """ Normal forward pass without checkpointing. @@ -424,19 +485,25 @@ def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor aggregated: [s, b, C] - aggregated input for layer computation h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda """ + # Split into 3 views to avoid extra grad accumulations in backward + hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply( + hidden_states, self._fused_add_3_op + ) + # Compute mappings - h_pre, h_post, h_res = self.compute_mappings(hidden_states) + h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings) # Aggregate for layer input with torch.cuda.nvtx.range("HyperConnection::aggregate"): - aggregated = self.aggregate(hidden_states, h_pre) + aggregated = self.aggregate(hs_for_aggregate, h_pre) - return aggregated, h_res, h_post + return aggregated, h_res, h_post, hs_for_residual def _forward_with_checkpoint( self, hidden_states: Tensor, manager: 'CheckpointManager' - ) -> Tuple[Tensor, Tensor, Tensor]: + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """ Forward pass with checkpointing for memory efficiency. @@ -453,17 +520,23 @@ def _forward_with_checkpoint( aggregated: [s, b, C] - aggregated input for layer computation h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) h_post: [s, b, n] - expansion weights + residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda """ from megatron.core.tensor_parallel.random import CheckpointWithoutOutput - h_pre, h_post, h_res = self.compute_mappings(hidden_states) + # Split into 3 views to avoid extra grad accumulations in backward + hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply( + hidden_states, self._fused_add_3_op + ) + + h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings) # Checkpoint aggregate - auto-registers to manager aggregated = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( - self.aggregate, hidden_states, h_pre + self.aggregate, hs_for_aggregate, h_pre ) - return aggregated, h_res, h_post + return aggregated, h_res, h_post, hs_for_residual # ==================== Block-level utilities ==================== diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bd9777087bb..deefd9a6234 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1075,13 +1075,13 @@ class TransformerConfig(ModelParallelConfig): """Initial value of Gating Factor (alpha in paper).""" use_fused_mhc: bool = False - """Use cuTile fused kernels for mHC operations. + """Use unified fused kernels for mHC operations. When True, attempts to replace the reference mHC modules (SinkhornKnopp, - H_aggregate, H_post_bda, ProjRms) with fused cuda.tile (cuTile) autograd - functions for better performance on supported GPUs. Requires cuTile to be - installed; if cuTile is unavailable the flag is silently reset to False and - a warning is emitted. + H_aggregate, H_post_bda, ProjRms) with fused/autograd implementations for + better performance on supported GPUs. Backend selection is internal and + op-specific: Triton for Sinkhorn and H_post_bda backward when available, + cuTile for the remaining fused kernels when available, then torch fallback. """ mhc_recompute_layer_num: Optional[int] = None @@ -1928,23 +1928,6 @@ def __post_init__(self): if self.use_fused_mhc: if not self.enable_hyper_connections: raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") - try: - from megatron.core.fusions.fused_mhc_kernels import is_cutile_available - - if not is_cutile_available(): - warnings.warn( - "use_fused_mhc is enabled but cuda.tile (cuTile) is not installed. " - "Falling back to reference mHC implementations.", - UserWarning, - ) - self.use_fused_mhc = False - except ImportError: - warnings.warn( - "use_fused_mhc is enabled but fused_mhc_kernels module could not be " - "imported. Falling back to reference mHC implementations.", - UserWarning, - ) - self.use_fused_mhc = False if self.fine_grained_activation_offloading: assert ( diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 353979f97f9..1f0ad904fbd 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1650,11 +1650,11 @@ def _forward_attention( """Forward attention with hyper connection pre/post processing on self-attention.""" inference_context = deprecate_inference_params(inference_context, inference_params) - residual = hidden_states - nvtx_range_push(suffix="self_attention_hyper_connection") - hidden_states, self_attn_h_res, self_attn_hc_h_post = self.self_attention_hyper_connection( - hidden_states, mhc_recompute_manager=mhc_recompute_manager + hidden_states, self_attn_h_res, self_attn_hc_h_post, residual = ( + self.self_attention_hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) ) nvtx_range_pop(suffix="self_attention_hyper_connection") @@ -1748,10 +1748,8 @@ def _forward_mlp( ) mhc_mlp_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager - residual = hidden_states - nvtx_range_push(suffix="mlp_hyper_connection") - hidden_states, mlp_h_res, mlp_hc_h_post = self.mlp_hyper_connection( + hidden_states, mlp_h_res, mlp_hc_h_post, residual = self.mlp_hyper_connection( hidden_states, mhc_recompute_manager=mhc_recompute_manager ) nvtx_range_pop(suffix="mlp_hyper_connection") diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 91bfa8808df..9e5d65b41fc 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Unit tests for fused mHC kernels (cuTile) and native implementations. +"""Unit tests for unified fused mHC kernels and native implementations. Each test compares the fused kernel's forward output AND backward gradients against a pure-PyTorch differentiable reference to catch numerical drift @@ -14,7 +14,7 @@ import torch from torch import Tensor -from megatron.core.fusions.fused_mhc_kernels import is_cutile_available +from megatron.core.fusions.fused_mhc_kernels import is_cutile_available, is_triton_available from megatron.core.transformer.hyper_connection import ( native_h_aggregate, native_h_post_bda, @@ -23,6 +23,7 @@ ) _require_cutile = pytest.mark.skipif(not is_cutile_available(), reason="cuTile not installed") +_require_triton = pytest.mark.skipif(not is_triton_available(), reason="Triton not installed") @pytest.fixture(autouse=True) @@ -56,7 +57,14 @@ def _rand(*shape, **kwargs): def _info(): - backend = "cuTile" if is_cutile_available() else "native" + if is_triton_available() and is_cutile_available(): + backend = "triton+cuTile" + elif is_triton_available(): + backend = "triton+native" + elif is_cutile_available(): + backend = "cuTile" + else: + backend = "native" print(f"\n [backend: {backend}]") @@ -98,6 +106,28 @@ def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): return proj, r +def _ref_proj_rms_compute_h( + x: Tensor, weight: Tensor, + alpha_pre: Tensor, alpha_post: Tensor, alpha_res: Tensor, + bias: Tensor, n: int, eps: float = 1e-6, +): + """Reference: fused proj_rms + compute_h.""" + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = norm / math.sqrt(K) # [M, 1] + N = proj.shape[-1] + alpha = torch.cat([ + alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), + ]) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res = h[..., 2 * n :] + y = torch.cat([h_pre, h_post, h_res], dim=-1) + return y, r + + # ============================================================================ # Sinkhorn # ============================================================================ @@ -135,7 +165,7 @@ class TestFusedSinkhorn: @_require_cutile @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) def test_fwd_bwd_vs_reference(self, s, b, n, iters): - """E2E: fused cuTile fwd output and bwd grad must match the PyTorch reference.""" + """E2E: public fused fwd output and bwd grad must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import fused_sinkhorn _info() @@ -194,7 +224,7 @@ class TestFusedHAggregate: @_require_cutile @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) def test_fwd_bwd_vs_reference(self, s, b, n, C): - """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import fused_h_aggregate _info() @@ -276,7 +306,7 @@ class TestFusedHPostBDA: @pytest.mark.parametrize("with_bias", [True, False]) @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) def test_fwd_bwd_vs_reference(self, s, b, n, C, with_bias): - """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda _info() @@ -321,6 +351,204 @@ def _make_inputs(): ) +class TestTritonHPostBDABwd: + """Tests for Triton h_post_bda backward kernels against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)]) + def test_bwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb backward grads must match the PyTorch reference.""" + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + # -- Triton backward -- + g_hr_t, g_res_t, g_hp_t, g_x_t, g_bias_t = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + ) + + # -- Reference backward via autograd -- + hr_r = hr_data.clone().requires_grad_(True) + orig_r = orig_data.clone().requires_grad_(True) + hp_r = hp_data.clone().requires_grad_(True) + x_r = x_data.clone().requires_grad_(True) + bi_r = bias_data.clone().requires_grad_(True) if with_bias else None + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + for name, gt, gr in [ + ("h_res", g_hr_t, hr_r.grad), + ("orig_res", g_res_t, orig_r.grad), + ("h_post", g_hp_t, hp_r.grad), + ("x", g_x_t, x_r.grad), + ]: + torch.testing.assert_close( + gt, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"Triton backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + g_bias_t, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, + msg="Triton backward mismatch on bias", + ) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_triton_vs_cutile(self, s, b, n, C, with_bias): + """Triton and cuTile backward must produce identical results.""" + from megatron.core.fusions.fused_mhc_kernels import _cutile_h_post_bda_bwd + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + triton_out = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + ) + cutile_out = _cutile_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + ) + + for i, name in enumerate(["h_res", "orig_res", "h_post", "x", "bias"]): + if triton_out[i] is None: + continue + torch.testing.assert_close( + triton_out[i], cutile_out[i], + atol=BWD_ATOL, rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name}", + ) + + +class TestTritonHPostBDABwdE2EDebug: + """Debug: run E2E forward, then compare cuTile vs Triton backward per-output.""" + + @_require_triton + @_require_cutile + def test_e2e_inputs_no_nan(self): + """Feed actual E2E backward inputs to Triton kernel and check for NaN.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + # Run E2E forward to produce realistic backward inputs + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + x_2d = hs.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post_val = h[..., n : 2 * n].sigmoid() * 2 + h_res = fused_sinkhorn(h[..., 2 * n :].view(s, b, n, n), sinkhorn_iters, eps) + _ = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post_val, layer_out_data, layer_bias_data + ) + go = torch.ones_like(output) + + # Capture inputs (detach from graph) + hr = h_res.detach() + orig = hs.view(s, b, n, C).detach() + hp = h_post_val.detach() + x = layer_out_data.detach() + bias = layer_bias_data.detach() + + # Compare cuTile vs Triton per-output + ct_out = _cutile_h_post_bda_bwd(go, hr, orig, hp, x, bias) + tr_out = _triton_h_post_bda_bwd(go, hr, orig, hp, x, bias) + + names = ["g_hr", "g_res", "g_hp", "g_x", "g_bias"] + for name, ct_t, tr_t in zip(names, ct_out, tr_out): + if tr_t is None: + continue + assert not tr_t.isnan().any(), f"Triton {name} has NaN" + assert not tr_t.isinf().any(), f"Triton {name} has Inf" + torch.testing.assert_close( + tr_t, ct_t, atol=BWD_ATOL, rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name} (E2E inputs)", + ) + + + +# ============================================================================ +# Triton: Sinkhorn +# ============================================================================ + + +class TestTritonSinkhorn: + @_require_triton + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10), (8, 4, 4, 20)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_f = data.clone().requires_grad_(True) + out_f = triton_fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_f.grad, inp_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5)]) + def test_triton_vs_cutile(self, s, b, n, iters): + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_sinkhorn_bwd, + _cutile_sinkhorn_fwd, + ) + from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_t = data.clone().requires_grad_(True) + out_t = triton_fused_sinkhorn(inp_t, iters, eps) + out_t.backward(grad_out) + + out_c, M_init = _cutile_sinkhorn_fwd(data.clone(), iters, eps) + grad_c = _cutile_sinkhorn_bwd(grad_out, M_init, iters, eps) + + torch.testing.assert_close(out_t, out_c, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_t.grad, grad_c, atol=BWD_ATOL, rtol=BWD_RTOL) + + # ============================================================================ # Proj RMS # ============================================================================ @@ -363,7 +591,7 @@ class TestFusedProjRms: @_require_cutile @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) def test_fwd_bwd_vs_reference(self, M, N, K): - """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms _info() @@ -395,6 +623,74 @@ def test_fwd_bwd_vs_reference(self, M, N, K): ) +# ============================================================================ +# Proj RMS + Compute H (fused) +# ============================================================================ + + +class TestFusedProjRmsComputeH: + @pytest.mark.parametrize( + "M,n,K", + [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)], + ) + def test_fwd_bwd_vs_reference(self, M, n, K): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h + + _info() + N = n * n + 2 * n + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + grad_y = _rand(M, N) + grad_r = _rand(M, 1) + + def _make_inputs(): + return ( + x_data.clone().requires_grad_(True), + w_data.clone().requires_grad_(True), + ap_data.clone().requires_grad_(True), + apo_data.clone().requires_grad_(True), + ar_data.clone().requires_grad_(True), + bias_data.clone().requires_grad_(True), + ) + + # -- fused path -- + xf, wf, apf, apof, arf, bf = _make_inputs() + y_f, r_f = fused_proj_rms_compute_h(xf, wf, apf, apof, arf, bf, n, eps) + (y_f * grad_y + r_f * grad_r).sum().backward() + + # -- reference path -- + xr, wr, apr, apor, arr, br = _make_inputs() + y_r, r_r = _ref_proj_rms_compute_h(xr, wr, apr, apor, arr, br, n, eps) + (y_r * grad_y + r_r * grad_r).sum().backward() + + torch.testing.assert_close(y_f, y_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="y_activated mismatch") + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="r mismatch") + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + _assert_cosine_similar( + apf.grad, apr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_pre" + ) + _assert_cosine_similar( + apof.grad, apor.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_post" + ) + _assert_cosine_similar( + arf.grad, arr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_res" + ) + _assert_cosine_similar( + bf.grad, br.grad, COSINE_SIM_THRESH, msg="backward mismatch on bias" + ) + + # ============================================================================ # End-to-end pipeline (all four kernels chained) # ============================================================================ @@ -483,7 +779,7 @@ def _run_inline_ref(): class TestEndToEndFused: - """Full mHC pipeline using fused cuTile kernels (requires cuTile).""" + """Full mHC pipeline using the public fused API.""" @pytest.mark.flaky_in_dev @_require_cutile @@ -496,7 +792,7 @@ def test_full_pipeline_fwd_bwd(self): ) _info() - s, b, n, C = 2, 4, 4, 1024 + s, b, n, C = 8, 4, 4, 1024 eps = 1e-6 sinkhorn_iters = 5 @@ -528,7 +824,7 @@ def _run_fused(): loss = output.sum() + aggregated.sum() loss.backward() - return output.detach(), aggregated.detach(), hs.grad.clone() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() def _run_ref(): hs = hs_data.clone().requires_grad_(True) @@ -551,19 +847,464 @@ def _run_ref(): h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data ) + loss = output.sum() + aggregated.sum() + loss.backward() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() + + + proj_f, r_f, out_f, agg_f, grad_f = _run_fused() + proj_r, r_r, out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL + ) + torch.testing.assert_close( + r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL + ) + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + def test_full_pipeline_fused_compute_h(self): + """E2E: fused proj_rms_compute_h replaces separate proj_rms + compute_h.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + y_activated, r = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps, + ) + + h_pre = y_activated[:, :n].view(s, b, n) + h_post = y_activated[:, n : 2 * n].view(s, b, n) + h_res_logits = y_activated[:, 2 * n :].view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + y_ref, r_ref = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps, + ) + + h_pre = y_ref[:, :n].view(s, b, n) + h_post = y_ref[:, n : 2 * n].view(s, b, n) + h_res_logits = y_ref[:, 2 * n :].view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + loss = output.sum() + aggregated.sum() loss.backward() return output.detach(), aggregated.detach(), hs.grad.clone() - out_f, agg_f, grad_f = _run_fused() + out_f, agg_f, grad_f = _run_fused_compute_h() out_r, agg_r, grad_r = _run_ref() torch.testing.assert_close( - agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated output mismatch" + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch" ) torch.testing.assert_close( out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" ) _assert_cosine_similar( - grad_f, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + grad_f, grad_r, COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h)", + ) + + +# ============================================================================ +# fused_add_3 kernel tests +# ============================================================================ + + +class TestFusedAdd3: + """Tests for fused_add_3 (torch.compile backend, no cuTile dependency).""" + + def test_fused_add_3_forward_bf16(self): + """fused_add_3 matches a + b + c for bf16 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + b = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + c = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_forward_fp32(self): + """fused_add_3 matches a + b + c for fp32 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + b = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + c = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + result = fused_add_3(a, b, c) + expected = a + b + c + torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) + + def test_fused_add_3_large_tensor(self): + """fused_add_3 handles large tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + b = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + c = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_gradient(self): + """fused_add_3 produces correct gradients.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + b = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + c = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + result = fused_add_3(a, b, c) + result.sum().backward() + torch.testing.assert_close(a.grad, torch.ones_like(a)) + torch.testing.assert_close(b.grad, torch.ones_like(b)) + torch.testing.assert_close(c.grad, torch.ones_like(c)) + + +# ============================================================================ +# End-to-end pipeline with BroadcastTensorFused +# ============================================================================ + + +class TestEndToEndNativeBroadcast: + """Full mHC pipeline using native modules + BroadcastTensorFused. + + Same pipeline as TestEndToEndNative but hidden_states is split via + BroadcastTensorFused so each consumer (proj_rms/compute_h, aggregate, + h_post_bda) gets a distinct autograd graph node. Verifies gradient + correctness versus the inline reference that uses the tensor directly. + """ + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + # Split via BroadcastTensorFused + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, native_fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_b, agg_b, grad_b = _run_broadcast() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_b, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (broadcast)" + ) + torch.testing.assert_close( + out_b, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch (broadcast)" + ) + _assert_cosine_similar( + grad_b, grad_r, COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, broadcast)", + ) + + +class TestEndToEndFusedBroadcast: + """Full mHC pipeline using the public fused API + BroadcastTensorFused.""" + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (fused broadcast)" + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused broadcast)", + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused broadcast)", + ) + + def test_full_pipeline_fused_compute_h_broadcast(self): + """E2E: fused proj_rms_compute_h + BroadcastTensorFused.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + y_activated, r = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps, + ) + + h_pre = y_activated[:, :n].view(s, b, n) + h_post = y_activated[:, n : 2 * n].view(s, b, n) + h_res_logits = y_activated[:, 2 * n :].view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + y_ref, r_ref = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps, + ) + + h_pre = y_ref[:, :n].view(s, b, n) + h_post = y_ref[:, n : 2 * n].view(s, b, n) + h_res_logits = y_ref[:, 2 * n :].view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_compute_h_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, + msg="aggregated mismatch (fused compute_h broadcast)", + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused compute_h broadcast)", + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h broadcast)", ) diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py b/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py new file mode 100644 index 00000000000..fd42be33c05 --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py @@ -0,0 +1,237 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for cuTile mHC kernel autotuning correctness. + +Each test clears the autotune cache, runs the fused kernel (triggering a fresh +autotune search), and verifies the output matches a pure-PyTorch reference. +This ensures every config in the search space that the autotuner might select +produces numerically correct results. + +Usage: + pytest tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py -s +""" + +import math + +import pytest +import torch +from torch import Tensor + +from megatron.core.fusions.fused_mhc_kernels import ( + _CUTILE_EXPERIMENTAL_AVAILABLE, + is_cutile_available, +) + +_require_cutile_experimental = pytest.mark.skipif( + not (is_cutile_available() and _CUTILE_EXPERIMENTAL_AVAILABLE), + reason="cuTile + cuda.tile_experimental required for autotune tests", +) + + +@pytest.fixture(autouse=True) +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" +FWD_ATOL, FWD_RTOL = 2e-2, 2e-2 +BWD_ATOL, BWD_RTOL = 5e-2, 5e-2 + + +def _rand(*shape, requires_grad=False): + return ( + torch.empty(*shape, dtype=DTYPE, device=DEVICE) + .uniform_(-0.1, 0.1) + .requires_grad_(requires_grad) + ) + + +# --------------------------------------------------------------------------- +# Pure-PyTorch references +# --------------------------------------------------------------------------- + + +def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: + row_max = logits.max(dim=-1, keepdim=True).values + M = torch.exp(logits - row_max) + for _ in range(num_iters): + M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps) + M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps) + return M + + +def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = 1.0 / (norm / math.sqrt(K) + eps) + return proj, r + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _clear_autotune_caches(): + """Clear all cuTile mHC autotune caches to force a fresh search.""" + from megatron.core.fusions import fused_mhc_kernels as mod + + mod._sinkhorn_fwd_best_cfg.clear() + mod._sinkhorn_bwd_best_cfg.clear() + mod._proj_rms_fwd_best_cfg.clear() + mod._proj_rms_bwd_best_cfg.clear() + + +# ============================================================================ +# Sinkhorn autotune +# ============================================================================ + + +class TestAutotunesSinkhorn: + @_require_cutile_experimental + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (4, 8, 4, 20)]) + def test_sinkhorn_fwd_autotune(self, s, b, n, iters): + """Autotuned sinkhorn fwd must match PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import _cutile_sinkhorn_fwd + + _clear_autotune_caches() + + eps = 1e-6 + data = _rand(s, b, n, n) + + out_fused, _ = _cutile_sinkhorn_fwd(data.clone(), iters, eps) + + out_ref = _ref_sinkhorn(data.clone(), iters, eps) + + torch.testing.assert_close(out_fused, out_ref, atol=FWD_ATOL, rtol=FWD_RTOL) + + @_require_cutile_experimental + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (4, 8, 4, 20)]) + def test_sinkhorn_bwd_autotune(self, s, b, n, iters): + """Autotuned sinkhorn bwd grads must match PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_sinkhorn_bwd, + _cutile_sinkhorn_fwd, + ) + + _clear_autotune_caches() + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + out_f, M_init = _cutile_sinkhorn_fwd(data.clone(), iters, eps) + grad_f = _cutile_sinkhorn_bwd(grad_out, M_init, iters, eps) + + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + @_require_cutile_experimental + def test_sinkhorn_autotune_cache_populated(self): + """After first call, autotune cache should contain the best config.""" + from megatron.core.fusions import fused_mhc_kernels as mod + from megatron.core.fusions.fused_mhc_kernels import _cutile_sinkhorn_fwd + + _clear_autotune_caches() + + s, b, n, iters = 2, 4, 4, 5 + data = _rand(s, b, n, n) + + # First call triggers autotune + _cutile_sinkhorn_fwd(data, iters) + N_batch = s * b + assert (N_batch, n, iters) in mod._sinkhorn_fwd_best_cfg + + # Second call should use cache (no error) + _cutile_sinkhorn_fwd(data, iters) + + +# ============================================================================ +# Proj RMS autotune +# ============================================================================ + + +class TestAutotunesProjRms: + @_require_cutile_experimental + @pytest.mark.parametrize("M,N,K", [(256, 24, 4096), (64, 8, 512)]) + def test_proj_rms_fwd_autotune(self, M, N, K): + """Autotuned proj_rms fwd must match PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms + + _clear_autotune_caches() + + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + + xf = x_data.clone() + wf = w_data.clone() + proj_f, r_f = fused_proj_rms(xf, wf, eps) + + xr = x_data.clone() + wr = w_data.clone() + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + @_require_cutile_experimental + @pytest.mark.parametrize("M,N,K", [(256, 24, 16384), (64, 8, 512)]) + def test_proj_rms_bwd_autotune(self, M, N, K): + """Autotuned proj_rms bwd grads must match PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms + + _clear_autotune_caches() + + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = fused_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + @_require_cutile_experimental + def test_proj_rms_autotune_cache_populated(self): + """After first call, autotune caches should contain the best configs.""" + from megatron.core.fusions import fused_mhc_kernels as mod + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms + + _clear_autotune_caches() + + M, N, K = 256, 24, 16384 + eps = 1e-6 + x = _rand(M, K, requires_grad=True) + w = _rand(N, K, requires_grad=True) + + proj, r = fused_proj_rms(x, w, eps) + assert (M, N, K) in mod._proj_rms_fwd_best_cfg, "fwd cache not populated" + + (proj.sum() + r.sum()).backward() + # bwd autotune only triggers for K >= 8192 + assert (M, N, K) in mod._proj_rms_bwd_best_cfg, "bwd cache not populated" diff --git a/tests/unit_tests/fusions/verify_compute_h_bwd.py b/tests/unit_tests/fusions/verify_compute_h_bwd.py new file mode 100644 index 00000000000..6d29d090c8a --- /dev/null +++ b/tests/unit_tests/fusions/verify_compute_h_bwd.py @@ -0,0 +1,198 @@ +"""Standalone verification of compute_h backward formulas. + +Pure PyTorch — no cuTile needed. Compares hand-derived gradients against autograd. +""" + +import math + +import torch + +torch.manual_seed(42) +DEVICE = "cpu" +DTYPE = torch.float64 # float64 for precise gradient checking + + +def forward_proj_rms_compute_h(x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps): + """Forward: proj_rms + compute_h. All differentiable.""" + proj = x @ weight.t() # [M, N] + norm = x.norm(dim=-1, keepdim=True) # [M, 1] + K = x.shape[-1] + N = proj.shape[-1] + r = norm / math.sqrt(K) # [M, 1] + + # Build alpha vector from 3 scalars + alpha = torch.cat([ + alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), + ]) + + # compute_h + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + h_post = h[..., n:2 * n].sigmoid() * 2 + h_res = h[..., 2 * n:] + y = torch.cat([h_pre, h_post, h_res], dim=-1) + return y, r, proj, norm + + +def manual_backward_compute_h( + grad_y, grad_r_ext, y_activated, proj, r, + alpha_pre, alpha_post, alpha_res, n, eps, +): + """Hand-derived backward through compute_h + r conversion. + + Returns: grad_proj, grad_r_for_bwd, grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias + """ + M, N = proj.shape + + # Build alpha vector + alpha = torch.cat([ + alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), + ]) + + # 1. Activation backward -> grad_h + y_pre = y_activated[:, :n] + y_post = y_activated[:, n:2 * n] + gy_pre = grad_y[:, :n] + gy_post = grad_y[:, n:2 * n] + gy_res = grad_y[:, 2 * n:] + + gh_pre = gy_pre * y_pre * (1.0 - y_pre) + gh_post = gy_post * (y_post / 2.0) * (1.0 - y_post / 2.0) * 2.0 + gh_res = gy_res + + grad_h = torch.cat([gh_pre, gh_post, gh_res], dim=-1) + + # 2. Linear backward: h = proj * alpha / (r + eps) + bias + r_eps = r + eps + alpha_row = alpha.unsqueeze(0) + + grad_proj = grad_h * alpha_row / r_eps + grad_r_from_h = -(grad_h * proj * alpha_row / (r_eps * r_eps)).sum(dim=-1, keepdim=True) + + # Per-column grad_alpha, then reduce to 3 scalars + ga_all = (grad_h * proj / r_eps).sum(dim=0) # [N] + grad_alpha_pre = ga_all[:n].sum() + grad_alpha_post = ga_all[n:2 * n].sum() + grad_alpha_res = ga_all[2 * n:].sum() + grad_bias = grad_h.sum(dim=0) + + # 3. Total grad_r = from compute_h + external + grad_r_total = grad_r_from_h + if grad_r_ext is not None: + grad_r_total = grad_r_total + grad_r_ext + + # 4. Convert grad_r from r_new=norm/sqrt(K) to r_old=1/(r_new+eps) + grad_r_for_bwd = grad_r_total * (-r_eps * r_eps) + + return grad_proj, grad_r_for_bwd, grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias + + +def manual_backward_proj_rms(grad_proj, grad_r_for_bwd, x, weight, norm, eps): + """Hand-derived backward through proj_rms (reference, no kernel). + + grad_r_for_bwd is w.r.t. r_old = 1/(norm/sqrt(K) + eps). + """ + K = x.shape[-1] + inv_sqrt_k = 1.0 / math.sqrt(K) + + # grad through proj = x @ weight.T + grad_x_from_proj = grad_proj @ weight # [M, K] + grad_weight = grad_proj.t() @ x # [N, K] + + # grad through r_old = 1/(norm/sqrt(K) + eps) + # d(r_old)/d(norm) = -1/(norm/sqrt(K)+eps)^2 * (1/sqrt(K)) + # d(norm)/d(x_i) = x_i / norm + u = norm * inv_sqrt_k + eps + coeff = -1.0 / (u * u) * inv_sqrt_k # [M, 1] + inv_norm = torch.where(norm > 0, 1.0 / norm, torch.zeros_like(norm)) + grad_x_from_r = grad_r_for_bwd * coeff * x * inv_norm # [M, K] + + grad_x = grad_x_from_proj + grad_x_from_r + return grad_x, grad_weight + + +def test_backward(M, n, K, eps=1e-6): + """Compare manual backward against autograd.""" + N = n * n + 2 * n + + x = torch.randn(M, K, dtype=DTYPE, device=DEVICE, requires_grad=True) + weight = torch.randn(N, K, dtype=DTYPE, device=DEVICE, requires_grad=True) + alpha_pre = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) + alpha_post = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) + alpha_res = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) + bias = torch.randn(N, dtype=DTYPE, device=DEVICE, requires_grad=True) + + grad_y = torch.randn(M, N, dtype=DTYPE, device=DEVICE) + grad_r_ext = torch.randn(M, 1, dtype=DTYPE, device=DEVICE) + + # ---- Autograd path ---- + y_auto, r_auto, _, _ = forward_proj_rms_compute_h( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, + ) + loss = (y_auto * grad_y).sum() + (r_auto * grad_r_ext).sum() + loss.backward() + + auto_grad_x = x.grad.clone() + auto_grad_w = weight.grad.clone() + auto_grad_ap = alpha_pre.grad.clone() + auto_grad_apo = alpha_post.grad.clone() + auto_grad_ar = alpha_res.grad.clone() + auto_grad_b = bias.grad.clone() + + # ---- Manual path ---- + x2 = x.detach().clone() + w2 = weight.detach().clone() + ap2 = alpha_pre.detach().clone() + apo2 = alpha_post.detach().clone() + ar2 = alpha_res.detach().clone() + b2 = bias.detach().clone() + + y_man, r_man, proj_man, norm_man = forward_proj_rms_compute_h( + x2, w2, ap2, apo2, ar2, b2, n, eps, + ) + + # Step 1: manual compute_h backward + ( + grad_proj, grad_r_for_bwd, + grad_alpha_pre, grad_alpha_post, grad_alpha_res, + grad_bias, + ) = manual_backward_compute_h( + grad_y, grad_r_ext, y_man, proj_man, r_man, ap2, apo2, ar2, n, eps, + ) + + # Step 2: manual proj_rms backward + grad_x_man, grad_w_man = manual_backward_proj_rms( + grad_proj, grad_r_for_bwd, x2, w2, norm_man, eps, + ) + + # ---- Compare ---- + def check(name, manual, auto, atol=1e-10, rtol=1e-7): + diff = (manual - auto).abs().max().item() + rel = diff / (auto.abs().max().item() + 1e-15) + ok = diff < atol or rel < rtol + status = "PASS" if ok else "FAIL" + print(f" {status}: {name:14s} max_abs_diff={diff:.3e} rel={rel:.3e}") + return ok + + print(f"\n--- M={M}, n={n}, K={K}, N={N} ---") + all_ok = True + all_ok &= check("grad_x", grad_x_man, auto_grad_x) + all_ok &= check("grad_weight", grad_w_man, auto_grad_w) + all_ok &= check("grad_alpha_pre", grad_alpha_pre, auto_grad_ap) + all_ok &= check("grad_alpha_post", grad_alpha_post, auto_grad_apo) + all_ok &= check("grad_alpha_res", grad_alpha_res, auto_grad_ar) + all_ok &= check("grad_bias", grad_bias, auto_grad_b) + return all_ok + + +if __name__ == "__main__": + print("Verifying manual backward formulas against autograd (float64)...") + results = [] + for M, n, K in [(32, 2, 64), (64, 4, 256), (16, 3, 128), (128, 4, 512)]: + results.append(test_backward(M, n, K)) + + print("\n" + "=" * 50) + if all(results): + print("ALL PASSED") + else: + print("SOME FAILED") diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index cf44f2d7cd0..c5f6ced37bb 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -72,7 +72,9 @@ def test_forward_normal_vs_checkpoint_correctness(self): # Forward without checkpoint (reference) torch.manual_seed(42) torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref = module._forward_normal(hidden_states) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal( + hidden_states + ) mixed_ref = module.apply_h_res(h_res_ref, residual) loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() loss_ref.backward() @@ -83,8 +85,8 @@ def test_forward_normal_vs_checkpoint_correctness(self): torch.manual_seed(42) torch.cuda.manual_seed(42) manager = CheckpointManager() - aggregated_ckpt, h_res_ckpt, h_post_ckpt = module._forward_with_checkpoint( - hidden_states_ckpt, manager + aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt_out = ( + module._forward_with_checkpoint(hidden_states_ckpt, manager) ) mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) # Calculate loss before discarding outputs @@ -181,7 +183,7 @@ def test_forward_with_manager_parameter(self): # Reference: forward without manager (uses _forward_normal) torch.manual_seed(42) torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref = module.forward( + aggregated_ref, h_res_ref, h_post_ref, _ = module.forward( hidden_states, mhc_recompute_manager=None ) loss_ref = aggregated_ref.sum() + h_res_ref.sum() + h_post_ref.sum() @@ -192,7 +194,7 @@ def test_forward_with_manager_parameter(self): torch.manual_seed(42) torch.cuda.manual_seed(42) manager = CheckpointManager() - aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + aggregated_ckpt, h_res_ckpt, h_post_ckpt, _ = module.forward( hidden_states_ckpt, mhc_recompute_manager=manager ) loss_ckpt = aggregated_ckpt.sum() + h_res_ckpt.sum() + h_post_ckpt.sum() @@ -260,7 +262,7 @@ def test_multiple_hyper_connections_in_chain(self): h = hidden_states_ref r = residual_ref for module in modules: - agg, h_res, h_post = module.forward(h, mhc_recompute_manager=None) + agg, h_res, h_post, _ = module.forward(h, mhc_recompute_manager=None) agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=None) mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] h = agg + mixed @@ -280,7 +282,7 @@ def test_multiple_hyper_connections_in_chain(self): h = hidden_states_ckpt r = residual_ckpt for module in modules: - agg, h_res, h_post = module.forward(h, mhc_recompute_manager=manager) + agg, h_res, h_post, _ = module.forward(h, mhc_recompute_manager=manager) agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=manager) mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] h = agg + mixed @@ -338,7 +340,7 @@ def test_partial_checkpoint_last_layer_not_checkpointed(self): # Reference torch.manual_seed(42) torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref = module.forward( + aggregated_ref, h_res_ref, h_post_ref, _ = module.forward( hidden_states_ref, mhc_recompute_manager=None ) aggregated_ref, _ = module.apply_h_post( @@ -357,7 +359,7 @@ def test_partial_checkpoint_last_layer_not_checkpointed(self): torch.manual_seed(42) torch.cuda.manual_seed(42) manager = CheckpointManager() - aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + aggregated_ckpt, h_res_ckpt, h_post_ckpt, _ = module.forward( hidden_states_ckpt, mhc_recompute_manager=manager ) @@ -404,5 +406,438 @@ def test_config_enable_mhc_recompute(self): assert config.enable_hyper_connections is True +class TestHyperConnectionBroadcastCheckpoint: + """Test HyperConnectionModule checkpoint with BroadcastTensorFused. + + These tests mirror TestHyperConnectionCheckpoint but use the returned + residual (BroadcastTensor view) from forward() for apply_h_res, exercising + all 3 gradient paths through a single hidden_states tensor. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4): + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_residual_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1) + module.cuda() + return module + + def test_broadcast_normal_vs_checkpoint_correctness(self): + """ + _forward_normal vs _forward_with_checkpoint using BroadcastTensor residual. + + All 3 gradient paths (compute_mappings, aggregate, apply_h_res) flow + through a single hidden_states tensor via BroadcastTensorFused. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference: normal forward, use returned residual + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal( + hidden_states + ) + mixed_ref = module.apply_h_res(h_res_ref, residual_ref) + loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_ref = hidden_states.grad.clone() + + # Checkpoint: use returned residual + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt = ( + module._forward_with_checkpoint(hidden_states_ckpt, manager) + ) + mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) + loss_ckpt = aggregated_ckpt.sum() + mixed_ckpt.sum() + h_post_ckpt.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_ckpt = hidden_states_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( + f"Broadcast checkpoint gradients mismatch" + ) + + def test_broadcast_forward_with_manager(self): + """forward() with/without manager using BroadcastTensor residual.""" + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + agg_ref, h_res_ref, h_post_ref, res_ref = module.forward( + hidden_states, mhc_recompute_manager=None + ) + loss_ref = agg_ref.sum() + module.apply_h_res(h_res_ref, res_ref).sum() + h_post_ref.sum() + loss_ref.backward() + grad_ref = hidden_states.grad.clone() + + # With manager + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + agg_ckpt, h_res_ckpt, h_post_ckpt, res_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + loss_ckpt = ( + agg_ckpt.sum() + module.apply_h_res(h_res_ckpt, res_ckpt).sum() + h_post_ckpt.sum() + ) + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_ckpt = hidden_states_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5) + + def test_broadcast_fused_h_res_h_post_bda(self): + """ + Full TransformerLayer-like flow: forward() -> simulated layer -> + fused_h_res_h_post_bda() using BroadcastTensor residual. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module.forward( + hidden_states, mhc_recompute_manager=None + ) + layer_output = aggregated_ref * 0.1 # simulated layer + output_ref = module.fused_h_res_h_post_bda( + h_res_ref, residual_ref, h_post_ref, + (layer_output, None), 0.0, True, False, + ) + loss_ref = output_ref.sum() + loss_ref.backward() + grad_ref = hidden_states.grad.clone() + + # With checkpoint manager + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + layer_output_ckpt = aggregated_ckpt * 0.1 + output_ckpt = module.fused_h_res_h_post_bda( + h_res_ckpt, residual_ckpt, h_post_ckpt, + (layer_output_ckpt, None), 0.0, True, False, manager, + ) + loss_ckpt = output_ckpt.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_ckpt = hidden_states_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( + f"Broadcast fused_h_res_h_post_bda gradients mismatch" + ) + + +class TestMHCBroadcastChainIntegration: + """Test multiple chained HyperConnectionModules with BroadcastTensor residual.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_chained_broadcast_fwd_bwd(self): + """ + Chain 3 HyperConnectionModules using BroadcastTensor residual throughout. + Compare ref (no ckpt) vs checkpoint path. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + config = TransformerConfig( + num_layers=4, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + modules = [ + HyperConnectionModule(config=config, layer_number=i + 1).cuda() for i in range(3) + ] + + hidden_states_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + + # Reference: no checkpoint, use BroadcastTensor residual + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + h = hidden_states_ref + for module in modules: + agg, h_res, h_post, res = module.forward(h, mhc_recompute_manager=None) + layer_out = agg * 0.1 + h = module.fused_h_res_h_post_bda( + h_res, res, h_post, (layer_out, None), 0.0, True, False, + ) + + loss_ref = h.sum() + loss_ref.backward() + grad_ref = hidden_states_ref.grad.clone() + + # With checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointManager() + h = hidden_states_ckpt + for module in modules: + agg, h_res, h_post, res = module.forward(h, mhc_recompute_manager=manager) + layer_out = agg * 0.1 + h = module.fused_h_res_h_post_bda( + h_res, res, h_post, (layer_out, None), 0.0, True, False, manager, + ) + + loss_ckpt = h.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_ckpt = hidden_states_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-4), ( + f"Chained broadcast HyperConnection gradients mismatch" + ) + + def test_partial_checkpoint_broadcast(self): + """ + Last layer not checkpointed, all using BroadcastTensor residual. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1).cuda() + + hidden_states_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + agg_ref, h_res_ref, h_post_ref, res_ref = module.forward( + hidden_states_ref, mhc_recompute_manager=None + ) + layer_out_ref = agg_ref * 0.1 + output_ref = module.fused_h_res_h_post_bda( + h_res_ref, res_ref, h_post_ref, (layer_out_ref, None), 0.0, True, False, + ) + loss_ref = output_ref.sum() + loss_ref.backward() + grad_ref = hidden_states_ref.grad.clone() + + # With manager, simulate last layer not checkpointed + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + agg_ckpt, h_res_ckpt, h_post_ckpt, res_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + layer_out_ckpt = agg_ckpt * 0.1 + # Last layer: no manager for BDA + output_ckpt = module.fused_h_res_h_post_bda( + h_res_ckpt, res_ckpt, h_post_ckpt, (layer_out_ckpt, None), 0.0, True, False, + ) + manager.discard_all_outputs_and_register_unified_recompute(output_ckpt) + loss_ckpt = output_ckpt.sum() + loss_ckpt.backward() + grad_ckpt = hidden_states_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( + f"Partial checkpoint broadcast gradients mismatch" + ) + + +class TestBroadcastTensorFused: + """Test BroadcastTensorFused autograd.Function and native_fused_add_3.""" + + def test_native_fused_add_3_correctness(self): + """Test native_fused_add_3 matches plain addition.""" + from megatron.core.transformer.hyper_connection import native_fused_add_3 + + a = torch.randn(8, 16, device='cuda') + b = torch.randn(8, 16, device='cuda') + c = torch.randn(8, 16, device='cuda') + result = native_fused_add_3(a, b, c) + expected = a + b + c + assert torch.allclose(result, expected, atol=1e-6) + + def test_native_fused_add_3_gradient(self): + """Test native_fused_add_3 produces correct gradients.""" + from megatron.core.transformer.hyper_connection import native_fused_add_3 + + a = torch.randn(4, 8, device='cuda', requires_grad=True) + b = torch.randn(4, 8, device='cuda', requires_grad=True) + c = torch.randn(4, 8, device='cuda', requires_grad=True) + result = native_fused_add_3(a, b, c) + result.sum().backward() + assert torch.allclose(a.grad, torch.ones_like(a)) + assert torch.allclose(b.grad, torch.ones_like(b)) + assert torch.allclose(c.grad, torch.ones_like(c)) + + def test_broadcast_tensor_returns_views(self): + """Test BroadcastTensorFused forward returns 3 views sharing storage.""" + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + x = torch.randn(4, 8, device='cuda', requires_grad=True) + a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) + assert a.data_ptr() == x.data_ptr() + assert b.data_ptr() == x.data_ptr() + assert c.data_ptr() == x.data_ptr() + + def test_broadcast_tensor_gradient_correctness(self): + """Test BroadcastTensorFused backward sums gradients correctly.""" + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + x = torch.randn(4, 8, device='cuda', requires_grad=True) + a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) + # Each branch applies a different scalar: loss = a.sum() + 2*b.sum() + 3*c.sum() + loss = a.sum() + 2 * b.sum() + 3 * c.sum() + loss.backward() + # x.grad should be (1 + 2 + 3) * ones + expected = torch.full_like(x, 6.0) + assert torch.allclose(x.grad, expected, atol=1e-6), ( + f"Expected grad {expected}, got {x.grad}" + ) + + def test_broadcast_tensor_partial_none_grads(self): + """Test BroadcastTensorFused handles None gradients from unused branches.""" + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + x = torch.randn(4, 8, device='cuda', requires_grad=True) + a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) + # Only use one branch + loss = a.sum() + loss.backward() + expected = torch.ones_like(x) + assert torch.allclose(x.grad, expected, atol=1e-6) + + def test_broadcast_tensor_end_to_end_with_hyper_connection(self): + """ + End-to-end test: HyperConnectionModule.forward returns 4-tuple and + gradients through all 3 paths (compute_mappings, aggregate, + fused_h_res_h_post_bda) are correct. + """ + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + try: + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1).cuda() + + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, + device='cuda', requires_grad=True, + ) + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Forward returns 4-tuple now + aggregated, h_res, h_post, residual = module.forward(hidden_states) + + # Simulate full flow: aggregate + apply_h_res on residual + mixed = module.apply_h_res(h_res, residual) + loss = aggregated.sum() + mixed.sum() + h_post.sum() + loss.backward() + + # Verify gradient exists and is finite + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() + finally: + Utils.destroy_model_parallel() + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From a93892a0976254c1507ed9714832839b08d12a38 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 30 Apr 2026 14:00:15 +0800 Subject: [PATCH 02/34] feat: align mHC fused kernel backends --- megatron/core/fusions/fused_mhc_kernels.py | 73 +++++++- megatron/core/fusions/triton_mhc_kernels.py | 160 +++++++++++++++++- .../fusions/test_fused_mhc_kernels.py | 41 ++++- 3 files changed, 262 insertions(+), 12 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index c18e38e8a6a..cb56cf6d733 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1848,9 +1848,10 @@ def _cutile_fused_compute_h_proj_rms_bwd( # The public fused API chooses the fastest validated backend per operation: # # sinkhorn fwd/bwd: Triton -> cuTile -> torch -# h_post_bda fwd: cuTile -> torch -# h_post_bda bwd: Triton -> cuTile -> torch -# h_aggregate/proj_rms: cuTile -> torch +# h_post_bda fwd/bwd: Triton -> cuTile -> torch +# h_aggregate fwd: Triton -> cuTile -> torch +# h_aggregate bwd: cuTile -> torch +# proj_rms/proj_rms_compute_h: cuTile -> torch # # Runtime CUDA launch failures are intentionally not swallowed; after such an # error the CUDA context may not be safely reusable for fallback work. @@ -1876,6 +1877,28 @@ def _get_triton_sinkhorn(): return None +def _get_triton_h_aggregate_fwd(): + if not _TRITON_AVAILABLE: + return None + try: + from megatron.core.fusions.triton_mhc_kernels import _triton_h_aggregate_fwd + + return _triton_h_aggregate_fwd + except ImportError: + return None + + +def _get_triton_h_post_bda_fwd(): + if not _TRITON_AVAILABLE: + return None + try: + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_fwd + + return _triton_h_post_bda_fwd + except ImportError: + return None + + def _get_triton_h_post_bda_bwd(): if not _TRITON_AVAILABLE: return None @@ -1887,6 +1910,14 @@ def _get_triton_h_post_bda_bwd(): return None +def _torch_h_aggregate_bwd( + grad_output: Tensor, x: Tensor, h_pre: Tensor +) -> Tuple[Tensor, Tensor]: + grad_x = grad_output.unsqueeze(2) * h_pre.unsqueeze(-1) + grad_h = torch.sum(grad_output.unsqueeze(2) * x, dim=-1) + return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) + + def _torch_h_post_bda_bwd( grad_output: Tensor, h_res: Tensor, @@ -2038,8 +2069,31 @@ def backward(ctx, grad_y, grad_r_ext): return grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None +class FusedHAggregate(torch.autograd.Function): + """H_aggregate with Triton/cuTile/torch forward and cuTile/torch backward.""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + triton_fwd = _get_triton_h_aggregate_fwd() + if triton_fwd is not None: + output = triton_fwd(x, h_pre) + elif _CUTILE_AVAILABLE: + output = _cutile_h_aggregate_fwd(x, h_pre) + else: + output = native_h_aggregate(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + x, h_pre = ctx.saved_tensors + if _CUTILE_AVAILABLE: + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + return _torch_h_aggregate_bwd(grad_output, x, h_pre) + + class FusedHPostBDA(torch.autograd.Function): - """H_post_bda with cuTile/torch forward and Triton/cuTile/torch backward.""" + """H_post_bda with Triton/cuTile/torch forward and backward.""" @staticmethod def forward( @@ -2050,7 +2104,10 @@ def forward( x: Tensor, bias: Optional[Tensor], ): - if _CUTILE_AVAILABLE: + triton_fwd = _get_triton_h_post_bda_fwd() + if triton_fwd is not None: + output = triton_fwd(h_res, original_residual, h_post, x, bias) + elif _CUTILE_AVAILABLE: output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) else: output = native_h_post_bda(h_res, original_residual, h_post, x, bias) @@ -2089,9 +2146,9 @@ def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: - """Weighted n-stream to 1-stream aggregation using cuTile, then torch.""" - if _CUTILE_AVAILABLE: - return CutileHAggregate.apply(x, h_pre) + """Weighted n-stream to 1-stream aggregation using Triton/cuTile/torch.""" + if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: + return FusedHAggregate.apply(x, h_pre) return native_h_aggregate(x, h_pre) diff --git a/megatron/core/fusions/triton_mhc_kernels.py b/megatron/core/fusions/triton_mhc_kernels.py index d21ea7905d0..26ee56f5e1c 100644 --- a/megatron/core/fusions/triton_mhc_kernels.py +++ b/megatron/core/fusions/triton_mhc_kernels.py @@ -5,7 +5,8 @@ Only the Triton kernels that are faster than the cuTile implementations are kept here: - Sinkhorn forward/backward - - H_post_bda backward + - H_aggregate forward + - H_post_bda forward/backward The public mHC API lives in ``fused_mhc_kernels.py`` and handles fallback. """ @@ -185,10 +186,163 @@ def triton_fused_sinkhorn( # ============================================================================ -# H_post BDA backward +# H_aggregate forward # ============================================================================ +@triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], +) +@triton.jit +def _triton_h_agg_fwd_kernel( + x_ptr, h_ptr, out_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_x_s, stride_x_n, stride_x_c, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, +): + """out[s, c] = sum_i x[s, i, c] * h[s, i].""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for i in tl.static_range(N): + x_i = tl.load( + x_ptr + offs_s[:, None] * stride_x_s + i * stride_x_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + h_i = tl.load(h_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + acc += h_i[:, None] * x_i + tl.store( + out_ptr + offs_s[:, None] * C + offs_c[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + +def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.contiguous().view(sb, n, C) + h_flat = h_pre.contiguous().view(sb, n) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_h_agg_fwd_kernel[grid]( + x_flat, h_flat, out, sb, C, n, + x_flat.stride(0), x_flat.stride(1), x_flat.stride(2), + ) + return out.view(s, b, C) + + +# ============================================================================ +# H_post BDA +# ============================================================================ + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], +) +@triton.jit +def _triton_hpb_fwd_kernel( + hr_ptr, orig_ptr, hp_ptr, x_ptr, bias_ptr, out_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_hr_s, stride_hr_i, stride_hr_j, + stride_orig_s, stride_orig_n, stride_orig_c, + stride_out_s, stride_out_n, stride_out_c, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, +): + """out = hr @ orig + hp * (x + bias).""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + hp_i = tl.load(hp_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + out_i = hp_i[:, None] * x_tile + + for j in tl.static_range(N): + hr_ij = tl.load( + hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + out_i += hr_ij[:, None] * orig_j + + tl.store( + out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], + out_i.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + +def _triton_h_post_bda_fwd( + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], +) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=dev) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_fwd_kernel[grid]( + hr_flat, orig_flat, hp_flat, x_flat, bias if bias is not None else x_flat, out, + sb, C, n, + hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), + orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), + out.stride(0), out.stride(1), out.stride(2), + HAS_BIAS=(bias is not None), + ) + return out.view(s, b, n, C) + + @triton.autotune( configs=[ triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) @@ -390,6 +544,8 @@ def _triton_h_post_bda_bwd( __all__ = [ + "_triton_h_aggregate_fwd", + "_triton_h_post_bda_fwd", "_triton_h_post_bda_bwd", "triton_fused_sinkhorn", ] diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 9e5d65b41fc..da50649d3e4 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -249,6 +249,24 @@ def test_fwd_bwd_vs_reference(self, s, b, n, C): torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) +class TestTritonHAggregate: + """Tests for Triton h_aggregate forward against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C): + from megatron.core.fusions.triton_mhc_kernels import _triton_h_aggregate_fwd + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + + out_t = _triton_h_aggregate_fwd(x_data, h_data) + out_r = _ref_h_aggregate(x_data, h_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + # ============================================================================ # H_post BDA # ============================================================================ @@ -351,8 +369,27 @@ def _make_inputs(): ) -class TestTritonHPostBDABwd: - """Tests for Triton h_post_bda backward kernels against PyTorch reference.""" +class TestTritonHPostBDA: + """Tests for Triton h_post_bda kernels against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb forward output must match the PyTorch reference.""" + from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_fwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + + out_t = _triton_h_post_bda_fwd(hr_data, orig_data, hp_data, x_data, bias_data) + out_r = _ref_h_post_bda(hr_data, orig_data, hp_data, x_data, bias_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) @_require_triton @pytest.mark.parametrize("with_bias", [True, False]) From e6baf1441217727c767a161bf44a4caa5a891161 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 30 Apr 2026 14:49:26 +0800 Subject: [PATCH 03/34] fix: make Triton sinkhorn constant compile-safe --- megatron/core/fusions/triton_mhc_kernels.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/fusions/triton_mhc_kernels.py b/megatron/core/fusions/triton_mhc_kernels.py index 26ee56f5e1c..798d880bc5d 100644 --- a/megatron/core/fusions/triton_mhc_kernels.py +++ b/megatron/core/fusions/triton_mhc_kernels.py @@ -18,8 +18,6 @@ import triton.language as tl from torch import Tensor -LOG2E = 1.4426950408889634 - # ============================================================================ # Sinkhorn-Knopp @@ -48,7 +46,7 @@ def _triton_sinkhorn_fwd_kernel( logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) row_max = tl.max(logits, axis=1) - M = tl.exp2((logits - row_max[:, None]) * LOG2E) + M = tl.exp2((logits - row_max[:, None]) * 1.4426950408889634) tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) for _ in range(NUM_ITERS): From 2f9d80f2c45d8f42e6c02d24b5a4b8935f2317c0 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 30 Apr 2026 18:35:40 +0800 Subject: [PATCH 04/34] perf: remove torch zero fills from mHC backward --- megatron/core/fusions/fused_mhc_kernels.py | 159 ++++++++++++++++----- 1 file changed, 120 insertions(+), 39 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index cb56cf6d733..e0e7314033c 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1304,6 +1304,7 @@ def _ct_fused_grad_h_proj_kernel( eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, + HAS_GRAD_R_EXT: ConstInt, ): """Precompute grad_h, grad_proj, and grad_r_total for downstream backward kernels. @@ -1365,11 +1366,14 @@ def _ct_fused_grad_h_proj_kernel( grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True, ) - grad_r_ext_tile = ct.load( - GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, - ) - grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + if HAS_GRAD_R_EXT: + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + else: + grad_r_ext_tile = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) grad_r_total = grad_r_from_h + grad_r_ext_tile ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h.astype(GRAD_H.dtype)) @@ -1448,14 +1452,14 @@ def _ct_fused_grad_x_weight_kernel( ct.store(GRAD_WEIGHT, index=(0, tile_k_id), tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype)) @ct.kernel - def _ct_scalar_grads_kernel( + def _ct_scalar_grads_partials_kernel( GRAD_H, # [M, TILE_SIZE_N] precomputed PROJ, # [M, N] R, # [M, 1] - GRAD_ALPHA_PRE, # [1] output (atomic) - GRAD_ALPHA_POST, # [1] output (atomic) - GRAD_ALPHA_RES, # [1] output (atomic) - GRAD_BIAS, # [TILE_SIZE_N] output (atomic) + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] output + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] output M: int, N: int, n: int, @@ -1463,7 +1467,7 @@ def _ct_scalar_grads_kernel( TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, ): - """Compute scalar gradients (grad_alpha_pre/post/res, grad_bias). + """Compute per-M-tile scalar-gradient partials. Grid: (ceil(M / TILE_SIZE_M),). Each block processes one M-tile. """ @@ -1499,11 +1503,68 @@ def _ct_scalar_grads_kernel( inv_r_eps = 1.0 / r_eps ga_all = grad_h * proj_tile * inv_r_eps - ct.atomic_add(GRAD_ALPHA_PRE, 0, ct.sum(ga_all * mask_pre_2d).astype(GRAD_ALPHA_PRE.dtype)) - ct.atomic_add(GRAD_ALPHA_POST, 0, ct.sum(ga_all * mask_post_2d).astype(GRAD_ALPHA_POST.dtype)) - ct.atomic_add(GRAD_ALPHA_RES, 0, ct.sum(ga_all * mask_res_2d).astype(GRAD_ALPHA_RES.dtype)) + ga_pre = ct.reshape(ct.sum(ga_all * mask_pre_2d), (1, 1)) + ga_post = ct.reshape(ct.sum(ga_all * mask_post_2d), (1, 1)) + ga_res = ct.reshape(ct.sum(ga_all * mask_res_2d), (1, 1)) partial_gb = ct.sum(grad_h, axis=0, keepdims=False) - ct.atomic_add(GRAD_BIAS, offsets, partial_gb.astype(GRAD_BIAS.dtype)) + ct.store( + GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), + tile=ga_pre.astype(GRAD_ALPHA_PRE_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), + tile=ga_post.astype(GRAD_ALPHA_POST_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), + tile=ga_res.astype(GRAD_ALPHA_RES_PARTIALS.dtype), + ) + ct.store( + GRAD_BIAS_PARTIALS, index=(bid_m, 0), + tile=ct.reshape(partial_gb, (1, TILE_SIZE_N)).astype(GRAD_BIAS_PARTIALS.dtype), + ) + + @ct.kernel + def _ct_scalar_grads_reduce_kernel( + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] + GRAD_ALPHA_PRE, # [1, 1] output + GRAD_ALPHA_POST, # [1, 1] output + GRAD_ALPHA_RES, # [1, 1] output + GRAD_BIAS, # [1, TILE_SIZE_N] output + NUM_M_BLOCKS: int, + TILE_SIZE_N: ConstInt, + ): + """Reduce scalar-gradient partials and write final dtype outputs.""" + acc_pre = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_post = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_res = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_bias = ct.full((1, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for bid_m in range(NUM_M_BLOCKS): + acc_pre += ct.load( + GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), shape=(1, 1), + padding_mode=PAD_ZERO, + ).astype(ct.float32) + acc_post += ct.load( + GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), shape=(1, 1), + padding_mode=PAD_ZERO, + ).astype(ct.float32) + acc_res += ct.load( + GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), shape=(1, 1), + padding_mode=PAD_ZERO, + ).astype(ct.float32) + acc_bias += ct.load( + GRAD_BIAS_PARTIALS, index=(bid_m, 0), shape=(1, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ).astype(ct.float32) + + ct.store(GRAD_ALPHA_PRE, index=(0, 0), tile=acc_pre.astype(GRAD_ALPHA_PRE.dtype)) + ct.store(GRAD_ALPHA_POST, index=(0, 0), tile=acc_post.astype(GRAD_ALPHA_POST.dtype)) + ct.store(GRAD_ALPHA_RES, index=(0, 0), tile=acc_res.astype(GRAD_ALPHA_RES.dtype)) + ct.store(GRAD_BIAS, index=(0, 0), tile=acc_bias.astype(GRAD_BIAS.dtype)) @ct.kernel def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( @@ -1525,13 +1586,14 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( n: int, eps: float, TILE_N_SIZE: ConstInt, + HAS_GRAD_R_EXT: ConstInt, ): """Fused backward (small K path) with work-stealing. Grid: (num_sms, 2). bid(1)==0: grad_weight via work-stealing over K-tiles, loops M. bid(1)==1: grad_x via work-stealing over (M×K) tiles. - Scalar gradients computed by separate _ct_scalar_grads_kernel. + Scalar gradients are computed by the separate partial/reduce kernels. """ zero_pad = ct.PaddingMode.ZERO @@ -1652,11 +1714,14 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True, ) - grad_r_ext_tile = ct.load( - GRAD_R_EXT, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, - ) - grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + if HAS_GRAD_R_EXT: + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + else: + grad_r_ext_tile = ct.full((TILE_DA_SIZE_M, 1), 0.0, dtype=ct.float32) grad_r_total = grad_r_from_h + grad_r_ext_tile x_tile = ct.load( @@ -1722,13 +1787,9 @@ def _cutile_fused_compute_h_proj_rms_bwd( grad_x = torch.empty_like(x) grad_weight = torch.empty_like(weight) - grad_alpha_pre_acc = torch.zeros(1, dtype=torch.float32, device=dev) - grad_alpha_post_acc = torch.zeros(1, dtype=torch.float32, device=dev) - grad_alpha_res_acc = torch.zeros(1, dtype=torch.float32, device=dev) - grad_bias_acc = torch.zeros(TILE_N, dtype=torch.float32, device=dev) - - if grad_r_ext is None: - grad_r_ext = torch.zeros(M, 1, dtype=x.dtype, device=dev) + has_grad_r_ext = grad_r_ext is not None + has_grad_r_ext_flag = int(has_grad_r_ext) + grad_r_ext_arg = grad_r_ext if has_grad_r_ext else r # 0. Precompute grad_h, grad_proj, grad_r_total grad_h_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) @@ -1741,11 +1802,11 @@ def _cutile_fused_compute_h_proj_rms_bwd( (math.ceil(M / tile_m_precomp),), _ct_fused_grad_h_proj_kernel, ( - grad_y, y_activated, proj, r, grad_r_ext, + grad_y, y_activated, proj, r, grad_r_ext_arg, alpha_pre, alpha_post, alpha_res, grad_h_buf, grad_proj_buf, grad_r_total_buf, M, N, n, eps, - tile_m_precomp, TILE_N, + tile_m_precomp, TILE_N, has_grad_r_ext_flag, ), ) @@ -1808,37 +1869,57 @@ def _cutile_fused_compute_h_proj_rms_bwd( (num_sms, 2, 1), _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, ( - x, weight, grad_y, y_activated, proj, r, grad_r_ext, + x, weight, grad_y, y_activated, proj, r, grad_r_ext_arg, alpha_pre, alpha_post, alpha_res, grad_x, grad_weight, M, N, K, n, eps, - TILE_N, + TILE_N, has_grad_r_ext_flag, ), ) # 2. Separate lightweight kernel for scalar gradients (grad_alpha, grad_bias) tile_m_scalar = min(128, M) num_m_blocks = math.ceil(M / tile_m_scalar) + grad_alpha_pre_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_post_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_res_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_bias_partials = torch.empty(num_m_blocks, TILE_N, dtype=torch.float32, device=dev) + grad_alpha_pre = torch.empty(1, 1, dtype=alpha_pre.dtype, device=dev) + grad_alpha_post = torch.empty(1, 1, dtype=alpha_post.dtype, device=dev) + grad_alpha_res = torch.empty(1, 1, dtype=alpha_res.dtype, device=dev) + grad_bias = torch.empty(1, TILE_N, dtype=bias.dtype, device=dev) + ct.launch( stream, (num_m_blocks,), - _ct_scalar_grads_kernel, + _ct_scalar_grads_partials_kernel, ( grad_h_buf, proj, r, - grad_alpha_pre_acc, grad_alpha_post_acc, grad_alpha_res_acc, - grad_bias_acc, + grad_alpha_pre_partials, grad_alpha_post_partials, grad_alpha_res_partials, + grad_bias_partials, M, N, n, eps, tile_m_scalar, TILE_N, ), ) + ct.launch( + stream, + (1,), + _ct_scalar_grads_reduce_kernel, + ( + grad_alpha_pre_partials, grad_alpha_post_partials, grad_alpha_res_partials, + grad_bias_partials, + grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias, + num_m_blocks, TILE_N, + ), + ) return ( grad_x, grad_weight, - grad_alpha_pre_acc.to(alpha_pre.dtype), - grad_alpha_post_acc.to(alpha_post.dtype), - grad_alpha_res_acc.to(alpha_res.dtype), - grad_bias_acc[:N].to(bias.dtype), + grad_alpha_pre.view_as(alpha_pre), + grad_alpha_post.view_as(alpha_post), + grad_alpha_res.view_as(alpha_res), + grad_bias.view(-1)[:N], ) From 8b417f624f81cb7d040c19bf0068a071b7e0fb6d Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 30 Apr 2026 18:54:04 +0800 Subject: [PATCH 05/34] fix: avoid unmasked cuTile proj_rms stores --- megatron/core/fusions/fused_mhc_kernels.py | 189 +++++++++++++-------- 1 file changed, 118 insertions(+), 71 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index e0e7314033c..5b5e8b297be 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -949,7 +949,19 @@ def _proj_rms_fwd_autotune_configs(N): for tile_m in tile_ms: for tile_k in tile_ks: for split_k in split_ks: - yield {"TILE_M": tile_m, "TILE_N": TILE_N, "TILE_K": tile_k, 'SPLIT_K': split_k} + yield {"TILE_M": tile_m, "TILE_N": TILE_N, "TILE_K": tile_k, 'SPLIT_K': split_k} + + def _default_tile_m(M: int) -> int: + """Pick a tile size that avoids unmasked stores past the M dimension.""" + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m <= M and M % tile_m == 0: + return tile_m + return 1 + + def _default_proj_rms_fwd_config(M: int, K: int, TILE_N: int): + """Static fallback for skinny MHC projection when autotune cache is absent.""" + split_k = 16 if K >= 16384 else 8 if K >= 8192 else 1 + return _default_tile_m(M), TILE_N, min(128, K), split_k # Cache the best config across calls (keyed by M, N, K). _proj_rms_fwd_best_cfg: dict = {} @@ -971,7 +983,7 @@ def _cutile_proj_rms_fwd( if cached is not None: tm, tn, tk, split_k = cached else: - tm, tn, tk, split_k = 128, TILE_N, 128, 1 + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) # print(cached) # print('hahahahahah') @@ -995,43 +1007,56 @@ def _cutile_proj_rms_fwd( configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] # filter out configs with TILE_K > K or TILE_M > M - configs = [cfg for cfg in configs if cfg.TILE_K <= K and cfg.TILE_M <= M] - # if no valid config, set a default config with TILE_M = M and TILE_K = K and SPLIT_K = 1 - # if len(configs) == 0: - # configs = [SimpleNamespace(TILE_M=M, TILE_N=TILE_N, TILE_K=K, SPLIT_K=1)] - mx_split_k = max(cfg.SPLIT_K for cfg in configs) - proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) - norm = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) - r = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) - tuned = ct_experimental.autotune_launch( - stream, - grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), - kernel=_ct_proj_rms_fwd_kernel, - args_fn=lambda cfg: ( - x, weight, proj, norm, r, M, N, K, eps, - cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, - ), - search_space=configs, - ) - best = tuned.tuned_config - # print(best) - # print('hahahahahah') - # best.SPLIT_K = 1 - _proj_rms_fwd_best_cfg[cache_key] = (best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K) - proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) - norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) - r = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) - # Re-launch with best config for correct output. - ct.launch( - stream, - (math.ceil(M / best.TILE_M), best.SPLIT_K), - _ct_proj_rms_fwd_kernel, - (x, weight, proj, norm, r, M, N, K, eps, - best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K), - ) + configs = [cfg for cfg in configs if cfg.TILE_K <= K and M % cfg.TILE_M == 0] + if len(configs) == 0: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), + ) + proj = proj.view(split_k, M, N).sum(dim=0) + norm = norm.view(split_k, M, 1).sum(dim=0) + else: + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, weight, proj, norm, r, M, N, K, eps, + cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + # print(best) + # print('hahahahahah') + # best.SPLIT_K = 1 + _proj_rms_fwd_best_cfg[cache_key] = ( + best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + ) + proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) + norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + r = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + # Re-launch with best config for correct output. + ct.launch( + stream, + (math.ceil(M / best.TILE_M), best.SPLIT_K), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, + best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K), + ) - proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) - norm = norm.view(best.SPLIT_K, M, 1).sum(dim=0) + proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) + norm = norm.view(best.SPLIT_K, M, 1).sum(dim=0) norm = torch.sqrt(norm) # r = r.view(split_k, M, 1).sum(dim=0) r = 1.0 / (norm / math.sqrt(K) + eps) @@ -1041,8 +1066,8 @@ def _cutile_proj_rms_fwd( def _reduce_compute_h_autotune_configs(M): """Generate autotune search space for reduce_compute_h kernel.""" - for tile_m in (1, 2, 4, 8, 16, 32, 64, 128): - if tile_m <= M: + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m <= M and M % tile_m == 0: yield tile_m _reduce_compute_h_best_cfg: dict = {} @@ -1059,6 +1084,7 @@ def _cutile_reduce_compute_h( N: int, K: int, eps: float, + tile_m: int, tile_n: int, split_k: int, ) -> Tuple[Tensor, Tensor, Tensor]: @@ -1078,7 +1104,7 @@ def _cutile_reduce_compute_h( r_out = torch.empty(M, 1, dtype=proj_acc.dtype, device=dev) proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) - cache_key = (M, N, K, n, split_k) + cache_key = (M, N, K, n, tile_m, split_k) cached = _reduce_compute_h_best_cfg.get(cache_key) def _make_args(tm): @@ -1091,7 +1117,7 @@ def _make_args(tm): ) if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: - tm = cached if cached is not None else min(128, M) + tm = cached if cached is not None else tile_m ct.launch(stream, (math.ceil(M / tm),), _ct_reduce_compute_h_kernel, _make_args(tm)) else: from types import SimpleNamespace @@ -1099,6 +1125,7 @@ def _make_args(tm): configs = [ SimpleNamespace(TILE_M=tm) for tm in _reduce_compute_h_autotune_configs(M) + if tm == tile_m ] tuned = ct_experimental.autotune_launch( stream, @@ -1151,7 +1178,7 @@ def _cutile_proj_rms_compute_h_fwd( if cached is not None: tm, tn, tk, split_k = cached else: - tm, tn, tk, split_k = 128, TILE_N, 128, 1 + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) @@ -1167,41 +1194,61 @@ def _cutile_proj_rms_compute_h_fwd( from types import SimpleNamespace configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] - configs = [cfg for cfg in configs if cfg.TILE_K <= K and cfg.TILE_M <= M] - mx_split_k = max(cfg.SPLIT_K for cfg in configs) - proj_acc = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) - norm_acc = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) - r_placeholder = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) - tuned = ct_experimental.autotune_launch( - stream, - grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), - kernel=_ct_proj_rms_fwd_kernel, - args_fn=lambda cfg: ( - x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, - cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, - ), - search_space=configs, - ) - best = tuned.tuned_config - _proj_rms_fwd_best_cfg[cache_key] = (best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K) - tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + configs = [cfg for cfg in configs if cfg.TILE_K <= K and M % cfg.TILE_M == 0] + if len(configs) == 0: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, + tm, tn, tk, split_k, + ), + ) + else: + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj_acc = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, + cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_fwd_best_cfg[cache_key] = ( + best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + ) + tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K - proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) - norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) - r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) - ct.launch( - stream, - (math.ceil(M / tm), split_k), - _ct_proj_rms_fwd_kernel, - (x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, tm, tn, tk, split_k), - ) + proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, + tm, tn, tk, split_k, + ), + ) # Launch reduce + compute_h kernel y_activated, r, proj_reduced = _cutile_reduce_compute_h( proj_acc, norm_acc, bias, alpha_pre, alpha_post, alpha_res, n, M, N, K, eps, - TILE_N, split_k, + tm, TILE_N, split_k, ) return y_activated, r, proj_reduced From e712bacab63711871174f6ad4d711e52f542a979 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 5 May 2026 10:14:13 +0800 Subject: [PATCH 06/34] perf: split mHC compute_h outputs --- megatron/core/fusions/fused_mhc_kernels.py | 402 +++++++++--------- megatron/core/transformer/hyper_connection.py | 8 +- .../fusions/test_fused_mhc_kernels.py | 66 ++- 3 files changed, 252 insertions(+), 224 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 5b5e8b297be..52b519cab6b 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -716,13 +716,15 @@ def _ct_reduce_compute_h_kernel( Alpha_pre, Alpha_post, Alpha_res, - Y, + H_PRE, + H_POST, + H_RES, R, PROJ_OUT, M: int, N: int, K: int, - n: int, + n: ConstInt, eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, @@ -741,17 +743,23 @@ def _ct_reduce_compute_h_kernel( alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() - # 1. Reduce split-K partials - y_accum = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0, dtype=ct.float32) + # 1. Reduce split-K partials for each logical output segment. + pre_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + post_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) r_accum = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) for split_idx in ct.static_iter(range(SPLIT_K)): bid_m_k = bid_m + split_idx * num_bid_m - y_tile = ct.load( - Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + pre_tile = ct.load( + Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + post_tile = ct.load( + Y_acc, index=(bid_m_k, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) - y_accum = y_accum + ct.astype(y_tile, ct.float32) + pre_accum = pre_accum + ct.astype(pre_tile, ct.float32) + post_accum = post_accum + ct.astype(post_tile, ct.float32) r_tile = ct.load( R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), @@ -759,8 +767,9 @@ def _ct_reduce_compute_h_kernel( ) r_accum = r_accum + ct.astype(r_tile, ct.float32) - # Store reduced proj for backward - ct.store(PROJ_OUT, index=(bid_m, 0), tile=y_accum.astype(PROJ_OUT.dtype)) + # Store reduced projection segments for backward. + ct.store(PROJ_OUT, index=(bid_m, 0), tile=pre_accum.astype(PROJ_OUT.dtype)) + ct.store(PROJ_OUT, index=(bid_m, 1), tile=post_accum.astype(PROJ_OUT.dtype)) # 2. Compute r = norm / sqrt(K) denom = ct.full((TILE_SIZE_M, 1), K * 1.0, dtype=ct.float32) @@ -771,36 +780,42 @@ def _ct_reduce_compute_h_kernel( ct.store(R, index=(bid_m, 0), tile=r_val.astype(R.dtype)) - # 3. Build per-column scale from 3 scalar alphas via masks - offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) - - one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) - zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) - - mask_pre = ct.where(ct.less(offsets, n), one, zero) - mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre - mask_res = one - mask_pre - mask_post - - scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res - scale = ct.reshape(scale, (1, TILE_SIZE_N)) - - mask_pre = ct.reshape(mask_pre, (1, TILE_SIZE_N)) - mask_post = ct.reshape(mask_post, (1, TILE_SIZE_N)) - mask_res = ct.reshape(mask_res, (1, TILE_SIZE_N)) - - # 4. Apply compute_h: linear = proj * scale / (r + eps) + bias - bias_tile = ct.load( - Bias, index=(0, 0), shape=(1, TILE_SIZE_N), padding_mode=PAD_ZERO, + # 3. Apply compute_h directly into split outputs. + inv_r_eps = 1.0 / (r_val + eps) + bias_pre = ct.load( + Bias, index=(0, 0), shape=(1, n), padding_mode=PAD_ZERO, ) - bias_tile = ct.astype(bias_tile, ct.float32) - - linear = y_accum * scale / (r_val + eps) + bias_tile - - # 5. Mask-based activations - sig = _ct_sigmoid(linear) - out = sig * mask_pre + 2.0 * sig * mask_post + linear * mask_res + bias_post = ct.load( + Bias, index=(0, 1), shape=(1, n), padding_mode=PAD_ZERO, + ) + bias_pre = ct.astype(bias_pre, ct.float32) + bias_post = ct.astype(bias_post, ct.float32) + + h_pre_linear = pre_accum * alpha_pre * inv_r_eps + bias_pre + h_post_linear = post_accum * alpha_post * inv_r_eps + bias_post + h_pre = _ct_sigmoid(h_pre_linear) + h_post = _ct_sigmoid(h_post_linear) * 2.0 + + ct.store(H_PRE, index=(bid_m, 0), tile=h_pre.astype(H_PRE.dtype)) + ct.store(H_POST, index=(bid_m, 0), tile=h_post.astype(H_POST.dtype)) + + for res_chunk in ct.static_iter(range(n)): + res_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + res_tile = ct.load( + Y_acc, index=(bid_m_k, 2 + res_chunk), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + res_accum = res_accum + ct.astype(res_tile, ct.float32) - ct.store(Y, index=(bid_m, 0), tile=out.astype(Y.dtype)) + bias_res = ct.load( + Bias, index=(0, 2 + res_chunk), shape=(1, n), padding_mode=PAD_ZERO, + ) + bias_res = ct.astype(bias_res, ct.float32) + h_res = res_accum * alpha_res * inv_r_eps + bias_res + ct.store(PROJ_OUT, index=(bid_m, 2 + res_chunk), tile=res_accum.astype(PROJ_OUT.dtype)) + ct.store(H_RES, index=(bid_m, res_chunk), tile=h_res.astype(H_RES.dtype)) @ct.kernel def _ct_proj_rms_bwd_kernel( @@ -1087,11 +1102,13 @@ def _cutile_reduce_compute_h( tile_m: int, tile_n: int, split_k: int, - ) -> Tuple[Tensor, Tensor, Tensor]: + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Launch reduce split-K + compute_h kernel. Returns: - y_activated: [M, N] with h_pre | h_post | h_res activated + h_pre: [M, n] sigmoid-activated pre weights + h_post: [M, n] 2*sigmoid-activated post weights + h_res: [M, n*n] residual logits r: [M, 1] r = norm / sqrt(K) proj_reduced: [M, N] reduced projection (for backward) """ @@ -1100,7 +1117,9 @@ def _cutile_reduce_compute_h( bias_2d = bias.unsqueeze(0).contiguous() # [1, N] - y_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) + h_pre_out = torch.empty(M, n, dtype=proj_acc.dtype, device=dev) + h_post_out = torch.empty(M, n, dtype=proj_acc.dtype, device=dev) + h_res_out = torch.empty(M, N - 2 * n, dtype=proj_acc.dtype, device=dev) r_out = torch.empty(M, 1, dtype=proj_acc.dtype, device=dev) proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) @@ -1111,7 +1130,7 @@ def _make_args(tm): return ( proj_acc, norm_acc, bias_2d, alpha_pre, alpha_post, alpha_res, - y_out, r_out, proj_out, + h_pre_out, h_post_out, h_res_out, r_out, proj_out, M, N, K, n, eps, tm, tile_n, split_k, ) @@ -1141,7 +1160,7 @@ def _make_args(tm): _ct_reduce_compute_h_kernel, _make_args(best_tm), ) - return y_out, r_out, proj_out + return h_pre_out, h_post_out, h_res_out, r_out, proj_out # -- Combined proj_rms + compute_h forward -------------------------------- @@ -1154,14 +1173,16 @@ def _cutile_proj_rms_compute_h_fwd( alpha_res: Tensor, n: int, eps: float, - ) -> Tuple[Tensor, Tensor, Tensor]: + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Fused proj_rms + compute_h forward. Launches the existing _ct_proj_rms_fwd_kernel (split-K matmul + partial norm), then _ct_reduce_compute_h_kernel (reduce + r + activations). Returns: - y_activated: [M, N] activated h values + h_pre: [M, n] activated pre weights + h_post: [M, n] activated post weights + h_res: [M, n*n] residual logits r: [M, 1] r = norm / sqrt(K) proj_reduced: [M, N] reduced projection (for backward) """ @@ -1244,13 +1265,13 @@ def _cutile_proj_rms_compute_h_fwd( ) # Launch reduce + compute_h kernel - y_activated, r, proj_reduced = _cutile_reduce_compute_h( + h_pre, h_post, h_res, r, proj_reduced = _cutile_reduce_compute_h( proj_acc, norm_acc, bias, alpha_pre, alpha_post, alpha_res, n, M, N, K, eps, tm, TILE_N, split_k, ) - return y_activated, r, proj_reduced + return h_pre, h_post, h_res, r, proj_reduced def _proj_rms_bwd_autotune_configs(N): """Generate autotune search space for proj_rms backward kernel (K >= 8192 path).""" @@ -1334,8 +1355,11 @@ def _cutile_proj_rms_bwd( @ct.kernel def _ct_fused_grad_h_proj_kernel( - GRAD_Y, # [M, N] - Y_ACT, # [M, N] + GRAD_H_PRE, # [M, n] + GRAD_H_POST, # [M, n] + GRAD_H_RES, # [M, n*n] + H_PRE, # [M, n] + H_POST, # [M, n] PROJ, # [M, N] R, # [M, 1] GRAD_R_EXT, # [M, 1] @@ -1347,10 +1371,13 @@ def _ct_fused_grad_h_proj_kernel( GRAD_R_TOTAL, # [M, 1] output M: int, N: int, - n: int, + n: ConstInt, eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, + HAS_GRAD_H_PRE: ConstInt, + HAS_GRAD_H_POST: ConstInt, + HAS_GRAD_H_RES: ConstInt, HAS_GRAD_R_EXT: ConstInt, ): """Precompute grad_h, grad_proj, and grad_r_total for downstream backward kernels. @@ -1363,56 +1390,105 @@ def _ct_fused_grad_h_proj_kernel( alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() - offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) - one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) - zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) - mask_pre = ct.where(ct.less(offsets, n), one, zero) - mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre - mask_res = one - mask_pre - mask_post + r_tile = ct.load( + R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + padding_mode=PAD_ZERO, + ) + r_tile = ct.astype(r_tile, ct.float32) - scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res - scale_2d = ct.reshape(scale, (1, TILE_SIZE_N)) - mask_pre_2d = ct.reshape(mask_pre, (1, TILE_SIZE_N)) - mask_post_2d = ct.reshape(mask_post, (1, TILE_SIZE_N)) - mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_r_from_h = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) - gy_tile = ct.load( - GRAD_Y, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + # Clear the padded columns once inside this kernel. Valid columns are + # overwritten by the segment stores below. + zero_full = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0, dtype=ct.float32) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_PRE: + gy_pre = ct.load( + GRAD_H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + gy_pre = ct.astype(gy_pre, ct.float32) + else: + gy_pre = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_pre = ct.load( + H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) - gy_tile = ct.astype(gy_tile, ct.float32) - y_tile = ct.load( - Y_ACT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + h_pre = ct.astype(h_pre, ct.float32) + proj_pre = ct.load( + PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) - y_tile = ct.astype(y_tile, ct.float32) - proj_tile = ct.load( - PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + proj_pre = ct.astype(proj_pre, ct.float32) + grad_h_pre = gy_pre * h_pre * (1.0 - h_pre) + grad_proj_pre = grad_h_pre * alpha_pre * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), + axis=1, keepdims=True, + ) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h_pre.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_pre.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_POST: + gy_post = ct.load( + GRAD_H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + gy_post = ct.astype(gy_post, ct.float32) + else: + gy_post = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_post = ct.load( + H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) - proj_tile = ct.astype(proj_tile, ct.float32) - r_tile = ct.load( - R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), + h_post = ct.astype(h_post, ct.float32) + proj_post = ct.load( + PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) - r_tile = ct.astype(r_tile, ct.float32) - - # Activation backward → grad_h - gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d - half_y = y_tile * 0.5 - gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d - gh_res = gy_tile * mask_res_2d - grad_h = gh_pre + gh_post + gh_res - - r_eps = r_tile + eps - inv_r_eps = 1.0 / r_eps - grad_proj_tile = grad_h * scale_2d * inv_r_eps - - # grad_r_total: reduction over N → [TILE_M, 1] - grad_r_from_h = ct.sum( - grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), + proj_post = ct.astype(proj_post, ct.float32) + half_h_post = h_post * 0.5 + grad_h_post = gy_post * half_h_post * (1.0 - half_h_post) * 2.0 + grad_proj_post = grad_h_post * alpha_post * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True, ) + ct.store(GRAD_H, index=(tile_m_id, 1), tile=grad_h_post.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 1), tile=grad_proj_post.astype(GRAD_PROJ.dtype)) + + for res_chunk in ct.static_iter(range(n)): + if HAS_GRAD_H_RES: + grad_h_res = ct.load( + GRAD_H_RES, index=(tile_m_id, res_chunk), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + grad_h_res = ct.astype(grad_h_res, ct.float32) + else: + grad_h_res = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + proj_res = ct.load( + PROJ, index=(tile_m_id, 2 + res_chunk), shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + proj_res = ct.astype(proj_res, ct.float32) + grad_proj_res = grad_h_res * alpha_res * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_res * proj_res * alpha_res * (-inv_r_eps * inv_r_eps), + axis=1, keepdims=True, + ) + ct.store( + GRAD_H, index=(tile_m_id, 2 + res_chunk), + tile=grad_h_res.astype(GRAD_H.dtype), + ) + ct.store( + GRAD_PROJ, index=(tile_m_id, 2 + res_chunk), + tile=grad_proj_res.astype(GRAD_PROJ.dtype), + ) + if HAS_GRAD_R_EXT: grad_r_ext_tile = ct.load( GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), @@ -1423,8 +1499,6 @@ def _ct_fused_grad_h_proj_kernel( grad_r_ext_tile = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) grad_r_total = grad_r_from_h + grad_r_ext_tile - ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h.astype(GRAD_H.dtype)) - ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_tile.astype(GRAD_PROJ.dtype)) ct.store(GRAD_R_TOTAL, index=(tile_m_id, 0), tile=grad_r_total.astype(GRAD_R_TOTAL.dtype)) @ct.kernel @@ -1617,23 +1691,15 @@ def _ct_scalar_grads_reduce_kernel( def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( X, # [M, K] WEIGHT, # [N, K] - GRAD_Y, # [M, N] - Y_ACT, # [M, N] - PROJ, # [M, N] + GRAD_PROJ, # [M, TILE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed R, # [M, 1] - GRAD_R_EXT, # [M, 1] - Alpha_pre, # [1] - Alpha_post, # [1] - Alpha_res, # [1] GRAD_X, # [M, K] output GRAD_WEIGHT, # [N, K] output M: int, N: int, K: int, - n: int, - eps: float, TILE_N_SIZE: ConstInt, - HAS_GRAD_R_EXT: ConstInt, ): """Fused backward (small K path) with work-stealing. @@ -1644,24 +1710,6 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( """ zero_pad = ct.PaddingMode.ZERO - alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() - alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() - alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() - - # Build masks and scale - offsets = ct.arange(TILE_N_SIZE, dtype=ct.int32) - one = ct.full((TILE_N_SIZE,), 1.0, dtype=ct.float32) - zero_v = ct.full((TILE_N_SIZE,), 0.0, dtype=ct.float32) - mask_pre = ct.where(ct.less(offsets, n), one, zero_v) - mask_post = ct.where(ct.less(offsets, 2 * n), one, zero_v) - mask_pre - mask_res = one - mask_pre - mask_post - - scale = alpha_pre * mask_pre + alpha_post * mask_post + alpha_res * mask_res - scale_2d = ct.reshape(scale, (1, TILE_N_SIZE)) - mask_pre_2d = ct.reshape(mask_pre, (1, TILE_N_SIZE)) - mask_post_2d = ct.reshape(mask_post, (1, TILE_N_SIZE)) - mask_res_2d = ct.reshape(mask_res, (1, TILE_N_SIZE)) - TILE_DB_SIZE_M = 128 TILE_DB_SIZE_K = 64 NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) @@ -1676,31 +1724,10 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( X, index=(m_tile, tile_id), shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), padding_mode=zero_pad, ) - # Compute grad_proj on-the-fly - gy_tile = ct.load( - GRAD_Y, index=(m_tile, 0), - shape=(TILE_DB_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, - ) - gy_tile = ct.astype(gy_tile, ct.float32) - y_tile = ct.load( - Y_ACT, index=(m_tile, 0), + grad_proj_tile = ct.load( + GRAD_PROJ, index=(m_tile, 0), shape=(TILE_DB_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, ) - y_tile = ct.astype(y_tile, ct.float32) - r_tile = ct.load( - R, index=(m_tile, 0), - shape=(TILE_DB_SIZE_M, 1), padding_mode=zero_pad, - ) - r_tile = ct.astype(r_tile, ct.float32) - - gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d - half_y = y_tile * 0.5 - gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d - gh_res = gy_tile * mask_res_2d - grad_h = gh_pre + gh_post + gh_res - r_eps = r_tile + eps - inv_r_eps = 1.0 / r_eps - grad_proj_tile = grad_h * scale_2d * inv_r_eps accumulator_db = ct.mma( x_tile.transpose().astype(ct.tfloat32), @@ -1725,52 +1752,20 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( b_tile_idx = tile_id % NUM_DA_K_TILES dd_tile_idx = tile_id // NUM_DA_K_TILES - # Compute grad_proj on-the-fly - gy_tile = ct.load( - GRAD_Y, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, - ) - gy_tile = ct.astype(gy_tile, ct.float32) - y_tile = ct.load( - Y_ACT, index=(dd_tile_idx, 0), + grad_proj_tile = ct.load( + GRAD_PROJ, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, ) - y_tile = ct.astype(y_tile, ct.float32) - proj_tile = ct.load( - PROJ, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + grad_r_total = ct.load( + GRAD_R_TOTAL, index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, ) - proj_tile = ct.astype(proj_tile, ct.float32) r_tile = ct.load( R, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, ) r_tile = ct.astype(r_tile, ct.float32) - gh_pre = gy_tile * y_tile * (1.0 - y_tile) * mask_pre_2d - half_y = y_tile * 0.5 - gh_post = gy_tile * half_y * (1.0 - half_y) * 2.0 * mask_post_2d - gh_res = gy_tile * mask_res_2d - grad_h = gh_pre + gh_post + gh_res - r_eps = r_tile + eps - inv_r_eps = 1.0 / r_eps - grad_proj_tile = grad_h * scale_2d * inv_r_eps - - # Simplified rms_dnorm: grad_r_total * x / (r * K) - grad_r_from_h = ct.sum( - grad_h * proj_tile * scale_2d * (-inv_r_eps * inv_r_eps), - axis=1, keepdims=True, - ) - if HAS_GRAD_R_EXT: - grad_r_ext_tile = ct.load( - GRAD_R_EXT, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, - ) - grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) - else: - grad_r_ext_tile = ct.full((TILE_DA_SIZE_M, 1), 0.0, dtype=ct.float32) - grad_r_total = grad_r_from_h + grad_r_ext_tile - x_tile = ct.load( X, index=(dd_tile_idx, b_tile_idx), shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), padding_mode=zero_pad, @@ -1803,8 +1798,12 @@ def _fused_grad_x_weight_autotune_configs(N): def _cutile_fused_compute_h_proj_rms_bwd( x: Tensor, weight: Tensor, - grad_y: Tensor, - y_activated: Tensor, + grad_h_pre: Tensor, + grad_h_post: Tensor, + grad_h_res: Tensor, + h_pre: Tensor, + h_post: Tensor, + h_res: Tensor, proj: Tensor, r: Tensor, grad_r_ext: Tensor, @@ -1837,23 +1836,32 @@ def _cutile_fused_compute_h_proj_rms_bwd( has_grad_r_ext = grad_r_ext is not None has_grad_r_ext_flag = int(has_grad_r_ext) grad_r_ext_arg = grad_r_ext if has_grad_r_ext else r + has_grad_h_pre = grad_h_pre is not None + has_grad_h_post = grad_h_post is not None + has_grad_h_res = grad_h_res is not None + grad_h_pre_arg = grad_h_pre if has_grad_h_pre else h_pre + grad_h_post_arg = grad_h_post if has_grad_h_post else h_post + grad_h_res_arg = grad_h_res if has_grad_h_res else h_res # 0. Precompute grad_h, grad_proj, grad_r_total grad_h_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) grad_proj_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) grad_r_total_buf = torch.empty(M, 1, dtype=torch.float32, device=dev) - tile_m_precomp = min(128, M) + tile_m_precomp = _default_tile_m(M) ct.launch( stream, (math.ceil(M / tile_m_precomp),), _ct_fused_grad_h_proj_kernel, ( - grad_y, y_activated, proj, r, grad_r_ext_arg, + grad_h_pre_arg, grad_h_post_arg, grad_h_res_arg, + h_pre, h_post, proj, r, grad_r_ext_arg, alpha_pre, alpha_post, alpha_res, grad_h_buf, grad_proj_buf, grad_r_total_buf, M, N, n, eps, - tile_m_precomp, TILE_N, has_grad_r_ext_flag, + tile_m_precomp, TILE_N, + int(has_grad_h_pre), int(has_grad_h_post), int(has_grad_h_res), + has_grad_r_ext_flag, ), ) @@ -1916,11 +1924,10 @@ def _cutile_fused_compute_h_proj_rms_bwd( (num_sms, 2, 1), _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, ( - x, weight, grad_y, y_activated, proj, r, grad_r_ext_arg, - alpha_pre, alpha_post, alpha_res, + x, weight, grad_proj_buf, grad_r_total_buf, r, grad_x, grad_weight, - M, N, K, n, eps, - TILE_N, has_grad_r_ext_flag, + M, N, K, + TILE_N, ), ) @@ -2086,7 +2093,7 @@ def _torch_proj_rms_compute_h( bias: Tensor, n: int, eps: float, -) -> Tuple[Tensor, Tensor]: +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: proj = torch.matmul(x, weight.t()) r = x.float().norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) alpha = torch.cat( @@ -2100,8 +2107,8 @@ def _torch_proj_rms_compute_h( h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) h_pre = h[..., :n].sigmoid() h_post = h[..., n : 2 * n].sigmoid() * 2 - y_activated = torch.cat([h_pre, h_post, h[..., 2 * n :]], dim=-1) - return y_activated, r.to(dtype=x.dtype) + h_res = h[..., 2 * n :] + return h_pre, h_post, h_res, r.to(dtype=x.dtype) if _CUTILE_AVAILABLE: @@ -2168,27 +2175,28 @@ def forward( n: int, eps: float = 1e-6, ): - y_activated, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( + h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, ) ctx.save_for_backward( - x, weight, y_activated, proj_reduced, r, + x, weight, h_pre, h_post, h_res, proj_reduced, r, alpha_pre, alpha_post, alpha_res, bias, ) ctx.n = n ctx.eps = eps - return y_activated, r + return h_pre, h_post, h_res, r @staticmethod - def backward(ctx, grad_y, grad_r_ext): + def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): ( - x, weight, y_activated, proj, r, + x, weight, h_pre, h_post, h_res, proj, r, alpha_pre, alpha_post, alpha_res, bias_param, ) = ctx.saved_tensors grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias = ( _cutile_fused_compute_h_proj_rms_bwd( - x, weight, grad_y, y_activated, proj, r, grad_r_ext, + x, weight, grad_h_pre, grad_h_post, grad_h_res, + h_pre, h_post, h_res, proj, r, grad_r_ext, alpha_pre, alpha_post, alpha_res, bias_param, ctx.n, ctx.eps, ) @@ -2305,8 +2313,8 @@ def fused_proj_rms_compute_h( bias: Tensor, n: int, eps: float = 1e-6, -) -> Tuple[Tensor, Tensor]: - """Projection + RMS norm + compute_h activations using cuTile, then torch.""" +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" if _CUTILE_AVAILABLE: return CutileProjRmsComputeH.apply( x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index b5ee204d890..0f122a50efb 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -301,14 +301,14 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: # Fused path: proj_rms + compute_h in one kernel launch sequence x_2d = x.reshape(s * b, self.n * self.hidden_size) with torch.cuda.nvtx.range("HyperConnection::fused_proj_rms_compute_h"): - y_activated, r = self._proj_rms_compute_h_op( + h_pre, h_post, h_res, _ = self._proj_rms_compute_h_op( x_2d, self.mapping_proj.weight, self.alpha_pre, self.alpha_post, self.alpha_res, self.bias, self.n, self.norm_eps, ) - h_pre = y_activated[:, : self.n].view(s, b, self.n) - h_post = y_activated[:, self.n : 2 * self.n].view(s, b, self.n) - h_res = y_activated[:, 2 * self.n :].view(s, b, self.n, self.n) + h_pre = h_pre.view(s, b, self.n) + h_post = h_post.view(s, b, self.n) + h_res = h_res.view(s, b, self.n, self.n) else: # Native path: separate proj_rms + _compute_h with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"): diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index da50649d3e4..1ca1f29574d 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -124,8 +124,7 @@ def _ref_proj_rms_compute_h( h_pre = h[..., :n].sigmoid() h_post = h[..., n : 2 * n].sigmoid() * 2 h_res = h[..., 2 * n :] - y = torch.cat([h_pre, h_post, h_res], dim=-1) - return y, r + return h_pre, h_post, h_res, r # ============================================================================ @@ -684,6 +683,9 @@ def test_fwd_bwd_vs_reference(self, M, n, K): ar_data = _rand(1) bias_data = _rand(N) grad_y = _rand(M, N) + grad_h_pre = grad_y[:, :n] + grad_h_post = grad_y[:, n : 2 * n] + grad_h_res = grad_y[:, 2 * n :] grad_r = _rand(M, 1) def _make_inputs(): @@ -698,15 +700,33 @@ def _make_inputs(): # -- fused path -- xf, wf, apf, apof, arf, bf = _make_inputs() - y_f, r_f = fused_proj_rms_compute_h(xf, wf, apf, apof, arf, bf, n, eps) - (y_f * grad_y + r_f * grad_r).sum().backward() + h_pre_f, h_post_f, h_res_f, r_f = fused_proj_rms_compute_h( + xf, wf, apf, apof, arf, bf, n, eps + ) + loss_f = ( + (h_pre_f * grad_h_pre).sum() + + (h_post_f * grad_h_post).sum() + + (h_res_f * grad_h_res).sum() + + (r_f * grad_r).sum() + ) + loss_f.backward() # -- reference path -- xr, wr, apr, apor, arr, br = _make_inputs() - y_r, r_r = _ref_proj_rms_compute_h(xr, wr, apr, apor, arr, br, n, eps) - (y_r * grad_y + r_r * grad_r).sum().backward() + h_pre_r, h_post_r, h_res_r, r_r = _ref_proj_rms_compute_h( + xr, wr, apr, apor, arr, br, n, eps + ) + loss_r = ( + (h_pre_r * grad_h_pre).sum() + + (h_post_r * grad_h_post).sum() + + (h_res_r * grad_h_res).sum() + + (r_r * grad_r).sum() + ) + loss_r.backward() - torch.testing.assert_close(y_f, y_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="y_activated mismatch") + torch.testing.assert_close(h_pre_f, h_pre_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_pre mismatch") + torch.testing.assert_close(h_post_f, h_post_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post mismatch") + torch.testing.assert_close(h_res_f, h_res_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_res mismatch") torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="r mismatch") torch.testing.assert_close( xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" @@ -941,13 +961,13 @@ def _run_fused_compute_h(): bias_p = bias_data.clone().requires_grad_(True) x_2d = hs.reshape(s * b, n * C) - y_activated, r = fused_proj_rms_compute_h( + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( x_2d, w, ap, apo, ar, bias_p, n, eps, ) - h_pre = y_activated[:, :n].view(s, b, n) - h_post = y_activated[:, n : 2 * n].view(s, b, n) - h_res_logits = y_activated[:, 2 * n :].view(s, b, n, n) + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) @@ -969,13 +989,13 @@ def _run_ref(): bias_p = bias_data.clone().requires_grad_(True) x_2d = hs.reshape(s * b, n * C) - y_ref, r_ref = _ref_proj_rms_compute_h( + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( x_2d, w, ap, apo, ar, bias_p, n, eps, ) - h_pre = y_ref[:, :n].view(s, b, n) - h_post = y_ref[:, n : 2 * n].view(s, b, n) - h_res_logits = y_ref[:, 2 * n :].view(s, b, n, n) + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) @@ -1283,13 +1303,13 @@ def _run_fused_compute_h_broadcast(): hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) x_2d = hs_map.reshape(s * b, n * C) - y_activated, r = fused_proj_rms_compute_h( + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( x_2d, w, ap, apo, ar, bias_p, n, eps, ) - h_pre = y_activated[:, :n].view(s, b, n) - h_post = y_activated[:, n : 2 * n].view(s, b, n) - h_res_logits = y_activated[:, 2 * n :].view(s, b, n, n) + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) @@ -1311,13 +1331,13 @@ def _run_ref(): bias_p = bias_data.clone().requires_grad_(True) x_2d = hs.reshape(s * b, n * C) - y_ref, r_ref = _ref_proj_rms_compute_h( + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( x_2d, w, ap, apo, ar, bias_p, n, eps, ) - h_pre = y_ref[:, :n].view(s, b, n) - h_post = y_ref[:, n : 2 * n].view(s, b, n) - h_res_logits = y_ref[:, 2 * n :].view(s, b, n, n) + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) From 3765234a6b4235833653c67e890a807408bea1d4 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 5 May 2026 13:10:05 +0800 Subject: [PATCH 07/34] chore: clean up mHC fused kernel files --- megatron/core/fusions/fused_mhc_kernels.py | 579 ++++++++++++++++-- megatron/core/fusions/triton_mhc_kernels.py | 549 ----------------- .../fusions/test_fused_mhc_kernels.py | 14 +- .../test_fused_mhc_kernels_autotune.py | 237 ------- .../fusions/verify_compute_h_bwd.py | 198 ------ .../test_hyper_connection_recompute.py | 433 ------------- 6 files changed, 543 insertions(+), 1467 deletions(-) delete mode 100644 megatron/core/fusions/triton_mhc_kernels.py delete mode 100644 tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py delete mode 100644 tests/unit_tests/fusions/verify_compute_h_bwd.py diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 52b519cab6b..ef62c4ebb02 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1,11 +1,11 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Fused cuTile kernels for mHC (Manifold-Constrained Hyper-Connections). +"""Fused kernels for mHC (Manifold-Constrained Hyper-Connections). -Requires cuda.tile (cuTile) for optimal performance on supported GPUs -(compute capability 10.x+). Reference (non-fused) implementations live in -``megatron.core.transformer.hyper_connection`` and are used when cuTile is -unavailable or when the ``use_fused_mhc`` config flag is False. +Uses Triton and cuda.tile (cuTile) kernels when available, with PyTorch +reference implementations as fallback. Reference (non-fused) implementations +live in ``megatron.core.transformer.hyper_connection`` and are used when fused +kernels are unavailable or when the ``use_fused_mhc`` config flag is False. Four fused operations: - sinkhorn: Sinkhorn-Knopp projection to doubly stochastic matrix @@ -62,6 +62,533 @@ def is_triton_available() -> bool: return _TRITON_AVAILABLE +# ============================================================================ +# Triton implementations (only defined when triton is available) +# ============================================================================ + +if _TRITON_AVAILABLE: + # ============================================================================ + # Sinkhorn-Knopp + # ============================================================================ + + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], + key=["HC", "NUM_ITERS"], + ) + @triton.jit + def _triton_sinkhorn_fwd_kernel( + inp_ptr, out_ptr, M_init_ptr, + N_batch, eps, + HC: tl.constexpr, NUM_ITERS: tl.constexpr, + ): + """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) + row_max = tl.max(logits, axis=1) + M = tl.exp2((logits - row_max[:, None]) * 1.4426950408889634) + tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) + + for _ in range(NUM_ITERS): + row_sum = tl.sum(M, axis=1) + M = M / (row_sum[:, None] + eps) + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + + tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) + + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], + key=["HC", "NUM_ITERS"], + ) + @triton.jit + def _triton_sinkhorn_bwd_kernel( + grad_out_ptr, M_init_ptr, grad_inp_ptr, + ws_M_ptr, ws_rs_ptr, ws_cs_ptr, + N_batch, eps, + HC: tl.constexpr, NUM_ITERS: tl.constexpr, + ): + """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + M_ws_base = pid * 2 * NUM_ITERS * HC * HC + v_ws_base = pid * NUM_ITERS + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + M = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + for t in range(NUM_ITERS): + ws_off = M_ws_base + (2 * t) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + row_sum = tl.sum(M, axis=1) + tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) + M = M / (row_sum[:, None] + eps) + + ws_off = M_ws_base + (2 * t + 1) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + col_sum = tl.sum(M, axis=0) + tl.store(ws_cs_ptr + (v_ws_base + t) * HC + offs_c, col_sum) + M = M / (col_sum[None, :] + eps) + + # M is the final forward output. It is the right value for the first VJP + # through the last column-normalization step. + grad = tl.load(grad_out_ptr + mat_ptrs).to(tl.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + + col_s = tl.load(ws_cs_ptr + (v_ws_base + t) * HC + offs_c).to(tl.float32) + grad = grad / (col_s[None, :] + eps) + col_corr = tl.sum(grad * M, axis=0) + grad = grad - col_corr[None, :] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t + 1) * HC * HC + + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) + grad = grad / (row_s[:, None] + eps) + row_corr = tl.sum(grad * M, axis=1) + grad = grad - row_corr[:, None] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t) * HC * HC + + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + grad = grad * M_init + tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) + + + def _triton_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-8 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.contiguous().view(N_batch, hc, hc) + _triton_sinkhorn_fwd_kernel[(N_batch,)]( + inp, out, M_init, N_batch, eps, hc, num_iterations, + ) + return out.view(original_shape), M_init.view(original_shape) + + + def _triton_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.contiguous().view(N_batch, hc, hc) + mi = M_init.contiguous().view(N_batch, hc, hc) + ws_M = torch.empty(N_batch * 2 * num_iterations * hc * hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + _triton_sinkhorn_bwd_kernel[(N_batch,)]( + go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations, + ) + return grad_input.view(original_shape) + + + class TritonFusedSinkhorn(torch.autograd.Function): + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, grad_output: Tensor): + (M_init,) = ctx.saved_tensors + grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + + def triton_fused_sinkhorn( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) + + + # ============================================================================ + # H_aggregate forward + # ============================================================================ + + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_h_agg_fwd_kernel( + x_ptr, h_ptr, out_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_x_s, stride_x_n, stride_x_c, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + ): + """out[s, c] = sum_i x[s, i, c] * h[s, i].""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for i in tl.static_range(N): + x_i = tl.load( + x_ptr + offs_s[:, None] * stride_x_s + i * stride_x_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + h_i = tl.load(h_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + acc += h_i[:, None] * x_i + tl.store( + out_ptr + offs_s[:, None] * C + offs_c[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + + def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.contiguous().view(sb, n, C) + h_flat = h_pre.contiguous().view(sb, n) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_h_agg_fwd_kernel[grid]( + x_flat, h_flat, out, sb, C, n, + x_flat.stride(0), x_flat.stride(1), x_flat.stride(2), + ) + return out.view(s, b, C) + + + # ============================================================================ + # H_post BDA + # ============================================================================ + + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_fwd_kernel( + hr_ptr, orig_ptr, hp_ptr, x_ptr, bias_ptr, out_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_hr_s, stride_hr_i, stride_hr_j, + stride_orig_s, stride_orig_n, stride_orig_c, + stride_out_s, stride_out_n, stride_out_c, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + ): + """out = hr @ orig + hp * (x + bias).""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + hp_i = tl.load(hp_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + out_i = hp_i[:, None] * x_tile + + for j in tl.static_range(N): + hr_ij = tl.load( + hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + out_i += hr_ij[:, None] * orig_j + + tl.store( + out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], + out_i.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + + def _triton_h_post_bda_fwd( + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=dev) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_fwd_kernel[grid]( + hr_flat, orig_flat, hp_flat, x_flat, bias if bias is not None else x_flat, out, + sb, C, n, + hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), + orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), + out.stride(0), out.stride(1), out.stride(2), + HAS_BIAS=(bias is not None), + ) + return out.view(s, b, n, C) + + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_x_orig_kernel( + go_ptr, hr_ptr, hp_ptr, g_orig_ptr, g_x_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_go_s, stride_go_n, stride_go_c, + stride_hr_s, stride_hr_i, stride_hr_j, + stride_orig_s, stride_orig_n, stride_orig_c, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + ): + """g_x = hp @ go, g_orig = hr.T @ go.""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + g_x_acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, other=0.0, + ).to(tl.float32) + hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) + g_x_acc += hp_j[:, None] * go_j + tl.store( + g_x_ptr + offs_s[:, None] * C + offs_c[None, :], + g_x_acc.to(g_x_ptr.dtype.element_ty), + mask=mask_2d, + ) + + for i in tl.static_range(N): + g_orig_i = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, other=0.0, + ).to(tl.float32) + hr_ji = tl.load( + hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, + mask=mask_s, other=0.0, + ).to(tl.float32) + g_orig_i += hr_ji[:, None] * go_j + tl.store( + g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], + g_orig_i.to(g_orig_ptr.dtype.element_ty), + mask=mask_2d, + ) + + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_hp_hr_kernel( + go_ptr, orig_ptr, x_ptr, bias_ptr, + g_hr_ptr, g_hp_ptr, + sb, C: tl.constexpr, N: tl.constexpr, + stride_go_s, stride_go_n, stride_go_c, + stride_orig_s, stride_orig_n, stride_orig_c, + stride_hr_s, stride_hr_i, stride_hr_j, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + ): + """g_hp = sum_c go*(x+bias), g_hr = go @ orig.T.""" + pid_s = tl.program_id(0) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + mask_s = offs_s < sb + + g_hp_acc = tl.zeros((BLOCK_S, N), dtype=tl.float32) + g_hr_acc = tl.zeros((BLOCK_S, N * N), dtype=tl.float32) + + for c_start in range(0, C, BLOCK_C): + offs_c = c_start + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + go_i = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + i * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hp = tl.sum(go_i * x_tile, axis=1) + g_hp_acc += tl.where( + tl.arange(0, N)[None, :] == i, + dot_hp[:, None], + tl.zeros((BLOCK_S, N), dtype=tl.float32), + ) + for j in tl.static_range(N): + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hr = tl.sum(go_i * orig_j, axis=1) + g_hr_acc += tl.where( + tl.arange(0, N * N)[None, :] == i * N + j, + dot_hr[:, None], + tl.zeros((BLOCK_S, N * N), dtype=tl.float32), + ) + + offs_n = tl.arange(0, N) + tl.store( + g_hp_ptr + offs_s[:, None] * N + offs_n[None, :], + g_hp_acc.to(g_hp_ptr.dtype.element_ty), + mask=mask_s[:, None], + ) + + nn_offs = tl.arange(0, N * N) + for i in tl.static_range(N): + for j in tl.static_range(N): + col_mask = (nn_offs == (i * N + j)).to(tl.float32) + val = tl.sum(g_hr_acc * col_mask[None, :], axis=1) + tl.store( + g_hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + val.to(g_hr_ptr.dtype.element_ty), + mask=mask_s, + ) + + + def _triton_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=dev) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=dev) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=dev) + g_x = torch.empty(sb, C, dtype=x.dtype, device=dev) + + go_flat = grad_output.contiguous().view(sb, n, C) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_bwd_g_x_orig_kernel[grid_a]( + go_flat, hr_flat, hp_flat, g_res, g_x, sb, C, n, + go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), + hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), + g_res.stride(0), g_res.stride(1), g_res.stride(2), + ) + + grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) + _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( + go_flat, orig_flat, x_flat, bias if bias is not None else x_flat, + g_hr, g_hp, sb, C, n, + go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), + orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), + g_hr.stride(0), g_hr.stride(1), g_hr.stride(2), + HAS_BIAS=(bias is not None), + ) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + # ============================================================================ # CuTile implementations (only defined when cuda.tile is available) # ============================================================================ @@ -688,16 +1215,10 @@ def _ct_proj_rms_fwd_kernel( a_tile.astype(ct.tfloat32), b_tile.transpose().astype(ct.tfloat32), acc=acc ) sum_sq += ct.sum(a_tile * a_tile, axis=1, keepdims=True) - # norm_tile = ct.sqrt(sum_sq) bid_m_k = tile_m_id + split_k_id * num_m_tiles ct.store(PROJ, index=(bid_m_k, 0), tile=acc.astype(PROJ.dtype)) ct.store(NORM, index=(bid_m_k, 0), tile=sum_sq.astype(NORM.dtype)) - # v = norm_tile / ct.sqrt(K) + eps - # r_tile = 1.0 / v - # ct.atomic_add(PROJ, proj_indices, acc.astype(PROJ.dtype)) - # ct.atomic_add(NORM, norm_indices, norm_tile.astype(NORM.dtype)) - # ct.store(R, index=(tile_m_id, 0), tile=r_tile.astype(R.dtype)) # -- Sigmoid helper for cuTile kernels ------------------------------------ @@ -1000,9 +1521,6 @@ def _cutile_proj_rms_fwd( else: tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) - # print(cached) - # print('hahahahahah') - # split_k = 1 proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) @@ -1015,7 +1533,6 @@ def _cutile_proj_rms_fwd( ) proj = proj.view(split_k, M, N).sum(dim=0).to(dtype=x.dtype) norm = norm.view(split_k, M, 1).sum(dim=0).to(dtype=x.dtype) - # r = r.view(split_k, M, 1).sum(dim=0) else: # Autotune on first call for this shape. from types import SimpleNamespace @@ -1052,9 +1569,6 @@ def _cutile_proj_rms_fwd( search_space=configs, ) best = tuned.tuned_config - # print(best) - # print('hahahahahah') - # best.SPLIT_K = 1 _proj_rms_fwd_best_cfg[cache_key] = ( best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K ) @@ -1073,7 +1587,6 @@ def _cutile_proj_rms_fwd( proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) norm = norm.view(best.SPLIT_K, M, 1).sum(dim=0) norm = torch.sqrt(norm) - # r = r.view(split_k, M, 1).sum(dim=0) r = 1.0 / (norm / math.sqrt(K) + eps) return proj, norm, r @@ -2004,45 +2517,25 @@ def _cutile_fused_compute_h_proj_rms_bwd( def _get_triton_sinkhorn(): if not _TRITON_AVAILABLE: return None - try: - from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn - - return triton_fused_sinkhorn - except ImportError: - return None + return globals().get("triton_fused_sinkhorn") def _get_triton_h_aggregate_fwd(): if not _TRITON_AVAILABLE: return None - try: - from megatron.core.fusions.triton_mhc_kernels import _triton_h_aggregate_fwd - - return _triton_h_aggregate_fwd - except ImportError: - return None + return globals().get("_triton_h_aggregate_fwd") def _get_triton_h_post_bda_fwd(): if not _TRITON_AVAILABLE: return None - try: - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_fwd - - return _triton_h_post_bda_fwd - except ImportError: - return None + return globals().get("_triton_h_post_bda_fwd") def _get_triton_h_post_bda_bwd(): if not _TRITON_AVAILABLE: return None - try: - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd - - return _triton_h_post_bda_bwd - except ImportError: - return None + return globals().get("_triton_h_post_bda_bwd") def _torch_h_aggregate_bwd( diff --git a/megatron/core/fusions/triton_mhc_kernels.py b/megatron/core/fusions/triton_mhc_kernels.py deleted file mode 100644 index 798d880bc5d..00000000000 --- a/megatron/core/fusions/triton_mhc_kernels.py +++ /dev/null @@ -1,549 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -"""Triton kernels used by the unified mHC fused path. - -Only the Triton kernels that are faster than the cuTile implementations are kept -here: - - Sinkhorn forward/backward - - H_aggregate forward - - H_post_bda forward/backward - -The public mHC API lives in ``fused_mhc_kernels.py`` and handles fallback. -""" - -from typing import Optional, Tuple - -import torch -import triton -import triton.language as tl -from torch import Tensor - - -# ============================================================================ -# Sinkhorn-Knopp -# ============================================================================ - - -@triton.autotune( - configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], - key=["HC", "NUM_ITERS"], -) -@triton.jit -def _triton_sinkhorn_fwd_kernel( - inp_ptr, out_ptr, M_init_ptr, - N_batch, eps, - HC: tl.constexpr, NUM_ITERS: tl.constexpr, -): - """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" - pid = tl.program_id(0) - if pid >= N_batch: - return - - base = pid * HC * HC - offs_r = tl.arange(0, HC) - offs_c = tl.arange(0, HC) - mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] - - logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) - row_max = tl.max(logits, axis=1) - M = tl.exp2((logits - row_max[:, None]) * 1.4426950408889634) - tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) - - for _ in range(NUM_ITERS): - row_sum = tl.sum(M, axis=1) - M = M / (row_sum[:, None] + eps) - col_sum = tl.sum(M, axis=0) - M = M / (col_sum[None, :] + eps) - - tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) - - -@triton.autotune( - configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], - key=["HC", "NUM_ITERS"], -) -@triton.jit -def _triton_sinkhorn_bwd_kernel( - grad_out_ptr, M_init_ptr, grad_inp_ptr, - ws_M_ptr, ws_rs_ptr, ws_cs_ptr, - N_batch, eps, - HC: tl.constexpr, NUM_ITERS: tl.constexpr, -): - """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" - pid = tl.program_id(0) - if pid >= N_batch: - return - - base = pid * HC * HC - M_ws_base = pid * 2 * NUM_ITERS * HC * HC - v_ws_base = pid * NUM_ITERS - offs_r = tl.arange(0, HC) - offs_c = tl.arange(0, HC) - mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] - - M = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) - for t in range(NUM_ITERS): - ws_off = M_ws_base + (2 * t) * HC * HC - tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) - - row_sum = tl.sum(M, axis=1) - tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) - M = M / (row_sum[:, None] + eps) - - ws_off = M_ws_base + (2 * t + 1) * HC * HC - tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) - - col_sum = tl.sum(M, axis=0) - tl.store(ws_cs_ptr + (v_ws_base + t) * HC + offs_c, col_sum) - M = M / (col_sum[None, :] + eps) - - # M is the final forward output. It is the right value for the first VJP - # through the last column-normalization step. - grad = tl.load(grad_out_ptr + mat_ptrs).to(tl.float32) - for t_rev in range(NUM_ITERS): - t = NUM_ITERS - 1 - t_rev - - col_s = tl.load(ws_cs_ptr + (v_ws_base + t) * HC + offs_c).to(tl.float32) - grad = grad / (col_s[None, :] + eps) - col_corr = tl.sum(grad * M, axis=0) - grad = grad - col_corr[None, :] - M = tl.load( - ws_M_ptr + M_ws_base + (2 * t + 1) * HC * HC - + offs_r[:, None] * HC + offs_c[None, :] - ).to(tl.float32) - - row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) - grad = grad / (row_s[:, None] + eps) - row_corr = tl.sum(grad * M, axis=1) - grad = grad - row_corr[:, None] - M = tl.load( - ws_M_ptr + M_ws_base + (2 * t) * HC * HC - + offs_r[:, None] * HC + offs_c[None, :] - ).to(tl.float32) - - M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) - grad = grad * M_init - tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) - - -def _triton_sinkhorn_fwd( - input_logits: Tensor, num_iterations: int, eps: float = 1e-8 -) -> Tuple[Tensor, Tensor]: - original_shape = input_logits.shape - hc = original_shape[-1] - N_batch = input_logits.numel() // (hc * hc) - dev = input_logits.device - out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) - M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) - inp = input_logits.contiguous().view(N_batch, hc, hc) - _triton_sinkhorn_fwd_kernel[(N_batch,)]( - inp, out, M_init, N_batch, eps, hc, num_iterations, - ) - return out.view(original_shape), M_init.view(original_shape) - - -def _triton_sinkhorn_bwd( - grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 -) -> Tensor: - original_shape = grad_output.shape - hc = original_shape[-1] - N_batch = grad_output.numel() // (hc * hc) - dev = grad_output.device - grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) - go = grad_output.contiguous().view(N_batch, hc, hc) - mi = M_init.contiguous().view(N_batch, hc, hc) - ws_M = torch.empty(N_batch * 2 * num_iterations * hc * hc, dtype=torch.float32, device=dev) - ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) - ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) - _triton_sinkhorn_bwd_kernel[(N_batch,)]( - go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations, - ) - return grad_input.view(original_shape) - - -class TritonFusedSinkhorn(torch.autograd.Function): - @staticmethod - def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): - out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) - ctx.save_for_backward(M_init) - ctx.num_iterations = num_iterations - ctx.eps = eps - return out - - @staticmethod - def backward(ctx, grad_output: Tensor): - (M_init,) = ctx.saved_tensors - grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) - return grad_input, None, None - - -def triton_fused_sinkhorn( - input_logits: Tensor, num_iterations: int, eps: float = 1e-6 -) -> Tensor: - return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) - - -# ============================================================================ -# H_aggregate forward -# ============================================================================ - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) - for bc in (64, 128, 256, 512) - for bs in (1, 2, 4, 8) - for nw in (2, 4, 8) - ], - key=["C", "N"], -) -@triton.jit -def _triton_h_agg_fwd_kernel( - x_ptr, h_ptr, out_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_x_s, stride_x_n, stride_x_c, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, -): - """out[s, c] = sum_i x[s, i, c] * h[s, i].""" - pid_s = tl.program_id(0) - pid_c = tl.program_id(1) - offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) - offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) - mask_s = offs_s < sb - mask_c = offs_c < C - mask_2d = mask_s[:, None] & mask_c[None, :] - - acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) - for i in tl.static_range(N): - x_i = tl.load( - x_ptr + offs_s[:, None] * stride_x_s + i * stride_x_n + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - h_i = tl.load(h_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) - acc += h_i[:, None] * x_i - tl.store( - out_ptr + offs_s[:, None] * C + offs_c[None, :], - acc.to(out_ptr.dtype.element_ty), - mask=mask_2d, - ) - - -def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: - s, b, n, C = x.shape - sb = s * b - out = torch.empty(sb, C, dtype=x.dtype, device=x.device) - x_flat = x.contiguous().view(sb, n, C) - h_flat = h_pre.contiguous().view(sb, n) - - grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) - _triton_h_agg_fwd_kernel[grid]( - x_flat, h_flat, out, sb, C, n, - x_flat.stride(0), x_flat.stride(1), x_flat.stride(2), - ) - return out.view(s, b, C) - - -# ============================================================================ -# H_post BDA -# ============================================================================ - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) - for bc in (64, 128, 256, 512) - for bs in (1, 2, 4, 8) - for nw in (2, 4, 8) - ], - key=["C", "N"], -) -@triton.jit -def _triton_hpb_fwd_kernel( - hr_ptr, orig_ptr, hp_ptr, x_ptr, bias_ptr, out_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_hr_s, stride_hr_i, stride_hr_j, - stride_orig_s, stride_orig_n, stride_orig_c, - stride_out_s, stride_out_n, stride_out_c, - HAS_BIAS: tl.constexpr, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, -): - """out = hr @ orig + hp * (x + bias).""" - pid_s = tl.program_id(0) - pid_c = tl.program_id(1) - offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) - offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) - mask_s = offs_s < sb - mask_c = offs_c < C - mask_2d = mask_s[:, None] & mask_c[None, :] - - x_tile = tl.load( - x_ptr + offs_s[:, None] * C + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - if HAS_BIAS: - bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) - x_tile += bias_tile[None, :] - - for i in tl.static_range(N): - hp_i = tl.load(hp_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) - out_i = hp_i[:, None] * x_tile - - for j in tl.static_range(N): - hr_ij = tl.load( - hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, - mask=mask_s, - other=0.0, - ).to(tl.float32) - orig_j = tl.load( - orig_ptr - + offs_s[:, None] * stride_orig_s - + j * stride_orig_n - + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - out_i += hr_ij[:, None] * orig_j - - tl.store( - out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], - out_i.to(out_ptr.dtype.element_ty), - mask=mask_2d, - ) - - -def _triton_h_post_bda_fwd( - h_res: Tensor, - original_residual: Tensor, - h_post: Tensor, - x: Tensor, - bias: Optional[Tensor], -) -> Tensor: - s, b, n, C = original_residual.shape - sb = s * b - dev = h_res.device - out = torch.empty(sb, n, C, dtype=h_res.dtype, device=dev) - hr_flat = h_res.contiguous().view(sb, n, n) - orig_flat = original_residual.contiguous().view(sb, n, C) - hp_flat = h_post.contiguous().view(sb, n) - x_flat = x.contiguous().view(sb, C) - - grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) - _triton_hpb_fwd_kernel[grid]( - hr_flat, orig_flat, hp_flat, x_flat, bias if bias is not None else x_flat, out, - sb, C, n, - hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), - orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), - out.stride(0), out.stride(1), out.stride(2), - HAS_BIAS=(bias is not None), - ) - return out.view(s, b, n, C) - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) - for bc in (64, 128, 256, 512) - for bs in (1, 2, 4, 8) - for nw in (2, 4, 8) - ], - key=["C", "N"], -) -@triton.jit -def _triton_hpb_bwd_g_x_orig_kernel( - go_ptr, hr_ptr, hp_ptr, g_orig_ptr, g_x_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_go_s, stride_go_n, stride_go_c, - stride_hr_s, stride_hr_i, stride_hr_j, - stride_orig_s, stride_orig_n, stride_orig_c, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, -): - """g_x = hp @ go, g_orig = hr.T @ go.""" - pid_s = tl.program_id(0) - pid_c = tl.program_id(1) - offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) - offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) - mask_s = offs_s < sb - mask_c = offs_c < C - mask_2d = mask_s[:, None] & mask_c[None, :] - - g_x_acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) - for j in tl.static_range(N): - go_j = tl.load( - go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], - mask=mask_2d, other=0.0, - ).to(tl.float32) - hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) - g_x_acc += hp_j[:, None] * go_j - tl.store( - g_x_ptr + offs_s[:, None] * C + offs_c[None, :], - g_x_acc.to(g_x_ptr.dtype.element_ty), - mask=mask_2d, - ) - - for i in tl.static_range(N): - g_orig_i = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) - for j in tl.static_range(N): - go_j = tl.load( - go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], - mask=mask_2d, other=0.0, - ).to(tl.float32) - hr_ji = tl.load( - hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, - mask=mask_s, other=0.0, - ).to(tl.float32) - g_orig_i += hr_ji[:, None] * go_j - tl.store( - g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], - g_orig_i.to(g_orig_ptr.dtype.element_ty), - mask=mask_2d, - ) - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) - for bc in (64, 128, 256, 512) - for bs in (1, 2, 4, 8) - for nw in (2, 4, 8) - ], - key=["C", "N"], -) -@triton.jit -def _triton_hpb_bwd_g_hp_hr_kernel( - go_ptr, orig_ptr, x_ptr, bias_ptr, - g_hr_ptr, g_hp_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_go_s, stride_go_n, stride_go_c, - stride_orig_s, stride_orig_n, stride_orig_c, - stride_hr_s, stride_hr_i, stride_hr_j, - HAS_BIAS: tl.constexpr, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, -): - """g_hp = sum_c go*(x+bias), g_hr = go @ orig.T.""" - pid_s = tl.program_id(0) - offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) - mask_s = offs_s < sb - - g_hp_acc = tl.zeros((BLOCK_S, N), dtype=tl.float32) - g_hr_acc = tl.zeros((BLOCK_S, N * N), dtype=tl.float32) - - for c_start in range(0, C, BLOCK_C): - offs_c = c_start + tl.arange(0, BLOCK_C) - mask_c = offs_c < C - mask_2d = mask_s[:, None] & mask_c[None, :] - - x_tile = tl.load( - x_ptr + offs_s[:, None] * C + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - if HAS_BIAS: - bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) - x_tile += bias_tile[None, :] - - for i in tl.static_range(N): - go_i = tl.load( - go_ptr + offs_s[:, None] * stride_go_s + i * stride_go_n + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - dot_hp = tl.sum(go_i * x_tile, axis=1) - g_hp_acc += tl.where( - tl.arange(0, N)[None, :] == i, - dot_hp[:, None], - tl.zeros((BLOCK_S, N), dtype=tl.float32), - ) - for j in tl.static_range(N): - orig_j = tl.load( - orig_ptr - + offs_s[:, None] * stride_orig_s - + j * stride_orig_n - + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) - dot_hr = tl.sum(go_i * orig_j, axis=1) - g_hr_acc += tl.where( - tl.arange(0, N * N)[None, :] == i * N + j, - dot_hr[:, None], - tl.zeros((BLOCK_S, N * N), dtype=tl.float32), - ) - - offs_n = tl.arange(0, N) - tl.store( - g_hp_ptr + offs_s[:, None] * N + offs_n[None, :], - g_hp_acc.to(g_hp_ptr.dtype.element_ty), - mask=mask_s[:, None], - ) - - nn_offs = tl.arange(0, N * N) - for i in tl.static_range(N): - for j in tl.static_range(N): - col_mask = (nn_offs == (i * N + j)).to(tl.float32) - val = tl.sum(g_hr_acc * col_mask[None, :], axis=1) - tl.store( - g_hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, - val.to(g_hr_ptr.dtype.element_ty), - mask=mask_s, - ) - - -def _triton_h_post_bda_bwd( - grad_output: Tensor, - h_res: Tensor, - original_residual: Tensor, - h_post: Tensor, - x: Tensor, - bias: Optional[Tensor], -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: - s, b, n, C = original_residual.shape - sb = s * b - dev = h_res.device - - g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=dev) - g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=dev) - g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=dev) - g_x = torch.empty(sb, C, dtype=x.dtype, device=dev) - - go_flat = grad_output.contiguous().view(sb, n, C) - hr_flat = h_res.contiguous().view(sb, n, n) - orig_flat = original_residual.contiguous().view(sb, n, C) - hp_flat = h_post.contiguous().view(sb, n) - x_flat = x.contiguous().view(sb, C) - - grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) - _triton_hpb_bwd_g_x_orig_kernel[grid_a]( - go_flat, hr_flat, hp_flat, g_res, g_x, sb, C, n, - go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), - hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), - g_res.stride(0), g_res.stride(1), g_res.stride(2), - ) - - grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) - _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( - go_flat, orig_flat, x_flat, bias if bias is not None else x_flat, - g_hr, g_hp, sb, C, n, - go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), - orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), - g_hr.stride(0), g_hr.stride(1), g_hr.stride(2), - HAS_BIAS=(bias is not None), - ) - - g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None - return ( - g_hr.view(s, b, n, n), - g_res.view(s, b, n, C), - g_hp.view(s, b, n), - g_x.view(s, b, C), - g_bias, - ) - - -__all__ = [ - "_triton_h_aggregate_fwd", - "_triton_h_post_bda_fwd", - "_triton_h_post_bda_bwd", - "triton_fused_sinkhorn", -] diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 1ca1f29574d..13e26e4a708 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -254,7 +254,7 @@ class TestTritonHAggregate: @_require_triton @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256), (64, 8, 4, 4096)]) def test_fwd_vs_reference(self, s, b, n, C): - from megatron.core.fusions.triton_mhc_kernels import _triton_h_aggregate_fwd + from megatron.core.fusions.fused_mhc_kernels import _triton_h_aggregate_fwd _info() x_data = _rand(s, b, n, C) @@ -376,7 +376,7 @@ class TestTritonHPostBDA: @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) def test_fwd_vs_reference(self, s, b, n, C, with_bias): """Triton hpb forward output must match the PyTorch reference.""" - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_fwd + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_fwd _info() hr_data = _rand(s, b, n, n) @@ -395,7 +395,7 @@ def test_fwd_vs_reference(self, s, b, n, C, with_bias): @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)]) def test_bwd_vs_reference(self, s, b, n, C, with_bias): """Triton hpb backward grads must match the PyTorch reference.""" - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd _info() hr_data = _rand(s, b, n, n) @@ -441,7 +441,7 @@ def test_bwd_vs_reference(self, s, b, n, C, with_bias): def test_triton_vs_cutile(self, s, b, n, C, with_bias): """Triton and cuTile backward must produce identical results.""" from megatron.core.fusions.fused_mhc_kernels import _cutile_h_post_bda_bwd - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd _info() hr_data = _rand(s, b, n, n) @@ -482,7 +482,7 @@ def test_e2e_inputs_no_nan(self): fused_proj_rms, fused_sinkhorn, ) - from megatron.core.fusions.triton_mhc_kernels import _triton_h_post_bda_bwd + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd s, b, n, C = 8, 4, 4, 1024 eps = 1e-6 @@ -543,7 +543,7 @@ class TestTritonSinkhorn: @_require_triton @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10), (8, 4, 4, 20)]) def test_fwd_bwd_vs_reference(self, s, b, n, iters): - from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn + from megatron.core.fusions.fused_mhc_kernels import triton_fused_sinkhorn eps = 1e-6 data = _rand(s, b, n, n) @@ -568,7 +568,7 @@ def test_triton_vs_cutile(self, s, b, n, iters): _cutile_sinkhorn_bwd, _cutile_sinkhorn_fwd, ) - from megatron.core.fusions.triton_mhc_kernels import triton_fused_sinkhorn + from megatron.core.fusions.fused_mhc_kernels import triton_fused_sinkhorn eps = 1e-6 data = _rand(s, b, n, n) diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py b/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py deleted file mode 100644 index fd42be33c05..00000000000 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py +++ /dev/null @@ -1,237 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -"""Unit tests for cuTile mHC kernel autotuning correctness. - -Each test clears the autotune cache, runs the fused kernel (triggering a fresh -autotune search), and verifies the output matches a pure-PyTorch reference. -This ensures every config in the search space that the autotuner might select -produces numerically correct results. - -Usage: - pytest tests/unit_tests/fusions/test_fused_mhc_kernels_autotune.py -s -""" - -import math - -import pytest -import torch -from torch import Tensor - -from megatron.core.fusions.fused_mhc_kernels import ( - _CUTILE_EXPERIMENTAL_AVAILABLE, - is_cutile_available, -) - -_require_cutile_experimental = pytest.mark.skipif( - not (is_cutile_available() and _CUTILE_EXPERIMENTAL_AVAILABLE), - reason="cuTile + cuda.tile_experimental required for autotune tests", -) - - -@pytest.fixture(autouse=True) -def _skip_without_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - -DTYPE = torch.bfloat16 -DEVICE = "cuda" -FWD_ATOL, FWD_RTOL = 2e-2, 2e-2 -BWD_ATOL, BWD_RTOL = 5e-2, 5e-2 - - -def _rand(*shape, requires_grad=False): - return ( - torch.empty(*shape, dtype=DTYPE, device=DEVICE) - .uniform_(-0.1, 0.1) - .requires_grad_(requires_grad) - ) - - -# --------------------------------------------------------------------------- -# Pure-PyTorch references -# --------------------------------------------------------------------------- - - -def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: - row_max = logits.max(dim=-1, keepdim=True).values - M = torch.exp(logits - row_max) - for _ in range(num_iters): - M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps) - M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps) - return M - - -def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): - proj = torch.matmul(x, weight.t()) - norm = x.norm(dim=-1, keepdim=True) - K = x.shape[-1] - r = 1.0 / (norm / math.sqrt(K) + eps) - return proj, r - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _clear_autotune_caches(): - """Clear all cuTile mHC autotune caches to force a fresh search.""" - from megatron.core.fusions import fused_mhc_kernels as mod - - mod._sinkhorn_fwd_best_cfg.clear() - mod._sinkhorn_bwd_best_cfg.clear() - mod._proj_rms_fwd_best_cfg.clear() - mod._proj_rms_bwd_best_cfg.clear() - - -# ============================================================================ -# Sinkhorn autotune -# ============================================================================ - - -class TestAutotunesSinkhorn: - @_require_cutile_experimental - @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (4, 8, 4, 20)]) - def test_sinkhorn_fwd_autotune(self, s, b, n, iters): - """Autotuned sinkhorn fwd must match PyTorch reference.""" - from megatron.core.fusions.fused_mhc_kernels import _cutile_sinkhorn_fwd - - _clear_autotune_caches() - - eps = 1e-6 - data = _rand(s, b, n, n) - - out_fused, _ = _cutile_sinkhorn_fwd(data.clone(), iters, eps) - - out_ref = _ref_sinkhorn(data.clone(), iters, eps) - - torch.testing.assert_close(out_fused, out_ref, atol=FWD_ATOL, rtol=FWD_RTOL) - - @_require_cutile_experimental - @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (4, 8, 4, 20)]) - def test_sinkhorn_bwd_autotune(self, s, b, n, iters): - """Autotuned sinkhorn bwd grads must match PyTorch reference.""" - from megatron.core.fusions.fused_mhc_kernels import ( - _cutile_sinkhorn_bwd, - _cutile_sinkhorn_fwd, - ) - - _clear_autotune_caches() - - eps = 1e-6 - data = _rand(s, b, n, n) - grad_out = _rand(s, b, n, n) - - out_f, M_init = _cutile_sinkhorn_fwd(data.clone(), iters, eps) - grad_f = _cutile_sinkhorn_bwd(grad_out, M_init, iters, eps) - - inp_r = data.clone().requires_grad_(True) - out_r = _ref_sinkhorn(inp_r, iters, eps) - out_r.backward(grad_out) - grad_r = inp_r.grad.clone() - - torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) - torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) - - @_require_cutile_experimental - def test_sinkhorn_autotune_cache_populated(self): - """After first call, autotune cache should contain the best config.""" - from megatron.core.fusions import fused_mhc_kernels as mod - from megatron.core.fusions.fused_mhc_kernels import _cutile_sinkhorn_fwd - - _clear_autotune_caches() - - s, b, n, iters = 2, 4, 4, 5 - data = _rand(s, b, n, n) - - # First call triggers autotune - _cutile_sinkhorn_fwd(data, iters) - N_batch = s * b - assert (N_batch, n, iters) in mod._sinkhorn_fwd_best_cfg - - # Second call should use cache (no error) - _cutile_sinkhorn_fwd(data, iters) - - -# ============================================================================ -# Proj RMS autotune -# ============================================================================ - - -class TestAutotunesProjRms: - @_require_cutile_experimental - @pytest.mark.parametrize("M,N,K", [(256, 24, 4096), (64, 8, 512)]) - def test_proj_rms_fwd_autotune(self, M, N, K): - """Autotuned proj_rms fwd must match PyTorch reference.""" - from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms - - _clear_autotune_caches() - - eps = 1e-6 - x_data = _rand(M, K) - w_data = _rand(N, K) - - xf = x_data.clone() - wf = w_data.clone() - proj_f, r_f = fused_proj_rms(xf, wf, eps) - - xr = x_data.clone() - wr = w_data.clone() - proj_r, r_r = _ref_proj_rms(xr, wr, eps) - - torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) - torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) - - @_require_cutile_experimental - @pytest.mark.parametrize("M,N,K", [(256, 24, 16384), (64, 8, 512)]) - def test_proj_rms_bwd_autotune(self, M, N, K): - """Autotuned proj_rms bwd grads must match PyTorch reference.""" - from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms - - _clear_autotune_caches() - - eps = 1e-6 - x_data = _rand(M, K) - w_data = _rand(N, K) - grad_proj = _rand(M, N) - grad_r = _rand(M, 1) - - xf = x_data.clone().requires_grad_(True) - wf = w_data.clone().requires_grad_(True) - proj_f, r_f = fused_proj_rms(xf, wf, eps) - (proj_f * grad_proj + r_f * grad_r).sum().backward() - - xr = x_data.clone().requires_grad_(True) - wr = w_data.clone().requires_grad_(True) - proj_r, r_r = _ref_proj_rms(xr, wr, eps) - (proj_r * grad_proj + r_r * grad_r).sum().backward() - - torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) - torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) - torch.testing.assert_close( - xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" - ) - torch.testing.assert_close( - wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" - ) - - @_require_cutile_experimental - def test_proj_rms_autotune_cache_populated(self): - """After first call, autotune caches should contain the best configs.""" - from megatron.core.fusions import fused_mhc_kernels as mod - from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms - - _clear_autotune_caches() - - M, N, K = 256, 24, 16384 - eps = 1e-6 - x = _rand(M, K, requires_grad=True) - w = _rand(N, K, requires_grad=True) - - proj, r = fused_proj_rms(x, w, eps) - assert (M, N, K) in mod._proj_rms_fwd_best_cfg, "fwd cache not populated" - - (proj.sum() + r.sum()).backward() - # bwd autotune only triggers for K >= 8192 - assert (M, N, K) in mod._proj_rms_bwd_best_cfg, "bwd cache not populated" diff --git a/tests/unit_tests/fusions/verify_compute_h_bwd.py b/tests/unit_tests/fusions/verify_compute_h_bwd.py deleted file mode 100644 index 6d29d090c8a..00000000000 --- a/tests/unit_tests/fusions/verify_compute_h_bwd.py +++ /dev/null @@ -1,198 +0,0 @@ -"""Standalone verification of compute_h backward formulas. - -Pure PyTorch — no cuTile needed. Compares hand-derived gradients against autograd. -""" - -import math - -import torch - -torch.manual_seed(42) -DEVICE = "cpu" -DTYPE = torch.float64 # float64 for precise gradient checking - - -def forward_proj_rms_compute_h(x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps): - """Forward: proj_rms + compute_h. All differentiable.""" - proj = x @ weight.t() # [M, N] - norm = x.norm(dim=-1, keepdim=True) # [M, 1] - K = x.shape[-1] - N = proj.shape[-1] - r = norm / math.sqrt(K) # [M, 1] - - # Build alpha vector from 3 scalars - alpha = torch.cat([ - alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), - ]) - - # compute_h - h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) - h_pre = h[..., :n].sigmoid() - h_post = h[..., n:2 * n].sigmoid() * 2 - h_res = h[..., 2 * n:] - y = torch.cat([h_pre, h_post, h_res], dim=-1) - return y, r, proj, norm - - -def manual_backward_compute_h( - grad_y, grad_r_ext, y_activated, proj, r, - alpha_pre, alpha_post, alpha_res, n, eps, -): - """Hand-derived backward through compute_h + r conversion. - - Returns: grad_proj, grad_r_for_bwd, grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias - """ - M, N = proj.shape - - # Build alpha vector - alpha = torch.cat([ - alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), - ]) - - # 1. Activation backward -> grad_h - y_pre = y_activated[:, :n] - y_post = y_activated[:, n:2 * n] - gy_pre = grad_y[:, :n] - gy_post = grad_y[:, n:2 * n] - gy_res = grad_y[:, 2 * n:] - - gh_pre = gy_pre * y_pre * (1.0 - y_pre) - gh_post = gy_post * (y_post / 2.0) * (1.0 - y_post / 2.0) * 2.0 - gh_res = gy_res - - grad_h = torch.cat([gh_pre, gh_post, gh_res], dim=-1) - - # 2. Linear backward: h = proj * alpha / (r + eps) + bias - r_eps = r + eps - alpha_row = alpha.unsqueeze(0) - - grad_proj = grad_h * alpha_row / r_eps - grad_r_from_h = -(grad_h * proj * alpha_row / (r_eps * r_eps)).sum(dim=-1, keepdim=True) - - # Per-column grad_alpha, then reduce to 3 scalars - ga_all = (grad_h * proj / r_eps).sum(dim=0) # [N] - grad_alpha_pre = ga_all[:n].sum() - grad_alpha_post = ga_all[n:2 * n].sum() - grad_alpha_res = ga_all[2 * n:].sum() - grad_bias = grad_h.sum(dim=0) - - # 3. Total grad_r = from compute_h + external - grad_r_total = grad_r_from_h - if grad_r_ext is not None: - grad_r_total = grad_r_total + grad_r_ext - - # 4. Convert grad_r from r_new=norm/sqrt(K) to r_old=1/(r_new+eps) - grad_r_for_bwd = grad_r_total * (-r_eps * r_eps) - - return grad_proj, grad_r_for_bwd, grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias - - -def manual_backward_proj_rms(grad_proj, grad_r_for_bwd, x, weight, norm, eps): - """Hand-derived backward through proj_rms (reference, no kernel). - - grad_r_for_bwd is w.r.t. r_old = 1/(norm/sqrt(K) + eps). - """ - K = x.shape[-1] - inv_sqrt_k = 1.0 / math.sqrt(K) - - # grad through proj = x @ weight.T - grad_x_from_proj = grad_proj @ weight # [M, K] - grad_weight = grad_proj.t() @ x # [N, K] - - # grad through r_old = 1/(norm/sqrt(K) + eps) - # d(r_old)/d(norm) = -1/(norm/sqrt(K)+eps)^2 * (1/sqrt(K)) - # d(norm)/d(x_i) = x_i / norm - u = norm * inv_sqrt_k + eps - coeff = -1.0 / (u * u) * inv_sqrt_k # [M, 1] - inv_norm = torch.where(norm > 0, 1.0 / norm, torch.zeros_like(norm)) - grad_x_from_r = grad_r_for_bwd * coeff * x * inv_norm # [M, K] - - grad_x = grad_x_from_proj + grad_x_from_r - return grad_x, grad_weight - - -def test_backward(M, n, K, eps=1e-6): - """Compare manual backward against autograd.""" - N = n * n + 2 * n - - x = torch.randn(M, K, dtype=DTYPE, device=DEVICE, requires_grad=True) - weight = torch.randn(N, K, dtype=DTYPE, device=DEVICE, requires_grad=True) - alpha_pre = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) - alpha_post = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) - alpha_res = torch.randn(1, dtype=DTYPE, device=DEVICE, requires_grad=True) - bias = torch.randn(N, dtype=DTYPE, device=DEVICE, requires_grad=True) - - grad_y = torch.randn(M, N, dtype=DTYPE, device=DEVICE) - grad_r_ext = torch.randn(M, 1, dtype=DTYPE, device=DEVICE) - - # ---- Autograd path ---- - y_auto, r_auto, _, _ = forward_proj_rms_compute_h( - x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, - ) - loss = (y_auto * grad_y).sum() + (r_auto * grad_r_ext).sum() - loss.backward() - - auto_grad_x = x.grad.clone() - auto_grad_w = weight.grad.clone() - auto_grad_ap = alpha_pre.grad.clone() - auto_grad_apo = alpha_post.grad.clone() - auto_grad_ar = alpha_res.grad.clone() - auto_grad_b = bias.grad.clone() - - # ---- Manual path ---- - x2 = x.detach().clone() - w2 = weight.detach().clone() - ap2 = alpha_pre.detach().clone() - apo2 = alpha_post.detach().clone() - ar2 = alpha_res.detach().clone() - b2 = bias.detach().clone() - - y_man, r_man, proj_man, norm_man = forward_proj_rms_compute_h( - x2, w2, ap2, apo2, ar2, b2, n, eps, - ) - - # Step 1: manual compute_h backward - ( - grad_proj, grad_r_for_bwd, - grad_alpha_pre, grad_alpha_post, grad_alpha_res, - grad_bias, - ) = manual_backward_compute_h( - grad_y, grad_r_ext, y_man, proj_man, r_man, ap2, apo2, ar2, n, eps, - ) - - # Step 2: manual proj_rms backward - grad_x_man, grad_w_man = manual_backward_proj_rms( - grad_proj, grad_r_for_bwd, x2, w2, norm_man, eps, - ) - - # ---- Compare ---- - def check(name, manual, auto, atol=1e-10, rtol=1e-7): - diff = (manual - auto).abs().max().item() - rel = diff / (auto.abs().max().item() + 1e-15) - ok = diff < atol or rel < rtol - status = "PASS" if ok else "FAIL" - print(f" {status}: {name:14s} max_abs_diff={diff:.3e} rel={rel:.3e}") - return ok - - print(f"\n--- M={M}, n={n}, K={K}, N={N} ---") - all_ok = True - all_ok &= check("grad_x", grad_x_man, auto_grad_x) - all_ok &= check("grad_weight", grad_w_man, auto_grad_w) - all_ok &= check("grad_alpha_pre", grad_alpha_pre, auto_grad_ap) - all_ok &= check("grad_alpha_post", grad_alpha_post, auto_grad_apo) - all_ok &= check("grad_alpha_res", grad_alpha_res, auto_grad_ar) - all_ok &= check("grad_bias", grad_bias, auto_grad_b) - return all_ok - - -if __name__ == "__main__": - print("Verifying manual backward formulas against autograd (float64)...") - results = [] - for M, n, K in [(32, 2, 64), (64, 4, 256), (16, 3, 128), (128, 4, 512)]: - results.append(test_backward(M, n, K)) - - print("\n" + "=" * 50) - if all(results): - print("ALL PASSED") - else: - print("SOME FAILED") diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index c5f6ced37bb..d8f54455e26 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -406,438 +406,5 @@ def test_config_enable_mhc_recompute(self): assert config.enable_hyper_connections is True -class TestHyperConnectionBroadcastCheckpoint: - """Test HyperConnectionModule checkpoint with BroadcastTensorFused. - - These tests mirror TestHyperConnectionCheckpoint but use the returned - residual (BroadcastTensor view) from forward() for apply_h_res, exercising - all 3 gradient paths through a single hidden_states tensor. - """ - - def setup_method(self, method): - Utils.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(123) - - def teardown_method(self, method): - Utils.destroy_model_parallel() - - def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4): - config = TransformerConfig( - num_layers=2, - hidden_size=hidden_size, - num_attention_heads=4, - use_cpu_initialization=True, - enable_hyper_connections=True, - num_residual_streams=num_residual_streams, - mhc_sinkhorn_iterations=5, - mhc_init_gating_factor=0.01, - ) - module = HyperConnectionModule(config=config, layer_number=1) - module.cuda() - return module - - def test_broadcast_normal_vs_checkpoint_correctness(self): - """ - _forward_normal vs _forward_with_checkpoint using BroadcastTensor residual. - - All 3 gradient paths (compute_mappings, aggregate, apply_h_res) flow - through a single hidden_states tensor via BroadcastTensorFused. - """ - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - - module = self._create_hyper_connection_module(hidden_size, num_streams) - - hidden_states = torch.randn( - seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True - ) - hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) - - # Reference: normal forward, use returned residual - torch.manual_seed(42) - torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal( - hidden_states - ) - mixed_ref = module.apply_h_res(h_res_ref, residual_ref) - loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() - loss_ref.backward() - grad_ref = hidden_states.grad.clone() - - # Checkpoint: use returned residual - torch.manual_seed(42) - torch.cuda.manual_seed(42) - manager = CheckpointManager() - aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt = ( - module._forward_with_checkpoint(hidden_states_ckpt, manager) - ) - mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) - loss_ckpt = aggregated_ckpt.sum() + mixed_ckpt.sum() + h_post_ckpt.sum() - manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) - loss_ckpt.backward() - grad_ckpt = hidden_states_ckpt.grad.clone() - - assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( - f"Broadcast checkpoint gradients mismatch" - ) - - def test_broadcast_forward_with_manager(self): - """forward() with/without manager using BroadcastTensor residual.""" - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - - module = self._create_hyper_connection_module(hidden_size, num_streams) - - hidden_states = torch.randn( - seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True - ) - hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) - - # Reference - torch.manual_seed(42) - torch.cuda.manual_seed(42) - agg_ref, h_res_ref, h_post_ref, res_ref = module.forward( - hidden_states, mhc_recompute_manager=None - ) - loss_ref = agg_ref.sum() + module.apply_h_res(h_res_ref, res_ref).sum() + h_post_ref.sum() - loss_ref.backward() - grad_ref = hidden_states.grad.clone() - - # With manager - torch.manual_seed(42) - torch.cuda.manual_seed(42) - manager = CheckpointManager() - agg_ckpt, h_res_ckpt, h_post_ckpt, res_ckpt = module.forward( - hidden_states_ckpt, mhc_recompute_manager=manager - ) - loss_ckpt = ( - agg_ckpt.sum() + module.apply_h_res(h_res_ckpt, res_ckpt).sum() + h_post_ckpt.sum() - ) - manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) - loss_ckpt.backward() - grad_ckpt = hidden_states_ckpt.grad.clone() - - assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5) - - def test_broadcast_fused_h_res_h_post_bda(self): - """ - Full TransformerLayer-like flow: forward() -> simulated layer -> - fused_h_res_h_post_bda() using BroadcastTensor residual. - """ - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - - module = self._create_hyper_connection_module(hidden_size, num_streams) - - hidden_states = torch.randn( - seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True - ) - hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) - - # Reference - torch.manual_seed(42) - torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref, residual_ref = module.forward( - hidden_states, mhc_recompute_manager=None - ) - layer_output = aggregated_ref * 0.1 # simulated layer - output_ref = module.fused_h_res_h_post_bda( - h_res_ref, residual_ref, h_post_ref, - (layer_output, None), 0.0, True, False, - ) - loss_ref = output_ref.sum() - loss_ref.backward() - grad_ref = hidden_states.grad.clone() - - # With checkpoint manager - torch.manual_seed(42) - torch.cuda.manual_seed(42) - manager = CheckpointManager() - aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt = module.forward( - hidden_states_ckpt, mhc_recompute_manager=manager - ) - layer_output_ckpt = aggregated_ckpt * 0.1 - output_ckpt = module.fused_h_res_h_post_bda( - h_res_ckpt, residual_ckpt, h_post_ckpt, - (layer_output_ckpt, None), 0.0, True, False, manager, - ) - loss_ckpt = output_ckpt.sum() - manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) - loss_ckpt.backward() - grad_ckpt = hidden_states_ckpt.grad.clone() - - assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( - f"Broadcast fused_h_res_h_post_bda gradients mismatch" - ) - - -class TestMHCBroadcastChainIntegration: - """Test multiple chained HyperConnectionModules with BroadcastTensor residual.""" - - def setup_method(self, method): - Utils.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(123) - - def teardown_method(self, method): - Utils.destroy_model_parallel() - - def test_chained_broadcast_fwd_bwd(self): - """ - Chain 3 HyperConnectionModules using BroadcastTensor residual throughout. - Compare ref (no ckpt) vs checkpoint path. - """ - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - n_channels = num_streams * hidden_size - - config = TransformerConfig( - num_layers=4, - hidden_size=hidden_size, - num_attention_heads=4, - use_cpu_initialization=True, - enable_hyper_connections=True, - num_residual_streams=num_streams, - mhc_sinkhorn_iterations=5, - mhc_init_gating_factor=0.01, - ) - modules = [ - HyperConnectionModule(config=config, layer_number=i + 1).cuda() for i in range(3) - ] - - hidden_states_ref = torch.randn( - seq_len, batch_size, n_channels, device='cuda', requires_grad=True - ) - hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) - - # Reference: no checkpoint, use BroadcastTensor residual - torch.manual_seed(42) - torch.cuda.manual_seed(42) - - h = hidden_states_ref - for module in modules: - agg, h_res, h_post, res = module.forward(h, mhc_recompute_manager=None) - layer_out = agg * 0.1 - h = module.fused_h_res_h_post_bda( - h_res, res, h_post, (layer_out, None), 0.0, True, False, - ) - - loss_ref = h.sum() - loss_ref.backward() - grad_ref = hidden_states_ref.grad.clone() - - # With checkpoint - torch.manual_seed(42) - torch.cuda.manual_seed(42) - - manager = CheckpointManager() - h = hidden_states_ckpt - for module in modules: - agg, h_res, h_post, res = module.forward(h, mhc_recompute_manager=manager) - layer_out = agg * 0.1 - h = module.fused_h_res_h_post_bda( - h_res, res, h_post, (layer_out, None), 0.0, True, False, manager, - ) - - loss_ckpt = h.sum() - manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) - loss_ckpt.backward() - grad_ckpt = hidden_states_ckpt.grad.clone() - - assert torch.allclose(grad_ckpt, grad_ref, atol=1e-4), ( - f"Chained broadcast HyperConnection gradients mismatch" - ) - - def test_partial_checkpoint_broadcast(self): - """ - Last layer not checkpointed, all using BroadcastTensor residual. - """ - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - - config = TransformerConfig( - num_layers=2, - hidden_size=hidden_size, - num_attention_heads=4, - use_cpu_initialization=True, - enable_hyper_connections=True, - num_residual_streams=num_streams, - mhc_sinkhorn_iterations=5, - mhc_init_gating_factor=0.01, - ) - module = HyperConnectionModule(config=config, layer_number=1).cuda() - - hidden_states_ref = torch.randn( - seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True - ) - hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) - - # Reference - torch.manual_seed(42) - torch.cuda.manual_seed(42) - agg_ref, h_res_ref, h_post_ref, res_ref = module.forward( - hidden_states_ref, mhc_recompute_manager=None - ) - layer_out_ref = agg_ref * 0.1 - output_ref = module.fused_h_res_h_post_bda( - h_res_ref, res_ref, h_post_ref, (layer_out_ref, None), 0.0, True, False, - ) - loss_ref = output_ref.sum() - loss_ref.backward() - grad_ref = hidden_states_ref.grad.clone() - - # With manager, simulate last layer not checkpointed - torch.manual_seed(42) - torch.cuda.manual_seed(42) - manager = CheckpointManager() - agg_ckpt, h_res_ckpt, h_post_ckpt, res_ckpt = module.forward( - hidden_states_ckpt, mhc_recompute_manager=manager - ) - layer_out_ckpt = agg_ckpt * 0.1 - # Last layer: no manager for BDA - output_ckpt = module.fused_h_res_h_post_bda( - h_res_ckpt, res_ckpt, h_post_ckpt, (layer_out_ckpt, None), 0.0, True, False, - ) - manager.discard_all_outputs_and_register_unified_recompute(output_ckpt) - loss_ckpt = output_ckpt.sum() - loss_ckpt.backward() - grad_ckpt = hidden_states_ckpt.grad.clone() - - assert torch.allclose(grad_ckpt, grad_ref, atol=1e-5), ( - f"Partial checkpoint broadcast gradients mismatch" - ) - - -class TestBroadcastTensorFused: - """Test BroadcastTensorFused autograd.Function and native_fused_add_3.""" - - def test_native_fused_add_3_correctness(self): - """Test native_fused_add_3 matches plain addition.""" - from megatron.core.transformer.hyper_connection import native_fused_add_3 - - a = torch.randn(8, 16, device='cuda') - b = torch.randn(8, 16, device='cuda') - c = torch.randn(8, 16, device='cuda') - result = native_fused_add_3(a, b, c) - expected = a + b + c - assert torch.allclose(result, expected, atol=1e-6) - - def test_native_fused_add_3_gradient(self): - """Test native_fused_add_3 produces correct gradients.""" - from megatron.core.transformer.hyper_connection import native_fused_add_3 - - a = torch.randn(4, 8, device='cuda', requires_grad=True) - b = torch.randn(4, 8, device='cuda', requires_grad=True) - c = torch.randn(4, 8, device='cuda', requires_grad=True) - result = native_fused_add_3(a, b, c) - result.sum().backward() - assert torch.allclose(a.grad, torch.ones_like(a)) - assert torch.allclose(b.grad, torch.ones_like(b)) - assert torch.allclose(c.grad, torch.ones_like(c)) - - def test_broadcast_tensor_returns_views(self): - """Test BroadcastTensorFused forward returns 3 views sharing storage.""" - from megatron.core.transformer.hyper_connection import ( - BroadcastTensorFused, - native_fused_add_3, - ) - - x = torch.randn(4, 8, device='cuda', requires_grad=True) - a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) - assert a.data_ptr() == x.data_ptr() - assert b.data_ptr() == x.data_ptr() - assert c.data_ptr() == x.data_ptr() - - def test_broadcast_tensor_gradient_correctness(self): - """Test BroadcastTensorFused backward sums gradients correctly.""" - from megatron.core.transformer.hyper_connection import ( - BroadcastTensorFused, - native_fused_add_3, - ) - - x = torch.randn(4, 8, device='cuda', requires_grad=True) - a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) - # Each branch applies a different scalar: loss = a.sum() + 2*b.sum() + 3*c.sum() - loss = a.sum() + 2 * b.sum() + 3 * c.sum() - loss.backward() - # x.grad should be (1 + 2 + 3) * ones - expected = torch.full_like(x, 6.0) - assert torch.allclose(x.grad, expected, atol=1e-6), ( - f"Expected grad {expected}, got {x.grad}" - ) - - def test_broadcast_tensor_partial_none_grads(self): - """Test BroadcastTensorFused handles None gradients from unused branches.""" - from megatron.core.transformer.hyper_connection import ( - BroadcastTensorFused, - native_fused_add_3, - ) - - x = torch.randn(4, 8, device='cuda', requires_grad=True) - a, b, c = BroadcastTensorFused.apply(x, native_fused_add_3) - # Only use one branch - loss = a.sum() - loss.backward() - expected = torch.ones_like(x) - assert torch.allclose(x.grad, expected, atol=1e-6) - - def test_broadcast_tensor_end_to_end_with_hyper_connection(self): - """ - End-to-end test: HyperConnectionModule.forward returns 4-tuple and - gradients through all 3 paths (compute_mappings, aggregate, - fused_h_res_h_post_bda) are correct. - """ - Utils.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(123) - try: - hidden_size = 64 - num_streams = 4 - seq_len = 8 - batch_size = 2 - - config = TransformerConfig( - num_layers=2, - hidden_size=hidden_size, - num_attention_heads=4, - use_cpu_initialization=True, - enable_hyper_connections=True, - num_residual_streams=num_streams, - mhc_sinkhorn_iterations=5, - mhc_init_gating_factor=0.01, - ) - module = HyperConnectionModule(config=config, layer_number=1).cuda() - - hidden_states = torch.randn( - seq_len, batch_size, num_streams * hidden_size, - device='cuda', requires_grad=True, - ) - - torch.manual_seed(42) - torch.cuda.manual_seed(42) - - # Forward returns 4-tuple now - aggregated, h_res, h_post, residual = module.forward(hidden_states) - - # Simulate full flow: aggregate + apply_h_res on residual - mixed = module.apply_h_res(h_res, residual) - loss = aggregated.sum() + mixed.sum() + h_post.sum() - loss.backward() - - # Verify gradient exists and is finite - assert hidden_states.grad is not None - assert torch.isfinite(hidden_states.grad).all() - finally: - Utils.destroy_model_parallel() - - if __name__ == "__main__": pytest.main([__file__, "-v"]) From 1a1c71cd1c0d0fe52fd2642cc4caaeaf89dea4f4 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 12 May 2026 10:56:46 +0800 Subject: [PATCH 08/34] Improve cuTile mHC reduce compute_h tiling --- megatron/core/fusions/fused_mhc_kernels.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index ef62c4ebb02..784225ddb6e 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1598,6 +1598,19 @@ def _reduce_compute_h_autotune_configs(M): if tile_m <= M and M % tile_m == 0: yield tile_m + def _default_reduce_compute_h_tile_m(M: int) -> int: + """Pick a reduce tile size with enough blocks to cover the GPU.""" + try: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + except Exception: + num_sms = 128 + + valid = [tm for tm in _reduce_compute_h_autotune_configs(M)] + for tm in valid: + if math.ceil(M / tm) >= num_sms: + return tm + return valid[-1] if valid else 1 + _reduce_compute_h_best_cfg: dict = {} def _cutile_reduce_compute_h( @@ -1612,7 +1625,7 @@ def _cutile_reduce_compute_h( N: int, K: int, eps: float, - tile_m: int, + _proj_tile_m: int, tile_n: int, split_k: int, ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: @@ -1636,7 +1649,8 @@ def _cutile_reduce_compute_h( r_out = torch.empty(M, 1, dtype=proj_acc.dtype, device=dev) proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) - cache_key = (M, N, K, n, tile_m, split_k) + default_tm = _default_reduce_compute_h_tile_m(M) + cache_key = (M, N, K, n, split_k) cached = _reduce_compute_h_best_cfg.get(cache_key) def _make_args(tm): @@ -1649,7 +1663,7 @@ def _make_args(tm): ) if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: - tm = cached if cached is not None else tile_m + tm = cached if cached is not None else default_tm ct.launch(stream, (math.ceil(M / tm),), _ct_reduce_compute_h_kernel, _make_args(tm)) else: from types import SimpleNamespace @@ -1657,7 +1671,6 @@ def _make_args(tm): configs = [ SimpleNamespace(TILE_M=tm) for tm in _reduce_compute_h_autotune_configs(M) - if tm == tile_m ] tuned = ct_experimental.autotune_launch( stream, From f33d1c4ef1c75b361a64cb12f925817273540579 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Wed, 13 May 2026 00:26:29 -0700 Subject: [PATCH 09/34] Fix: autoformat --- megatron/core/fusions/fused_mhc_kernels.py | 977 +++++++++++------- megatron/core/transformer/hyper_connection.py | 15 +- .../fusions/test_fused_mhc_kernels.py | 124 ++- .../test_hyper_connection_recompute.py | 4 +- 4 files changed, 712 insertions(+), 408 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 784225ddb6e..81e08d86b10 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -71,16 +71,12 @@ def is_triton_available() -> bool: # Sinkhorn-Knopp # ============================================================================ - @triton.autotune( - configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], - key=["HC", "NUM_ITERS"], + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] ) @triton.jit def _triton_sinkhorn_fwd_kernel( - inp_ptr, out_ptr, M_init_ptr, - N_batch, eps, - HC: tl.constexpr, NUM_ITERS: tl.constexpr, + inp_ptr, out_ptr, M_init_ptr, N_batch, eps, HC: tl.constexpr, NUM_ITERS: tl.constexpr ): """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" pid = tl.program_id(0) @@ -105,17 +101,21 @@ def _triton_sinkhorn_fwd_kernel( tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) - @triton.autotune( - configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], - key=["HC", "NUM_ITERS"], + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] ) @triton.jit def _triton_sinkhorn_bwd_kernel( - grad_out_ptr, M_init_ptr, grad_inp_ptr, - ws_M_ptr, ws_rs_ptr, ws_cs_ptr, - N_batch, eps, - HC: tl.constexpr, NUM_ITERS: tl.constexpr, + grad_out_ptr, + M_init_ptr, + grad_inp_ptr, + ws_M_ptr, + ws_rs_ptr, + ws_cs_ptr, + N_batch, + eps, + HC: tl.constexpr, + NUM_ITERS: tl.constexpr, ): """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" pid = tl.program_id(0) @@ -156,8 +156,11 @@ def _triton_sinkhorn_bwd_kernel( col_corr = tl.sum(grad * M, axis=0) grad = grad - col_corr[None, :] M = tl.load( - ws_M_ptr + M_ws_base + (2 * t + 1) * HC * HC - + offs_r[:, None] * HC + offs_c[None, :] + ws_M_ptr + + M_ws_base + + (2 * t + 1) * HC * HC + + offs_r[:, None] * HC + + offs_c[None, :] ).to(tl.float32) row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) @@ -165,15 +168,13 @@ def _triton_sinkhorn_bwd_kernel( row_corr = tl.sum(grad * M, axis=1) grad = grad - row_corr[:, None] M = tl.load( - ws_M_ptr + M_ws_base + (2 * t) * HC * HC - + offs_r[:, None] * HC + offs_c[None, :] + ws_M_ptr + M_ws_base + (2 * t) * HC * HC + offs_r[:, None] * HC + offs_c[None, :] ).to(tl.float32) M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) grad = grad * M_init tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) - def _triton_sinkhorn_fwd( input_logits: Tensor, num_iterations: int, eps: float = 1e-8 ) -> Tuple[Tensor, Tensor]: @@ -184,12 +185,9 @@ def _triton_sinkhorn_fwd( out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) inp = input_logits.contiguous().view(N_batch, hc, hc) - _triton_sinkhorn_fwd_kernel[(N_batch,)]( - inp, out, M_init, N_batch, eps, hc, num_iterations, - ) + _triton_sinkhorn_fwd_kernel[(N_batch,)](inp, out, M_init, N_batch, eps, hc, num_iterations) return out.view(original_shape), M_init.view(original_shape) - def _triton_sinkhorn_bwd( grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 ) -> Tensor: @@ -204,12 +202,13 @@ def _triton_sinkhorn_bwd( ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) _triton_sinkhorn_bwd_kernel[(N_batch,)]( - go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations, + go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations ) return grad_input.view(original_shape) - class TritonFusedSinkhorn(torch.autograd.Function): + """Autograd wrapper for Triton fused Sinkhorn.""" + @staticmethod def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) @@ -224,18 +223,16 @@ def backward(ctx, grad_output: Tensor): grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None - def triton_fused_sinkhorn( input_logits: Tensor, num_iterations: int, eps: float = 1e-6 ) -> Tensor: + """Apply Triton fused Sinkhorn with autograd support.""" return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) - # ============================================================================ # H_aggregate forward # ============================================================================ - @triton.autotune( configs=[ triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) @@ -247,10 +244,17 @@ def triton_fused_sinkhorn( ) @triton.jit def _triton_h_agg_fwd_kernel( - x_ptr, h_ptr, out_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_x_s, stride_x_n, stride_x_c, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + x_ptr, + h_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_x_s, + stride_x_n, + stride_x_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, ): """out[s, c] = sum_i x[s, i, c] * h[s, i].""" pid_s = tl.program_id(0) @@ -276,7 +280,6 @@ def _triton_h_agg_fwd_kernel( mask=mask_2d, ) - def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: s, b, n, C = x.shape sb = s * b @@ -286,17 +289,14 @@ def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) _triton_h_agg_fwd_kernel[grid]( - x_flat, h_flat, out, sb, C, n, - x_flat.stride(0), x_flat.stride(1), x_flat.stride(2), + x_flat, h_flat, out, sb, C, n, x_flat.stride(0), x_flat.stride(1), x_flat.stride(2) ) return out.view(s, b, C) - # ============================================================================ # H_post BDA # ============================================================================ - @triton.autotune( configs=[ triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) @@ -308,13 +308,27 @@ def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: ) @triton.jit def _triton_hpb_fwd_kernel( - hr_ptr, orig_ptr, hp_ptr, x_ptr, bias_ptr, out_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_hr_s, stride_hr_i, stride_hr_j, - stride_orig_s, stride_orig_n, stride_orig_c, - stride_out_s, stride_out_n, stride_out_c, + hr_ptr, + orig_ptr, + hp_ptr, + x_ptr, + bias_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_out_s, + stride_out_n, + stride_out_c, HAS_BIAS: tl.constexpr, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, ): """out = hr @ orig + hp * (x + bias).""" pid_s = tl.program_id(0) @@ -325,11 +339,9 @@ def _triton_hpb_fwd_kernel( mask_c = offs_c < C mask_2d = mask_s[:, None] & mask_c[None, :] - x_tile = tl.load( - x_ptr + offs_s[:, None] * C + offs_c[None, :], - mask=mask_2d, - other=0.0, - ).to(tl.float32) + x_tile = tl.load(x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0).to( + tl.float32 + ) if HAS_BIAS: bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) x_tile += bias_tile[None, :] @@ -360,13 +372,8 @@ def _triton_hpb_fwd_kernel( mask=mask_2d, ) - def _triton_h_post_bda_fwd( - h_res: Tensor, - original_residual: Tensor, - h_post: Tensor, - x: Tensor, - bias: Optional[Tensor], + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: s, b, n, C = original_residual.shape sb = s * b @@ -379,16 +386,28 @@ def _triton_h_post_bda_fwd( grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) _triton_hpb_fwd_kernel[grid]( - hr_flat, orig_flat, hp_flat, x_flat, bias if bias is not None else x_flat, out, - sb, C, n, - hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), - orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), - out.stride(0), out.stride(1), out.stride(2), + hr_flat, + orig_flat, + hp_flat, + x_flat, + bias if bias is not None else x_flat, + out, + sb, + C, + n, + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), HAS_BIAS=(bias is not None), ) return out.view(s, b, n, C) - @triton.autotune( configs=[ triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) @@ -400,12 +419,25 @@ def _triton_h_post_bda_fwd( ) @triton.jit def _triton_hpb_bwd_g_x_orig_kernel( - go_ptr, hr_ptr, hp_ptr, g_orig_ptr, g_x_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_go_s, stride_go_n, stride_go_c, - stride_hr_s, stride_hr_i, stride_hr_j, - stride_orig_s, stride_orig_n, stride_orig_c, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + go_ptr, + hr_ptr, + hp_ptr, + g_orig_ptr, + g_x_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, ): """g_x = hp @ go, g_orig = hr.T @ go.""" pid_s = tl.program_id(0) @@ -420,7 +452,8 @@ def _triton_hpb_bwd_g_x_orig_kernel( for j in tl.static_range(N): go_j = tl.load( go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], - mask=mask_2d, other=0.0, + mask=mask_2d, + other=0.0, ).to(tl.float32) hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) g_x_acc += hp_j[:, None] * go_j @@ -435,11 +468,13 @@ def _triton_hpb_bwd_g_x_orig_kernel( for j in tl.static_range(N): go_j = tl.load( go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], - mask=mask_2d, other=0.0, + mask=mask_2d, + other=0.0, ).to(tl.float32) hr_ji = tl.load( hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, - mask=mask_s, other=0.0, + mask=mask_s, + other=0.0, ).to(tl.float32) g_orig_i += hr_ji[:, None] * go_j tl.store( @@ -448,7 +483,6 @@ def _triton_hpb_bwd_g_x_orig_kernel( mask=mask_2d, ) - @triton.autotune( configs=[ triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) @@ -460,14 +494,27 @@ def _triton_hpb_bwd_g_x_orig_kernel( ) @triton.jit def _triton_hpb_bwd_g_hp_hr_kernel( - go_ptr, orig_ptr, x_ptr, bias_ptr, - g_hr_ptr, g_hp_ptr, - sb, C: tl.constexpr, N: tl.constexpr, - stride_go_s, stride_go_n, stride_go_c, - stride_orig_s, stride_orig_n, stride_orig_c, - stride_hr_s, stride_hr_i, stride_hr_j, + go_ptr, + orig_ptr, + x_ptr, + bias_ptr, + g_hr_ptr, + g_hp_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, HAS_BIAS: tl.constexpr, - BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, ): """g_hp = sum_c go*(x+bias), g_hr = go @ orig.T.""" pid_s = tl.program_id(0) @@ -483,9 +530,7 @@ def _triton_hpb_bwd_g_hp_hr_kernel( mask_2d = mask_s[:, None] & mask_c[None, :] x_tile = tl.load( - x_ptr + offs_s[:, None] * C + offs_c[None, :], - mask=mask_2d, - other=0.0, + x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0 ).to(tl.float32) if HAS_BIAS: bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) @@ -537,7 +582,6 @@ def _triton_hpb_bwd_g_hp_hr_kernel( mask=mask_s, ) - def _triton_h_post_bda_bwd( grad_output: Tensor, h_res: Tensor, @@ -563,19 +607,45 @@ def _triton_h_post_bda_bwd( grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) _triton_hpb_bwd_g_x_orig_kernel[grid_a]( - go_flat, hr_flat, hp_flat, g_res, g_x, sb, C, n, - go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), - hr_flat.stride(0), hr_flat.stride(1), hr_flat.stride(2), - g_res.stride(0), g_res.stride(1), g_res.stride(2), + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + g_res.stride(0), + g_res.stride(1), + g_res.stride(2), ) grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( - go_flat, orig_flat, x_flat, bias if bias is not None else x_flat, - g_hr, g_hp, sb, C, n, - go_flat.stride(0), go_flat.stride(1), go_flat.stride(2), - orig_flat.stride(0), orig_flat.stride(1), orig_flat.stride(2), - g_hr.stride(0), g_hr.stride(1), g_hr.stride(2), + go_flat, + orig_flat, + x_flat, + bias if bias is not None else x_flat, + g_hr, + g_hp, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + g_hr.stride(0), + g_hr.stride(1), + g_hr.stride(2), HAS_BIAS=(bias is not None), ) @@ -700,16 +770,13 @@ def _cutile_sinkhorn_fwd( from types import SimpleNamespace configs = [ - SimpleNamespace(TILE_SIZE=ts) - for ts in _sinkhorn_autotune_tile_sizes(N_batch) + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) ] tuned = ct_experimental.autotune_launch( stream, grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), kernel=_ct_sinkhorn_fwd_kernel, - args_fn=lambda cfg: ( - inp, out, M_init, eps, hc, num_iterations, cfg.TILE_SIZE, - ), + args_fn=lambda cfg: (inp, out, M_init, eps, hc, num_iterations, cfg.TILE_SIZE), search_space=configs, ) best_ts = tuned.tuned_config.TILE_SIZE @@ -739,7 +806,9 @@ def _cutile_sinkhorn_bwd( cached = _sinkhorn_bwd_best_cfg.get(cache_key) def _alloc_and_launch(ts): - ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) ct.launch( @@ -756,11 +825,12 @@ def _alloc_and_launch(ts): from types import SimpleNamespace configs = [ - SimpleNamespace(TILE_SIZE=ts) - for ts in _sinkhorn_autotune_tile_sizes(N_batch) + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) ] # Allocate workspace for largest tile size (all configs share same workspace shape). - ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) tuned = ct_experimental.autotune_launch( @@ -768,7 +838,15 @@ def _alloc_and_launch(ts): grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), kernel=_ct_sinkhorn_bwd_kernel, args_fn=lambda cfg: ( - go, mi, grad_input, ws_M, ws_rs, ws_cs, eps, hc, num_iterations, + go, + mi, + grad_input, + ws_M, + ws_rs, + ws_cs, + eps, + hc, + num_iterations, cfg.TILE_SIZE, ), search_space=configs, @@ -839,8 +917,9 @@ def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: tm, tc = cached else: tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) - ct.launch(stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, - (x_flat, h_flat, out, n, tm, tc)) + ct.launch( + stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, (x_flat, h_flat, out, n, tm, tc) + ) return out.view(s, b, C) @@ -864,8 +943,12 @@ def _cutile_h_aggregate_bwd( tm, tc = cached else: tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) - ct.launch(stream, (math.ceil(sb / tm),), _ct_h_agg_bwd_kernel, - (go_flat, x_flat, h_flat, gx, gh, n, tm, tc)) + ct.launch( + stream, + (math.ceil(sb / tm),), + _ct_h_agg_bwd_kernel, + (go_flat, x_flat, h_flat, gx, gh, n, tm, tc), + ) return gx.view(s, b, n, C), gh.view(s, b, n) @@ -922,9 +1005,7 @@ def _ct_hpb_fwd_bias_kernel( @ct.kernel def _ct_hpb_bwd_g_x_orig_kernel( - go, hr, hp, - g_orig, g_x, - N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, + go, hr, hp, g_orig, g_x, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): """Compute g_x = hp @ go and g_orig = hr.T @ go. @@ -948,16 +1029,15 @@ def _ct_hpb_bwd_g_x_orig_kernel( hr_row_j = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) g_orig_tile = g_orig_tile + ct.reshape(hr_row_j, (TILE_SIZE, N, 1)) * go_j ct.store( - g_x, index=(pid, ct_idx), + g_x, + index=(pid, ct_idx), tile=ct.reshape(g_x_tile, (TILE_SIZE, TILE_C)).astype(g_x.dtype), ) ct.store(g_orig, index=(pid, 0, ct_idx), tile=g_orig_tile.astype(g_orig.dtype)) @ct.kernel def _ct_hpb_bwd_g_hp_hr_kernel( - go, orig, x, - g_hr, g_hp, - N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, + go, orig, x, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): """Compute g_hp = sum(go * x) and g_hr = go @ orig.T (no bias). @@ -968,7 +1048,9 @@ def _ct_hpb_bwd_g_hp_hr_kernel( acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) for ct_idx in range(num_c_tiles): - x_tile = ct.load(x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) x_exp = ct.expand_dims(x_tile, axis=1) # [TS, 1, TC] go_tile = ct.load( go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO @@ -980,17 +1062,12 @@ def _ct_hpb_bwd_g_hp_hr_kernel( acc_g_hr = acc_g_hr + ct.sum( ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 ) - ct.store( - g_hp, index=(pid, 0), - tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype), - ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) @ct.kernel def _ct_hpb_bwd_g_hp_hr_bias_kernel( - go, orig, x, bias, - g_hr, g_hp, - N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt, + go, orig, x, bias, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): """Compute g_hp = sum(go * (x+bias)) and g_hr = go @ orig.T (with bias). @@ -1001,7 +1078,9 @@ def _ct_hpb_bwd_g_hp_hr_bias_kernel( acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) for ct_idx in range(num_c_tiles): - x_tile = ct.load(x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # [TS, 1, TC] go_tile = ct.load( @@ -1014,10 +1093,7 @@ def _ct_hpb_bwd_g_hp_hr_bias_kernel( acc_g_hr = acc_g_hr + ct.sum( ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 ) - ct.store( - g_hp, index=(pid, 0), - tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype), - ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) # -- H_post BDA autotune configs & caches -------------------------------- @@ -1093,7 +1169,8 @@ def _cutile_h_post_bda_bwd( else: ts, tc = 1, math.gcd(C, 1024) ct.launch( - stream, (math.ceil(sb / ts), math.ceil(C / tc)), + stream, + (math.ceil(sb / ts), math.ceil(C / tc)), _ct_hpb_bwd_g_x_orig_kernel, (go_flat, hr_flat, hp_flat, g_res, g_x, n, tc, ts), ) @@ -1106,14 +1183,22 @@ def _cutile_h_post_bda_bwd( grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE), math.ceil(C / cfg.TILE_C)), kernel=_ct_hpb_bwd_g_x_orig_kernel, args_fn=lambda cfg: ( - go_flat, hr_flat, hp_flat, g_res, g_x, n, cfg.TILE_C, cfg.TILE_SIZE, + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + n, + cfg.TILE_C, + cfg.TILE_SIZE, ), search_space=configs, ) best = tuned.tuned_config _hpb_bwd_g_x_orig_best_cfg[cache_key_a] = (best.TILE_SIZE, best.TILE_C) ct.launch( - stream, (math.ceil(sb / best.TILE_SIZE), math.ceil(C / best.TILE_C)), + stream, + (math.ceil(sb / best.TILE_SIZE), math.ceil(C / best.TILE_C)), _ct_hpb_bwd_g_x_orig_kernel, (go_flat, hr_flat, hp_flat, g_res, g_x, n, best.TILE_C, best.TILE_SIZE), ) @@ -1121,7 +1206,9 @@ def _cutile_h_post_bda_bwd( # --- Kernel B: g_hp, g_hr (1D grid, loops C-tiles) --- cache_key_b = ('hpb_bwd_g_hp_hr', sb, n, C, bias is not None) cached_b = _hpb_bwd_g_hp_hr_best_cfg.get(cache_key_b) - hp_hr_kernel = _ct_hpb_bwd_g_hp_hr_bias_kernel if bias is not None else _ct_hpb_bwd_g_hp_hr_kernel + hp_hr_kernel = ( + _ct_hpb_bwd_g_hp_hr_bias_kernel if bias is not None else _ct_hpb_bwd_g_hp_hr_kernel + ) if cached_b is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: if cached_b is not None: @@ -1195,7 +1282,8 @@ def _ct_proj_rms_fwd_kernel( SPLIT_K: ConstInt, ): ''' - Grid: (num_tiles_m, num_tiles_k). Fused matmul + norm + r: proj, norm, r in one pass over K. + Grid: (num_tiles_m, num_tiles_k). + Fused matmul + norm + r: proj, norm, r in one pass over K. ''' tile_m_id = ct.bid(0) split_k_id = ct.bid(1) @@ -1272,19 +1360,16 @@ def _ct_reduce_compute_h_kernel( for split_idx in ct.static_iter(range(SPLIT_K)): bid_m_k = bid_m + split_idx * num_bid_m pre_tile = ct.load( - Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) post_tile = ct.load( - Y_acc, index=(bid_m_k, 1), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + Y_acc, index=(bid_m_k, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) pre_accum = pre_accum + ct.astype(pre_tile, ct.float32) post_accum = post_accum + ct.astype(post_tile, ct.float32) r_tile = ct.load( - R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, + R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO ) r_accum = r_accum + ct.astype(r_tile, ct.float32) @@ -1303,12 +1388,8 @@ def _ct_reduce_compute_h_kernel( # 3. Apply compute_h directly into split outputs. inv_r_eps = 1.0 / (r_val + eps) - bias_pre = ct.load( - Bias, index=(0, 0), shape=(1, n), padding_mode=PAD_ZERO, - ) - bias_post = ct.load( - Bias, index=(0, 1), shape=(1, n), padding_mode=PAD_ZERO, - ) + bias_pre = ct.load(Bias, index=(0, 0), shape=(1, n), padding_mode=PAD_ZERO) + bias_post = ct.load(Bias, index=(0, 1), shape=(1, n), padding_mode=PAD_ZERO) bias_pre = ct.astype(bias_pre, ct.float32) bias_post = ct.astype(bias_post, ct.float32) @@ -1325,14 +1406,14 @@ def _ct_reduce_compute_h_kernel( for split_idx in ct.static_iter(range(SPLIT_K)): bid_m_k = bid_m + split_idx * num_bid_m res_tile = ct.load( - Y_acc, index=(bid_m_k, 2 + res_chunk), shape=(TILE_SIZE_M, n), + Y_acc, + index=(bid_m_k, 2 + res_chunk), + shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) res_accum = res_accum + ct.astype(res_tile, ct.float32) - bias_res = ct.load( - Bias, index=(0, 2 + res_chunk), shape=(1, n), padding_mode=PAD_ZERO, - ) + bias_res = ct.load(Bias, index=(0, 2 + res_chunk), shape=(1, n), padding_mode=PAD_ZERO) bias_res = ct.astype(bias_res, ct.float32) h_res = res_accum * alpha_res * inv_r_eps + bias_res ct.store(PROJ_OUT, index=(bid_m, 2 + res_chunk), tile=res_accum.astype(PROJ_OUT.dtype)) @@ -1563,14 +1644,28 @@ def _cutile_proj_rms_fwd( grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), kernel=_ct_proj_rms_fwd_kernel, args_fn=lambda cfg: ( - x, weight, proj, norm, r, M, N, K, eps, - cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + x, + weight, + proj, + norm, + r, + M, + N, + K, + eps, + cfg.TILE_M, + cfg.TILE_N, + cfg.TILE_K, + cfg.SPLIT_K, ), search_space=configs, ) best = tuned.tuned_config _proj_rms_fwd_best_cfg[cache_key] = ( - best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, ) proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) @@ -1580,8 +1675,21 @@ def _cutile_proj_rms_fwd( stream, (math.ceil(M / best.TILE_M), best.SPLIT_K), _ct_proj_rms_fwd_kernel, - (x, weight, proj, norm, r, M, N, K, eps, - best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K), + ( + x, + weight, + proj, + norm, + r, + M, + N, + K, + eps, + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, + ), ) proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) @@ -1655,11 +1763,25 @@ def _cutile_reduce_compute_h( def _make_args(tm): return ( - proj_acc, norm_acc, bias_2d, - alpha_pre, alpha_post, alpha_res, - h_pre_out, h_post_out, h_res_out, r_out, proj_out, - M, N, K, n, eps, - tm, tile_n, split_k, + proj_acc, + norm_acc, + bias_2d, + alpha_pre, + alpha_post, + alpha_res, + h_pre_out, + h_post_out, + h_res_out, + r_out, + proj_out, + M, + N, + K, + n, + eps, + tm, + tile_n, + split_k, ) if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: @@ -1668,10 +1790,7 @@ def _make_args(tm): else: from types import SimpleNamespace - configs = [ - SimpleNamespace(TILE_M=tm) - for tm in _reduce_compute_h_autotune_configs(M) - ] + configs = [SimpleNamespace(TILE_M=tm) for tm in _reduce_compute_h_autotune_configs(M)] tuned = ct_experimental.autotune_launch( stream, grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M),), @@ -1682,8 +1801,7 @@ def _make_args(tm): best_tm = tuned.tuned_config.TILE_M _reduce_compute_h_best_cfg[cache_key] = best_tm ct.launch( - stream, (math.ceil(M / best_tm),), - _ct_reduce_compute_h_kernel, _make_args(best_tm), + stream, (math.ceil(M / best_tm),), _ct_reduce_compute_h_kernel, _make_args(best_tm) ) return h_pre_out, h_post_out, h_res_out, r_out, proj_out @@ -1752,8 +1870,19 @@ def _cutile_proj_rms_compute_h_fwd( (math.ceil(M / tm), split_k), _ct_proj_rms_fwd_kernel, ( - x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, - tm, tn, tk, split_k, + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, ), ) else: @@ -1766,14 +1895,28 @@ def _cutile_proj_rms_compute_h_fwd( grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), kernel=_ct_proj_rms_fwd_kernel, args_fn=lambda cfg: ( - x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, - cfg.TILE_M, cfg.TILE_N, cfg.TILE_K, cfg.SPLIT_K, + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + cfg.TILE_M, + cfg.TILE_N, + cfg.TILE_K, + cfg.SPLIT_K, ), search_space=configs, ) best = tuned.tuned_config _proj_rms_fwd_best_cfg[cache_key] = ( - best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, ) tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K @@ -1785,17 +1928,38 @@ def _cutile_proj_rms_compute_h_fwd( (math.ceil(M / tm), split_k), _ct_proj_rms_fwd_kernel, ( - x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, - tm, tn, tk, split_k, + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, ), ) # Launch reduce + compute_h kernel h_pre, h_post, h_res, r, proj_reduced = _cutile_reduce_compute_h( - proj_acc, norm_acc, bias, - alpha_pre, alpha_post, alpha_res, - n, M, N, K, eps, - tm, TILE_N, split_k, + proj_acc, + norm_acc, + bias, + alpha_pre, + alpha_post, + alpha_res, + n, + M, + N, + K, + eps, + tm, + TILE_N, + split_k, ) return h_pre, h_post, h_res, r, proj_reduced @@ -1850,21 +2014,49 @@ def _cutile_proj_rms_bwd( grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K), 1), kernel=_ct_proj_rms_bwd_kernel, args_fn=lambda cfg: ( - x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, - cfg.TILE_SIZE_M, cfg.TILE_SIZE_N, cfg.TILE_SIZE_K, + x, + weight, + norm, + grad_proj, + grad_r, + da, + db, + M, + N, + K, + eps, + cfg.TILE_SIZE_M, + cfg.TILE_SIZE_N, + cfg.TILE_SIZE_K, ), search_space=configs, ) best = tuned.tuned_config _proj_rms_bwd_best_cfg[cache_key] = ( - best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, ) ct.launch( stream, (math.ceil(K / best.TILE_SIZE_K), 1), _ct_proj_rms_bwd_kernel, - (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, - best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K), + ( + x, + weight, + norm, + grad_proj, + grad_r, + da, + db, + M, + N, + K, + eps, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ), ) else: num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count @@ -1881,20 +2073,20 @@ def _cutile_proj_rms_bwd( @ct.kernel def _ct_fused_grad_h_proj_kernel( - GRAD_H_PRE, # [M, n] - GRAD_H_POST, # [M, n] - GRAD_H_RES, # [M, n*n] - H_PRE, # [M, n] - H_POST, # [M, n] - PROJ, # [M, N] - R, # [M, 1] - GRAD_R_EXT, # [M, 1] - Alpha_pre, # [1] - Alpha_post, # [1] - Alpha_res, # [1] - GRAD_H, # [M, TILE_SIZE_N] output - GRAD_PROJ, # [M, TILE_SIZE_N] output - GRAD_R_TOTAL, # [M, 1] output + GRAD_H_PRE, # [M, n] + GRAD_H_POST, # [M, n] + GRAD_H_RES, # [M, n*n] + H_PRE, # [M, n] + H_POST, # [M, n] + PROJ, # [M, N] + R, # [M, 1] + GRAD_R_EXT, # [M, 1] + Alpha_pre, # [1] + Alpha_post, # [1] + Alpha_res, # [1] + GRAD_H, # [M, TILE_SIZE_N] output + GRAD_PROJ, # [M, TILE_SIZE_N] output + GRAD_R_TOTAL, # [M, 1] output M: int, N: int, n: ConstInt, @@ -1916,10 +2108,7 @@ def _ct_fused_grad_h_proj_kernel( alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() - r_tile = ct.load( - R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, - ) + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) r_tile = ct.astype(r_tile, ct.float32) r_eps = r_tile + eps @@ -1934,55 +2123,45 @@ def _ct_fused_grad_h_proj_kernel( if HAS_GRAD_H_PRE: gy_pre = ct.load( - GRAD_H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + GRAD_H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) gy_pre = ct.astype(gy_pre, ct.float32) else: gy_pre = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) - h_pre = ct.load( - H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, - ) + h_pre = ct.load(H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO) h_pre = ct.astype(h_pre, ct.float32) proj_pre = ct.load( - PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) proj_pre = ct.astype(proj_pre, ct.float32) grad_h_pre = gy_pre * h_pre * (1.0 - h_pre) grad_proj_pre = grad_h_pre * alpha_pre * inv_r_eps grad_r_from_h += ct.sum( - grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), - axis=1, keepdims=True, + grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True ) ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h_pre.astype(GRAD_H.dtype)) ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_pre.astype(GRAD_PROJ.dtype)) if HAS_GRAD_H_POST: gy_post = ct.load( - GRAD_H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + GRAD_H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) gy_post = ct.astype(gy_post, ct.float32) else: gy_post = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) h_post = ct.load( - H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) h_post = ct.astype(h_post, ct.float32) proj_post = ct.load( - PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), - padding_mode=PAD_ZERO, + PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) proj_post = ct.astype(proj_post, ct.float32) half_h_post = h_post * 0.5 grad_h_post = gy_post * half_h_post * (1.0 - half_h_post) * 2.0 grad_proj_post = grad_h_post * alpha_post * inv_r_eps grad_r_from_h += ct.sum( - grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), - axis=1, keepdims=True, + grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True ) ct.store(GRAD_H, index=(tile_m_id, 1), tile=grad_h_post.astype(GRAD_H.dtype)) ct.store(GRAD_PROJ, index=(tile_m_id, 1), tile=grad_proj_post.astype(GRAD_PROJ.dtype)) @@ -1990,35 +2169,35 @@ def _ct_fused_grad_h_proj_kernel( for res_chunk in ct.static_iter(range(n)): if HAS_GRAD_H_RES: grad_h_res = ct.load( - GRAD_H_RES, index=(tile_m_id, res_chunk), shape=(TILE_SIZE_M, n), + GRAD_H_RES, + index=(tile_m_id, res_chunk), + shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) grad_h_res = ct.astype(grad_h_res, ct.float32) else: grad_h_res = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) proj_res = ct.load( - PROJ, index=(tile_m_id, 2 + res_chunk), shape=(TILE_SIZE_M, n), + PROJ, + index=(tile_m_id, 2 + res_chunk), + shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO, ) proj_res = ct.astype(proj_res, ct.float32) grad_proj_res = grad_h_res * alpha_res * inv_r_eps grad_r_from_h += ct.sum( - grad_h_res * proj_res * alpha_res * (-inv_r_eps * inv_r_eps), - axis=1, keepdims=True, + grad_h_res * proj_res * alpha_res * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True ) + ct.store(GRAD_H, index=(tile_m_id, 2 + res_chunk), tile=grad_h_res.astype(GRAD_H.dtype)) ct.store( - GRAD_H, index=(tile_m_id, 2 + res_chunk), - tile=grad_h_res.astype(GRAD_H.dtype), - ) - ct.store( - GRAD_PROJ, index=(tile_m_id, 2 + res_chunk), + GRAD_PROJ, + index=(tile_m_id, 2 + res_chunk), tile=grad_proj_res.astype(GRAD_PROJ.dtype), ) if HAS_GRAD_R_EXT: grad_r_ext_tile = ct.load( - GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, + GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO ) grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) else: @@ -2029,13 +2208,13 @@ def _ct_fused_grad_h_proj_kernel( @ct.kernel def _ct_fused_grad_x_weight_kernel( - X, # [M, K] - WEIGHT, # [N, K] - GRAD_PROJ, # [M, TILE_SIZE_N] precomputed - GRAD_R_TOTAL, # [M, 1] precomputed - R, # [M, 1] - GRAD_X, # [M, K] output - GRAD_WEIGHT, # [N, K] output + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_SIZE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output M: int, N: int, K: int, @@ -2054,38 +2233,35 @@ def _ct_fused_grad_x_weight_kernel( # Load weight tile once — only depends on K-tile weight_tile = ct.load( - WEIGHT, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), - padding_mode=PAD_ZERO, + WEIGHT, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=PAD_ZERO ) acc_grad_weight = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) for tile_m_id in range(NUM_M_TILES): grad_proj_tile = ct.load( - GRAD_PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), + GRAD_PROJ, + index=(tile_m_id, 0), + shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO, ) x_tile = ct.load( - X, index=(tile_m_id, tile_k_id), shape=(TILE_SIZE_M, TILE_SIZE_K), + X, + index=(tile_m_id, tile_k_id), + shape=(TILE_SIZE_M, TILE_SIZE_K), padding_mode=PAD_ZERO, ) grad_r_total = ct.load( - GRAD_R_TOTAL, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, - ) - r_tile = ct.load( - R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, + GRAD_R_TOTAL, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO ) + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) r_tile = ct.astype(r_tile, ct.float32) # grad_x = grad_proj @ weight + grad_r_total * x / (r * K) inv_rK = 1.0 / (r_tile * K) acc_grad_x = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) acc_grad_x = ct.mma( - grad_proj_tile.astype(ct.tfloat32), - weight_tile.astype(ct.tfloat32), - acc=acc_grad_x, + grad_proj_tile.astype(ct.tfloat32), weight_tile.astype(ct.tfloat32), acc=acc_grad_x ) ct.store(GRAD_X, index=(tile_m_id, tile_k_id), tile=acc_grad_x.astype(GRAD_X.dtype)) @@ -2096,17 +2272,21 @@ def _ct_fused_grad_x_weight_kernel( acc=acc_grad_weight, ) - ct.store(GRAD_WEIGHT, index=(0, tile_k_id), tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype)) + ct.store( + GRAD_WEIGHT, + index=(0, tile_k_id), + tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype), + ) @ct.kernel def _ct_scalar_grads_partials_kernel( - GRAD_H, # [M, TILE_SIZE_N] precomputed - PROJ, # [M, N] - R, # [M, 1] - GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] output + GRAD_H, # [M, TILE_SIZE_N] precomputed + PROJ, # [M, N] + R, # [M, 1] + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] output GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] output - GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] output - GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] output + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] output + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] output M: int, N: int, n: int, @@ -2132,18 +2312,13 @@ def _ct_scalar_grads_partials_kernel( mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) grad_h = ct.load( - GRAD_H, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), - padding_mode=PAD_ZERO, + GRAD_H, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO ) proj_tile = ct.load( - PROJ, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), - padding_mode=PAD_ZERO, + PROJ, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO ) proj_tile = ct.astype(proj_tile, ct.float32) - r_tile = ct.load( - R, index=(bid_m, 0), shape=(TILE_SIZE_M, 1), - padding_mode=PAD_ZERO, - ) + r_tile = ct.load(R, index=(bid_m, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) r_tile = ct.astype(r_tile, ct.float32) r_eps = r_tile + eps @@ -2155,32 +2330,36 @@ def _ct_scalar_grads_partials_kernel( ga_res = ct.reshape(ct.sum(ga_all * mask_res_2d), (1, 1)) partial_gb = ct.sum(grad_h, axis=0, keepdims=False) ct.store( - GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), + GRAD_ALPHA_PRE_PARTIALS, + index=(bid_m, 0), tile=ga_pre.astype(GRAD_ALPHA_PRE_PARTIALS.dtype), ) ct.store( - GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), + GRAD_ALPHA_POST_PARTIALS, + index=(bid_m, 0), tile=ga_post.astype(GRAD_ALPHA_POST_PARTIALS.dtype), ) ct.store( - GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), + GRAD_ALPHA_RES_PARTIALS, + index=(bid_m, 0), tile=ga_res.astype(GRAD_ALPHA_RES_PARTIALS.dtype), ) ct.store( - GRAD_BIAS_PARTIALS, index=(bid_m, 0), + GRAD_BIAS_PARTIALS, + index=(bid_m, 0), tile=ct.reshape(partial_gb, (1, TILE_SIZE_N)).astype(GRAD_BIAS_PARTIALS.dtype), ) @ct.kernel def _ct_scalar_grads_reduce_kernel( - GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] - GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] - GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] - GRAD_ALPHA_PRE, # [1, 1] output - GRAD_ALPHA_POST, # [1, 1] output - GRAD_ALPHA_RES, # [1, 1] output - GRAD_BIAS, # [1, TILE_SIZE_N] output + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] + GRAD_ALPHA_PRE, # [1, 1] output + GRAD_ALPHA_POST, # [1, 1] output + GRAD_ALPHA_RES, # [1, 1] output + GRAD_BIAS, # [1, TILE_SIZE_N] output NUM_M_BLOCKS: int, TILE_SIZE_N: ConstInt, ): @@ -2192,20 +2371,16 @@ def _ct_scalar_grads_reduce_kernel( for bid_m in range(NUM_M_BLOCKS): acc_pre += ct.load( - GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), shape=(1, 1), - padding_mode=PAD_ZERO, + GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO ).astype(ct.float32) acc_post += ct.load( - GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), shape=(1, 1), - padding_mode=PAD_ZERO, + GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO ).astype(ct.float32) acc_res += ct.load( - GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), shape=(1, 1), - padding_mode=PAD_ZERO, + GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO ).astype(ct.float32) acc_bias += ct.load( - GRAD_BIAS_PARTIALS, index=(bid_m, 0), shape=(1, TILE_SIZE_N), - padding_mode=PAD_ZERO, + GRAD_BIAS_PARTIALS, index=(bid_m, 0), shape=(1, TILE_SIZE_N), padding_mode=PAD_ZERO ).astype(ct.float32) ct.store(GRAD_ALPHA_PRE, index=(0, 0), tile=acc_pre.astype(GRAD_ALPHA_PRE.dtype)) @@ -2215,13 +2390,13 @@ def _ct_scalar_grads_reduce_kernel( @ct.kernel def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( - X, # [M, K] - WEIGHT, # [N, K] - GRAD_PROJ, # [M, TILE_N] precomputed - GRAD_R_TOTAL, # [M, 1] precomputed - R, # [M, 1] - GRAD_X, # [M, K] output - GRAD_WEIGHT, # [N, K] output + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output M: int, N: int, K: int, @@ -2247,12 +2422,16 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) for m_tile in range(NUM_M_TILES): x_tile = ct.load( - X, index=(m_tile, tile_id), - shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), padding_mode=zero_pad, + X, + index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), + padding_mode=zero_pad, ) grad_proj_tile = ct.load( - GRAD_PROJ, index=(m_tile, 0), - shape=(TILE_DB_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + GRAD_PROJ, + index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, ) accumulator_db = ct.mma( @@ -2262,7 +2441,8 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( ) ct.store( - GRAD_WEIGHT, index=(0, tile_id), + GRAD_WEIGHT, + index=(0, tile_id), tile=accumulator_db.transpose().astype(GRAD_WEIGHT.dtype), allow_tma=False, ) @@ -2279,36 +2459,47 @@ def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( dd_tile_idx = tile_id // NUM_DA_K_TILES grad_proj_tile = ct.load( - GRAD_PROJ, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, TILE_N_SIZE), padding_mode=zero_pad, + GRAD_PROJ, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, ) grad_r_total = ct.load( - GRAD_R_TOTAL, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, + GRAD_R_TOTAL, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), + padding_mode=zero_pad, ) r_tile = ct.load( - R, index=(dd_tile_idx, 0), - shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad, + R, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad ) r_tile = ct.astype(r_tile, ct.float32) x_tile = ct.load( - X, index=(dd_tile_idx, b_tile_idx), - shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), padding_mode=zero_pad, + X, + index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), + padding_mode=zero_pad, ) inv_rK = 1.0 / (r_tile * K) accumulator_da = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) weight_tile = ct.load( - WEIGHT, index=(0, b_tile_idx), - shape=(TILE_N_SIZE, TILE_DA_SIZE_K), padding_mode=zero_pad, + WEIGHT, + index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), + padding_mode=zero_pad, ) accumulator_da = ct.mma( grad_proj_tile.astype(ct.tfloat32), weight_tile.astype(ct.tfloat32), acc=accumulator_da, ) - ct.store(GRAD_X, index=(dd_tile_idx, b_tile_idx), tile=accumulator_da.astype(GRAD_X.dtype)) + ct.store( + GRAD_X, + index=(dd_tile_idx, b_tile_idx), + tile=accumulator_da.astype(GRAD_X.dtype), + ) def _fused_grad_x_weight_autotune_configs(N): """Autotune search space for fused grad_x + grad_weight kernel.""" @@ -2380,13 +2571,29 @@ def _cutile_fused_compute_h_proj_rms_bwd( (math.ceil(M / tile_m_precomp),), _ct_fused_grad_h_proj_kernel, ( - grad_h_pre_arg, grad_h_post_arg, grad_h_res_arg, - h_pre, h_post, proj, r, grad_r_ext_arg, - alpha_pre, alpha_post, alpha_res, - grad_h_buf, grad_proj_buf, grad_r_total_buf, - M, N, n, eps, - tile_m_precomp, TILE_N, - int(has_grad_h_pre), int(has_grad_h_post), int(has_grad_h_res), + grad_h_pre_arg, + grad_h_post_arg, + grad_h_res_arg, + h_pre, + h_post, + proj, + r, + grad_r_ext_arg, + alpha_pre, + alpha_post, + alpha_res, + grad_h_buf, + grad_proj_buf, + grad_r_total_buf, + M, + N, + n, + eps, + tile_m_precomp, + TILE_N, + int(has_grad_h_pre), + int(has_grad_h_post), + int(has_grad_h_res), has_grad_r_ext_flag, ), ) @@ -2406,10 +2613,19 @@ def _cutile_fused_compute_h_proj_rms_bwd( (math.ceil(K / tk),), _ct_fused_grad_x_weight_kernel, ( - x, weight, grad_proj_buf, grad_r_total_buf, r, - grad_x, grad_weight, - M, N, K, - tm, tn, tk, + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + tm, + tn, + tk, ), ) else: @@ -2421,26 +2637,46 @@ def _cutile_fused_compute_h_proj_rms_bwd( grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K),), kernel=_ct_fused_grad_x_weight_kernel, args_fn=lambda cfg: ( - x, weight, grad_proj_buf, grad_r_total_buf, r, - grad_x, grad_weight, - M, N, K, - cfg.TILE_SIZE_M, cfg.TILE_SIZE_N, cfg.TILE_SIZE_K, + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + cfg.TILE_SIZE_M, + cfg.TILE_SIZE_N, + cfg.TILE_SIZE_K, ), search_space=configs, ) best = tuned.tuned_config _fused_grad_x_weight_best_cfg[cache_key] = ( - best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, ) ct.launch( stream, (math.ceil(K / best.TILE_SIZE_K),), _ct_fused_grad_x_weight_kernel, ( - x, weight, grad_proj_buf, grad_r_total_buf, r, - grad_x, grad_weight, - M, N, K, - best.TILE_SIZE_M, best.TILE_SIZE_N, best.TILE_SIZE_K, + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, ), ) else: @@ -2450,9 +2686,16 @@ def _cutile_fused_compute_h_proj_rms_bwd( (num_sms, 2, 1), _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, ( - x, weight, grad_proj_buf, grad_r_total_buf, r, - grad_x, grad_weight, - M, N, K, + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, TILE_N, ), ) @@ -2474,11 +2717,19 @@ def _cutile_fused_compute_h_proj_rms_bwd( (num_m_blocks,), _ct_scalar_grads_partials_kernel, ( - grad_h_buf, proj, r, - grad_alpha_pre_partials, grad_alpha_post_partials, grad_alpha_res_partials, + grad_h_buf, + proj, + r, + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, grad_bias_partials, - M, N, n, eps, - tile_m_scalar, TILE_N, + M, + N, + n, + eps, + tile_m_scalar, + TILE_N, ), ) ct.launch( @@ -2486,10 +2737,16 @@ def _cutile_fused_compute_h_proj_rms_bwd( (1,), _ct_scalar_grads_reduce_kernel, ( - grad_alpha_pre_partials, grad_alpha_post_partials, grad_alpha_res_partials, + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, grad_bias_partials, - grad_alpha_pre, grad_alpha_post, grad_alpha_res, grad_bias, - num_m_blocks, TILE_N, + grad_alpha_pre, + grad_alpha_post, + grad_alpha_res, + grad_bias, + num_m_blocks, + TILE_N, ), ) @@ -2519,7 +2776,6 @@ def _cutile_fused_compute_h_proj_rms_bwd( # ============================================================================ from megatron.core.transformer.hyper_connection import ( - native_fused_add_3 as fused_add_3, native_h_aggregate, native_h_post_bda, native_proj_rms, @@ -2551,9 +2807,7 @@ def _get_triton_h_post_bda_bwd(): return globals().get("_triton_h_post_bda_bwd") -def _torch_h_aggregate_bwd( - grad_output: Tensor, x: Tensor, h_pre: Tensor -) -> Tuple[Tensor, Tensor]: +def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tuple[Tensor, Tensor]: grad_x = grad_output.unsqueeze(2) * h_pre.unsqueeze(-1) grad_h = torch.sum(grad_output.unsqueeze(2) * x, dim=-1) return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) @@ -2603,11 +2857,7 @@ def _torch_proj_rms_compute_h( proj = torch.matmul(x, weight.t()) r = x.float().norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) alpha = torch.cat( - [ - alpha_pre.expand(n), - alpha_post.expand(n), - alpha_res.expand(weight.shape[0] - 2 * n), - ], + [alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(weight.shape[0] - 2 * n)], dim=-1, ) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) @@ -2682,11 +2932,20 @@ def forward( eps: float = 1e-6, ): h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( - x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, + x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps ) ctx.save_for_backward( - x, weight, h_pre, h_post, h_res, proj_reduced, r, - alpha_pre, alpha_post, alpha_res, bias, + x, + weight, + h_pre, + h_post, + h_res, + proj_reduced, + r, + alpha_pre, + alpha_post, + alpha_res, + bias, ) ctx.n = n ctx.eps = eps @@ -2695,16 +2954,38 @@ def forward( @staticmethod def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): ( - x, weight, h_pre, h_post, h_res, proj, r, - alpha_pre, alpha_post, alpha_res, bias_param, + x, + weight, + h_pre, + h_post, + h_res, + proj, + r, + alpha_pre, + alpha_post, + alpha_res, + bias_param, ) = ctx.saved_tensors grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias = ( _cutile_fused_compute_h_proj_rms_bwd( - x, weight, grad_h_pre, grad_h_post, grad_h_res, - h_pre, h_post, h_res, proj, r, grad_r_ext, - alpha_pre, alpha_post, alpha_res, - bias_param, ctx.n, ctx.eps, + x, + weight, + grad_h_pre, + grad_h_post, + grad_h_res, + h_pre, + h_post, + h_res, + proj, + r, + grad_r_ext, + alpha_pre, + alpha_post, + alpha_res, + bias_param, + ctx.n, + ctx.eps, ) ) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 0f122a50efb..8be7bcf6ec3 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -302,9 +302,14 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: x_2d = x.reshape(s * b, self.n * self.hidden_size) with torch.cuda.nvtx.range("HyperConnection::fused_proj_rms_compute_h"): h_pre, h_post, h_res, _ = self._proj_rms_compute_h_op( - x_2d, self.mapping_proj.weight, - self.alpha_pre, self.alpha_post, self.alpha_res, - self.bias, self.n, self.norm_eps, + x_2d, + self.mapping_proj.weight, + self.alpha_pre, + self.alpha_post, + self.alpha_res, + self.bias, + self.n, + self.norm_eps, ) h_pre = h_pre.view(s, b, self.n) h_post = h_post.view(s, b, self.n) @@ -317,9 +322,7 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: h_pre, h_post, h_res = self._compute_h(proj, r) h_res = h_res.view(s, b, self.n, self.n) - h_res = self._sinkhorn_op( - h_res, self.sinkhorn_iterations, self.norm_eps - ) # [s, b, n, n] + h_res = self._sinkhorn_op(h_res, self.sinkhorn_iterations, self.norm_eps) # [s, b, n, n] return h_pre, h_post, h_res diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 13e26e4a708..a16d407340f 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -107,9 +107,14 @@ def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): def _ref_proj_rms_compute_h( - x: Tensor, weight: Tensor, - alpha_pre: Tensor, alpha_post: Tensor, alpha_res: Tensor, - bias: Tensor, n: int, eps: float = 1e-6, + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, ): """Reference: fused proj_rms + compute_h.""" proj = torch.matmul(x, weight.t()) @@ -117,9 +122,7 @@ def _ref_proj_rms_compute_h( K = x.shape[-1] r = norm / math.sqrt(K) # [M, 1] N = proj.shape[-1] - alpha = torch.cat([ - alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n), - ]) + alpha = torch.cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n)]) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) h_pre = h[..., :n].sigmoid() h_post = h[..., n : 2 * n].sigmoid() * 2 @@ -392,7 +395,9 @@ def test_fwd_vs_reference(self, s, b, n, C, with_bias): @_require_triton @pytest.mark.parametrize("with_bias", [True, False]) - @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)]) + @pytest.mark.parametrize( + "s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)] + ) def test_bwd_vs_reference(self, s, b, n, C, with_bias): """Triton hpb backward grads must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd @@ -407,7 +412,7 @@ def test_bwd_vs_reference(self, s, b, n, C, with_bias): # -- Triton backward -- g_hr_t, g_res_t, g_hp_t, g_x_t, g_bias_t = _triton_h_post_bda_bwd( - grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + grad_out, hr_data, orig_data, hp_data, x_data, bias_data ) # -- Reference backward via autograd -- @@ -430,7 +435,10 @@ def test_bwd_vs_reference(self, s, b, n, C, with_bias): ) if with_bias: torch.testing.assert_close( - g_bias_t, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, + g_bias_t, + bi_r.grad, + atol=BWD_ATOL, + rtol=BWD_RTOL, msg="Triton backward mismatch on bias", ) @@ -440,8 +448,10 @@ def test_bwd_vs_reference(self, s, b, n, C, with_bias): @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) def test_triton_vs_cutile(self, s, b, n, C, with_bias): """Triton and cuTile backward must produce identical results.""" - from megatron.core.fusions.fused_mhc_kernels import _cutile_h_post_bda_bwd - from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, + ) _info() hr_data = _rand(s, b, n, n) @@ -452,18 +462,20 @@ def test_triton_vs_cutile(self, s, b, n, C, with_bias): grad_out = _rand(s, b, n, C) triton_out = _triton_h_post_bda_bwd( - grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + grad_out, hr_data, orig_data, hp_data, x_data, bias_data ) cutile_out = _cutile_h_post_bda_bwd( - grad_out, hr_data, orig_data, hp_data, x_data, bias_data, + grad_out, hr_data, orig_data, hp_data, x_data, bias_data ) for i, name in enumerate(["h_res", "orig_res", "h_post", "x", "bias"]): if triton_out[i] is None: continue torch.testing.assert_close( - triton_out[i], cutile_out[i], - atol=BWD_ATOL, rtol=BWD_RTOL, + triton_out[i], + cutile_out[i], + atol=BWD_ATOL, + rtol=BWD_RTOL, msg=f"Triton vs cuTile mismatch on {name}", ) @@ -477,12 +489,12 @@ def test_e2e_inputs_no_nan(self): """Feed actual E2E backward inputs to Triton kernel and check for NaN.""" from megatron.core.fusions.fused_mhc_kernels import ( _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, fused_h_aggregate, fused_h_post_bda, fused_proj_rms, fused_sinkhorn, ) - from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd s, b, n, C = 8, 4, 4, 1024 eps = 1e-6 @@ -528,12 +540,14 @@ def test_e2e_inputs_no_nan(self): assert not tr_t.isnan().any(), f"Triton {name} has NaN" assert not tr_t.isinf().any(), f"Triton {name} has Inf" torch.testing.assert_close( - tr_t, ct_t, atol=BWD_ATOL, rtol=BWD_RTOL, + tr_t, + ct_t, + atol=BWD_ATOL, + rtol=BWD_RTOL, msg=f"Triton vs cuTile mismatch on {name} (E2E inputs)", ) - # ============================================================================ # Triton: Sinkhorn # ============================================================================ @@ -567,8 +581,8 @@ def test_triton_vs_cutile(self, s, b, n, iters): from megatron.core.fusions.fused_mhc_kernels import ( _cutile_sinkhorn_bwd, _cutile_sinkhorn_fwd, + triton_fused_sinkhorn, ) - from megatron.core.fusions.fused_mhc_kernels import triton_fused_sinkhorn eps = 1e-6 data = _rand(s, b, n, n) @@ -665,10 +679,7 @@ def test_fwd_bwd_vs_reference(self, M, N, K): class TestFusedProjRmsComputeH: - @pytest.mark.parametrize( - "M,n,K", - [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)], - ) + @pytest.mark.parametrize("M,n,K", [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)]) def test_fwd_bwd_vs_reference(self, M, n, K): """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h @@ -724,9 +735,15 @@ def _make_inputs(): ) loss_r.backward() - torch.testing.assert_close(h_pre_f, h_pre_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_pre mismatch") - torch.testing.assert_close(h_post_f, h_post_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post mismatch") - torch.testing.assert_close(h_res_f, h_res_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_res mismatch") + torch.testing.assert_close( + h_pre_f, h_pre_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_pre mismatch" + ) + torch.testing.assert_close( + h_post_f, h_post_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post mismatch" + ) + torch.testing.assert_close( + h_res_f, h_res_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_res mismatch" + ) torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="r mismatch") torch.testing.assert_close( xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" @@ -743,9 +760,7 @@ def _make_inputs(): _assert_cosine_similar( arf.grad, arr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_res" ) - _assert_cosine_similar( - bf.grad, br.grad, COSINE_SIM_THRESH, msg="backward mismatch on bias" - ) + _assert_cosine_similar(bf.grad, br.grad, COSINE_SIM_THRESH, msg="backward mismatch on bias") # ============================================================================ @@ -908,19 +923,12 @@ def _run_ref(): loss.backward() return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() - proj_f, r_f, out_f, agg_f, grad_f = _run_fused() proj_r, r_r, out_r, agg_r, grad_r = _run_ref() - torch.testing.assert_close( - proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL - ) - torch.testing.assert_close( - r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL - ) - torch.testing.assert_close( - agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL - ) + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL) torch.testing.assert_close( out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" ) @@ -962,7 +970,7 @@ def _run_fused_compute_h(): x_2d = hs.reshape(s * b, n * C) h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( - x_2d, w, ap, apo, ar, bias_p, n, eps, + x_2d, w, ap, apo, ar, bias_p, n, eps ) h_pre = h_pre.view(s, b, n) @@ -990,7 +998,7 @@ def _run_ref(): x_2d = hs.reshape(s * b, n * C) h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( - x_2d, w, ap, apo, ar, bias_p, n, eps, + x_2d, w, ap, apo, ar, bias_p, n, eps ) h_pre = h_pre.view(s, b, n) @@ -1018,7 +1026,9 @@ def _run_ref(): out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" ) _assert_cosine_similar( - grad_f, grad_r, COSINE_SIM_THRESH, + grad_f, + grad_r, + COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, fused compute_h)", ) @@ -1171,8 +1181,7 @@ def _run_inline_ref(): out_b, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch (broadcast)" ) _assert_cosine_similar( - grad_b, grad_r, COSINE_SIM_THRESH, - msg="hidden_states grad (E2E backward, broadcast)", + grad_b, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, broadcast)" ) @@ -1258,11 +1267,16 @@ def _run_ref(): agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (fused broadcast)" ) torch.testing.assert_close( - out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, msg="h_post_bda output mismatch (fused broadcast)", ) _assert_cosine_similar( - grad_f, grad_r, COSINE_SIM_THRESH, + grad_f, + grad_r, + COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, fused broadcast)", ) @@ -1304,7 +1318,7 @@ def _run_fused_compute_h_broadcast(): x_2d = hs_map.reshape(s * b, n * C) h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( - x_2d, w, ap, apo, ar, bias_p, n, eps, + x_2d, w, ap, apo, ar, bias_p, n, eps ) h_pre = h_pre.view(s, b, n) @@ -1332,7 +1346,7 @@ def _run_ref(): x_2d = hs.reshape(s * b, n * C) h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( - x_2d, w, ap, apo, ar, bias_p, n, eps, + x_2d, w, ap, apo, ar, bias_p, n, eps ) h_pre = h_pre.view(s, b, n) @@ -1354,14 +1368,22 @@ def _run_ref(): out_r, agg_r, grad_r = _run_ref() torch.testing.assert_close( - agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, + agg_f, + agg_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, msg="aggregated mismatch (fused compute_h broadcast)", ) torch.testing.assert_close( - out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, msg="h_post_bda output mismatch (fused compute_h broadcast)", ) _assert_cosine_similar( - grad_f, grad_r, COSINE_SIM_THRESH, + grad_f, + grad_r, + COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, fused compute_h broadcast)", ) diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index d8f54455e26..31fcbdff3d1 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -72,9 +72,7 @@ def test_forward_normal_vs_checkpoint_correctness(self): # Forward without checkpoint (reference) torch.manual_seed(42) torch.cuda.manual_seed(42) - aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal( - hidden_states - ) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal(hidden_states) mixed_ref = module.apply_h_res(h_res_ref, residual) loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() loss_ref.backward() From f31deaedcacaa68ada6a49199609e67cb69f8de7 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Wed, 13 May 2026 00:41:31 -0700 Subject: [PATCH 10/34] upd: fix lint --- megatron/core/fusions/fused_mhc_kernels.py | 14 ++++++++++++++ megatron/core/transformer/hyper_connection.py | 2 ++ 2 files changed, 16 insertions(+) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 81e08d86b10..ad040cce0ed 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -211,6 +211,7 @@ class TritonFusedSinkhorn(torch.autograd.Function): @staticmethod def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run Triton Sinkhorn forward and save initial matrix for backward.""" out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) ctx.save_for_backward(M_init) ctx.num_iterations = num_iterations @@ -219,6 +220,7 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output: Tensor): + """Run Triton Sinkhorn backward.""" (M_init,) = ctx.saved_tensors grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None @@ -2874,6 +2876,7 @@ class CutileSinkhornKnopp(torch.autograd.Function): @staticmethod def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run cuTile Sinkhorn forward and save initial matrix for backward.""" output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) ctx.save_for_backward(M_init) ctx.num_iterations = num_iterations @@ -2882,6 +2885,7 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output): + """Run cuTile Sinkhorn backward.""" (M_init,) = ctx.saved_tensors grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None @@ -2891,12 +2895,14 @@ class CutileHAggregate(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, h_pre: Tensor): + """Run cuTile h_aggregate forward.""" output = _cutile_h_aggregate_fwd(x, h_pre) ctx.save_for_backward(x, h_pre) return output @staticmethod def backward(ctx, grad_output): + """Run cuTile h_aggregate backward.""" x, h_pre = ctx.saved_tensors return _cutile_h_aggregate_bwd(grad_output, x, h_pre) @@ -2905,6 +2911,7 @@ class CutileProjRms(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): + """Run cuTile projection plus RMS normalization forward.""" proj, norm, r = _cutile_proj_rms_fwd(x, weight, eps) ctx.save_for_backward(x, weight, norm) ctx.eps = eps @@ -2912,6 +2919,7 @@ def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): @staticmethod def backward(ctx, grad_proj, grad_r): + """Run cuTile projection plus RMS normalization backward.""" x, weight, norm = ctx.saved_tensors grad_x, grad_weight = _cutile_proj_rms_bwd(grad_proj, grad_r, x, weight, norm, ctx.eps) return grad_x, grad_weight, None @@ -2931,6 +2939,7 @@ def forward( n: int, eps: float = 1e-6, ): + """Run fused cuTile projection, RMS normalization, and compute_h forward.""" h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps ) @@ -2953,6 +2962,7 @@ def forward( @staticmethod def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): + """Run fused cuTile projection, RMS normalization, and compute_h backward.""" ( x, weight, @@ -2997,6 +3007,7 @@ class FusedHAggregate(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, h_pre: Tensor): + """Run h_aggregate forward using the best available backend.""" triton_fwd = _get_triton_h_aggregate_fwd() if triton_fwd is not None: output = triton_fwd(x, h_pre) @@ -3009,6 +3020,7 @@ def forward(ctx, x: Tensor, h_pre: Tensor): @staticmethod def backward(ctx, grad_output): + """Run h_aggregate backward using the best available backend.""" x, h_pre = ctx.saved_tensors if _CUTILE_AVAILABLE: return _cutile_h_aggregate_bwd(grad_output, x, h_pre) @@ -3027,6 +3039,7 @@ def forward( x: Tensor, bias: Optional[Tensor], ): + """Run h_post_bda forward using the best available backend.""" triton_fwd = _get_triton_h_post_bda_fwd() if triton_fwd is not None: output = triton_fwd(h_res, original_residual, h_post, x, bias) @@ -3044,6 +3057,7 @@ def forward( @staticmethod def backward(ctx, grad_output): + """Run h_post_bda backward using the best available backend.""" if ctx.has_bias: h_res, orig_res, h_post, x, bias = ctx.saved_tensors else: diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 8be7bcf6ec3..e5c064c323f 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -108,11 +108,13 @@ class BroadcastTensorFused(torch.autograd.Function): @staticmethod def forward(ctx, x, fused_add_3_fn): + """Return three view aliases and save the fused gradient combiner.""" ctx.fused_add_3_fn = fused_add_3_fn return x.view_as(x), x.view_as(x), x.view_as(x) @staticmethod def backward(ctx, grad1, grad2, grad3): + """Combine gradients from the three broadcast aliases.""" grads = [g for g in (grad1, grad2, grad3) if g is not None] if len(grads) == 0: return None, None From 08f4183331a9bca6d975583d5ceda0852eeafe0c Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 14 May 2026 10:26:52 +0800 Subject: [PATCH 11/34] Constrain mHC reduce compute_h tile tuning --- megatron/core/fusions/fused_mhc_kernels.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index ad040cce0ed..f9981d21d23 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1704,7 +1704,10 @@ def _cutile_proj_rms_fwd( def _reduce_compute_h_autotune_configs(M): """Generate autotune search space for reduce_compute_h kernel.""" + min_tile_m = 16 if M >= 16 else 1 for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m < min_tile_m: + continue if tile_m <= M and M % tile_m == 0: yield tile_m From d27a4ba616cd32752088e0b52a66884ac71fbf8d Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 14 May 2026 13:04:17 +0800 Subject: [PATCH 12/34] Export fused_add_3 public wrapper --- megatron/core/fusions/fused_mhc_kernels.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index f9981d21d23..1a3e263da97 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -2781,6 +2781,7 @@ def _cutile_fused_compute_h_proj_rms_bwd( # ============================================================================ from megatron.core.transformer.hyper_connection import ( + native_fused_add_3, native_h_aggregate, native_h_post_bda, native_proj_rms, @@ -2788,6 +2789,11 @@ def _cutile_fused_compute_h_proj_rms_bwd( ) +def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: + """Add three tensors using the native torch.compile-backed implementation.""" + return native_fused_add_3(a, b, c) + + def _get_triton_sinkhorn(): if not _TRITON_AVAILABLE: return None From dce7d7655de11cf58d56b006c26674006664eba3 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 14 May 2026 15:36:09 +0800 Subject: [PATCH 13/34] Match proj rms compute_h torch fallback reference --- megatron/core/fusions/fused_mhc_kernels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 1a3e263da97..bf440177203 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -2866,7 +2866,7 @@ def _torch_proj_rms_compute_h( eps: float, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: proj = torch.matmul(x, weight.t()) - r = x.float().norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) + r = x.norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) alpha = torch.cat( [alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(weight.shape[0] - 2 * n)], dim=-1, @@ -2875,7 +2875,7 @@ def _torch_proj_rms_compute_h( h_pre = h[..., :n].sigmoid() h_post = h[..., n : 2 * n].sigmoid() * 2 h_res = h[..., 2 * n :] - return h_pre, h_post, h_res, r.to(dtype=x.dtype) + return h_pre, h_post, h_res, r if _CUTILE_AVAILABLE: From e09d1ef38e75e52dc1bec82116780d922a11f5ff Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 14 May 2026 18:02:21 +0800 Subject: [PATCH 14/34] Add mHC backend test overrides --- megatron/core/fusions/fused_mhc_kernels.py | 69 ++++++++++++++++++++-- 1 file changed, 63 insertions(+), 6 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index bf440177203..402e72ec772 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -15,11 +15,40 @@ """ import math +import os from typing import Optional, Tuple import torch from torch import Tensor + +def _env_flag(name: str) -> bool: + return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") + + +def _forced_backend() -> str: + value = os.getenv("MHC_FORCE_BACKEND", "auto").strip().lower() + value = value.replace("-", "_").replace("+", "_") + aliases = { + "auto": "auto", + "mixed": "auto", + "default": "auto", + "native": "native", + "torch": "native", + "pytorch": "native", + "none": "native", + "triton": "triton", + "triton_native": "triton", + "cutile": "cutile", + "cu_tile": "cutile", + "cuda_tile": "cutile", + } + if value not in aliases: + valid = ", ".join(sorted(aliases)) + raise ValueError(f"Unsupported MHC_FORCE_BACKEND={value!r}; expected one of: {valid}") + return aliases[value] + + # --------------------------------------------------------------------------- # Check cuTile availability # --------------------------------------------------------------------------- @@ -39,11 +68,6 @@ pass -def is_cutile_available() -> bool: - """Return True if cuTile fused kernels are available.""" - return _CUTILE_AVAILABLE - - # --------------------------------------------------------------------------- # Check Triton availability # --------------------------------------------------------------------------- @@ -57,8 +81,41 @@ def is_cutile_available() -> bool: pass +_MHC_FORCED_BACKEND = _forced_backend() + +if _MHC_FORCED_BACKEND == "native": + _TRITON_AVAILABLE = False + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "triton": + if not _TRITON_AVAILABLE: + raise RuntimeError("MHC_FORCE_BACKEND=triton was requested, but Triton is not available") + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "cutile": + if not _CUTILE_AVAILABLE: + raise RuntimeError("MHC_FORCE_BACKEND=cutile was requested, but cuTile is not available") + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_TRITON"): + if _MHC_FORCED_BACKEND == "triton": + raise ValueError("MHC_FORCE_BACKEND=triton conflicts with MHC_DISABLE_TRITON=1") + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_CUTILE"): + if _MHC_FORCED_BACKEND == "cutile": + raise ValueError("MHC_FORCE_BACKEND=cutile conflicts with MHC_DISABLE_CUTILE=1") + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False + + +def is_cutile_available() -> bool: + """Return True if cuTile fused kernels are enabled.""" + return _CUTILE_AVAILABLE + + def is_triton_available() -> bool: - """Return True if Triton is available for supported mHC kernels.""" + """Return True if Triton is enabled for supported mHC kernels.""" return _TRITON_AVAILABLE From 732d6141355968ed91007abc4f6cade0659ed362 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 19 May 2026 14:10:28 +0800 Subject: [PATCH 15/34] Address mHC fused review cleanup --- megatron/core/fusions/fused_mhc_kernels.py | 74 +++++++++++++++++-- megatron/core/transformer/hyper_connection.py | 5 ++ .../fusions/test_fused_mhc_kernels.py | 15 +++- 3 files changed, 84 insertions(+), 10 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 402e72ec772..9f96a5fcb0e 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -14,6 +14,7 @@ - proj_rms: fused projection + RMS normalization """ +import logging import math import os from typing import Optional, Tuple @@ -21,6 +22,11 @@ import torch from torch import Tensor +from megatron.core._rank_utils import log_single_rank + +logger = logging.getLogger(__name__) +LOG2E = math.log2(math.e) + def _env_flag(name: str) -> bool: return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") @@ -147,7 +153,8 @@ def _triton_sinkhorn_fwd_kernel( logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) row_max = tl.max(logits, axis=1) - M = tl.exp2((logits - row_max[:, None]) * 1.4426950408889634) + # Subtract row_max before exp to keep the exponent numerically stable. + M = tl.exp2((logits - row_max[:, None]) * LOG2E) tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) for _ in range(NUM_ITERS): @@ -725,7 +732,6 @@ def _triton_h_post_bda_bwd( if _CUTILE_AVAILABLE: ConstInt = ct.Constant[int] PAD_ZERO = ct.PaddingMode.ZERO - LOG2E = 1.4426950408889634 # -- Sinkhorn kernels ---------------------------------------------------- @@ -1343,6 +1349,8 @@ def _ct_proj_rms_fwd_kernel( ''' Grid: (num_tiles_m, num_tiles_k). Fused matmul + norm + r: proj, norm, r in one pass over K. + R is a retained signature placeholder; r is computed after split-K + reduction from NORM. ''' tile_m_id = ct.bid(0) split_k_id = ct.bid(1) @@ -1438,10 +1446,7 @@ def _ct_reduce_compute_h_kernel( # 2. Compute r = norm / sqrt(K) denom = ct.full((TILE_SIZE_M, 1), K * 1.0, dtype=ct.float32) - mean = ct.truediv(r_accum, denom) - rstd = ct.rsqrt(mean) - ones = ct.full((TILE_SIZE_M, 1), 1.0, dtype=ct.float32) - r_val = ct.truediv(ones, rstd) # norm / sqrt(K) + r_val = ct.sqrt(ct.truediv(r_accum, denom)) ct.store(R, index=(bid_m, 0), tile=r_val.astype(R.dtype)) @@ -1663,6 +1668,8 @@ def _cutile_proj_rms_fwd( proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # _ct_proj_rms_fwd_kernel keeps R in its signature; r is computed + # below from the reduced norm. r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) ct.launch( @@ -1684,6 +1691,7 @@ def _cutile_proj_rms_fwd( tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) proj = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. r = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) ct.launch( stream, @@ -1697,6 +1705,7 @@ def _cutile_proj_rms_fwd( mx_split_k = max(cfg.SPLIT_K for cfg in configs) proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) norm = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for autotune launches; not read. r = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) tuned = ct_experimental.autotune_launch( stream, @@ -1728,6 +1737,7 @@ def _cutile_proj_rms_fwd( ) proj = torch.empty(best.SPLIT_K * M, N, dtype=x.dtype, device=dev) norm = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. r = torch.empty(best.SPLIT_K * M, 1, dtype=x.dtype, device=dev) # Re-launch with best config for correct output. ct.launch( @@ -1909,6 +1919,8 @@ def _cutile_proj_rms_compute_h_fwd( proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # _ct_proj_rms_fwd_kernel keeps R in its signature; reduce_compute_h + # computes r from norm_acc. r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) ct.launch( @@ -1926,6 +1938,7 @@ def _cutile_proj_rms_compute_h_fwd( tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) ct.launch( stream, @@ -1951,6 +1964,7 @@ def _cutile_proj_rms_compute_h_fwd( mx_split_k = max(cfg.SPLIT_K for cfg in configs) proj_acc = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) norm_acc = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for autotune launches; not read. r_placeholder = torch.empty(mx_split_k * M, 1, dtype=x.dtype, device=dev) tuned = ct_experimental.autotune_launch( stream, @@ -1984,6 +1998,7 @@ def _cutile_proj_rms_compute_h_fwd( proj_acc = torch.empty(split_k * M, N, dtype=x.dtype, device=dev) norm_acc = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. r_placeholder = torch.empty(split_k * M, 1, dtype=x.dtype, device=dev) ct.launch( stream, @@ -2846,6 +2861,48 @@ def _cutile_fused_compute_h_proj_rms_bwd( ) +_BACKEND_INFO_LOGGED = False + + +def _select_triton_cutile_native(triton_impl) -> str: + if triton_impl is not None: + return "triton" + if _CUTILE_AVAILABLE: + return "cutile" + return "native" + + +def _mhc_backend_selection() -> str: + """Return a concise description of the selected mHC fused backends.""" + sinkhorn = _select_triton_cutile_native(_get_triton_sinkhorn()) + h_aggregate_fwd = _select_triton_cutile_native(_get_triton_h_aggregate_fwd()) + h_aggregate_bwd = "cutile" if _CUTILE_AVAILABLE else "native" + h_post_bda_fwd = _select_triton_cutile_native(_get_triton_h_post_bda_fwd()) + h_post_bda_bwd = _select_triton_cutile_native(_get_triton_h_post_bda_bwd()) + proj_rms = "cutile" if _CUTILE_AVAILABLE else "native" + return ( + f"MHC_FORCE_BACKEND={_MHC_FORCED_BACKEND}; " + f"sinkhorn={sinkhorn}; " + f"h_aggregate=fwd:{h_aggregate_fwd},bwd:{h_aggregate_bwd}; " + f"h_post_bda=fwd:{h_post_bda_fwd},bwd:{h_post_bda_bwd}; " + f"proj_rms={proj_rms}; " + f"proj_rms_compute_h={proj_rms}" + ) + + +def log_fused_mhc_backend_once() -> None: + """Log the fused mHC backend selection once per process.""" + global _BACKEND_INFO_LOGGED + if _BACKEND_INFO_LOGGED: + return + _BACKEND_INFO_LOGGED = True + log_single_rank( + logger, + logging.INFO, + f"[mHC] fused backend selection: {_mhc_backend_selection()}", + ) + + def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: """Add three tensors using the native torch.compile-backed implementation.""" return native_fused_add_3(a, b, c) @@ -2876,8 +2933,9 @@ def _get_triton_h_post_bda_bwd(): def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tuple[Tensor, Tensor]: - grad_x = grad_output.unsqueeze(2) * h_pre.unsqueeze(-1) - grad_h = torch.sum(grad_output.unsqueeze(2) * x, dim=-1) + grad_output_expanded = grad_output.unsqueeze(2) + grad_x = grad_output_expanded * h_pre.unsqueeze(-1) + grad_h = torch.sum(grad_output_expanded * x, dim=-1) return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index e5c064c323f..83321b11ab7 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -208,8 +208,10 @@ def __init__(self, config: TransformerConfig, layer_number: int): fused_proj_rms, fused_proj_rms_compute_h, fused_sinkhorn, + log_fused_mhc_backend_once, ) + log_fused_mhc_backend_once() self._sinkhorn_op = fused_sinkhorn self._h_aggregate_op = fused_h_aggregate self._h_post_bda_op = fused_h_post_bda @@ -469,6 +471,9 @@ def forward( When provided, uses _forward_with_checkpoint for memory-efficient execution. Returns: + A 4-tuple. This is an intentional breaking change from the older + 3-tuple API because fused_h_res_h_post_bda consumes the residual + branch created by BroadcastTensorFused. aggregated: [s, b, C] - aggregated input for layer computation h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) h_post: [s, b, n] - expansion weights diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index a16d407340f..20bee2eb812 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -4,7 +4,8 @@ Each test compares the fused kernel's forward output AND backward gradients against a pure-PyTorch differentiable reference to catch numerical drift -introduced by kernel fusion. +introduced by kernel fusion. The public fused API tests also exercise backend +dispatch and fallback selection through the same entry points used by mHC. """ import math @@ -163,6 +164,8 @@ def test_fwd_bwd_vs_torch_reference(self, s, b, n, iters): class TestFusedSinkhorn: + """Public fused sinkhorn dispatch/fallback plus numerical correctness.""" + @pytest.mark.flaky_in_dev @_require_cutile @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) @@ -222,6 +225,8 @@ def test_fwd_bwd_vs_torch_reference(self, s, b, n, C): class TestFusedHAggregate: + """Public fused h_aggregate dispatch/fallback plus numerical correctness.""" + @pytest.mark.flaky_in_dev @_require_cutile @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) @@ -321,6 +326,8 @@ def _make_inputs(): class TestFusedHPostBDA: + """Public fused h_post_bda dispatch/fallback plus numerical correctness.""" + @pytest.mark.flaky_in_dev @_require_cutile @pytest.mark.parametrize("with_bias", [True, False]) @@ -637,6 +644,8 @@ def test_fwd_bwd_vs_torch_reference(self, M, N, K): class TestFusedProjRms: + """Public fused proj_rms dispatch/fallback plus numerical correctness.""" + @pytest.mark.flaky_in_dev @_require_cutile @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) @@ -679,6 +688,8 @@ def test_fwd_bwd_vs_reference(self, M, N, K): class TestFusedProjRmsComputeH: + """Public fused proj_rms_compute_h dispatch/fallback plus numerical correctness.""" + @pytest.mark.parametrize("M,n,K", [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)]) def test_fwd_bwd_vs_reference(self, M, n, K): """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" @@ -1186,7 +1197,7 @@ def _run_inline_ref(): class TestEndToEndFusedBroadcast: - """Full mHC pipeline using the public fused API + BroadcastTensorFused.""" + """Full mHC pipeline through public fused dispatch + BroadcastTensorFused.""" def test_full_pipeline_fwd_bwd(self): from megatron.core.fusions.fused_mhc_kernels import ( From c9184f5f16edd6cc4d76fbd702af8f44a6d7b36d Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 19 May 2026 14:23:54 +0800 Subject: [PATCH 16/34] Fix mHC h_res transpose application --- megatron/core/fusions/fused_mhc_kernels.py | 48 ++++++++++--------- megatron/core/transformer/hyper_connection.py | 18 +++---- .../fusions/test_fused_mhc_kernels.py | 27 ++++++++++- .../test_hyper_connection_recompute.py | 15 ++++++ 4 files changed, 75 insertions(+), 33 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 9f96a5fcb0e..26c56790cd6 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -10,7 +10,7 @@ Four fused operations: - sinkhorn: Sinkhorn-Knopp projection to doubly stochastic matrix - h_aggregate: weighted n-stream -> 1-stream aggregation - - h_post_bda: fused H_res @ residual + H_post * (x + bias) + - h_post_bda: fused H_res.T @ residual + H_post * (x + bias) - proj_rms: fused projection + RMS normalization """ @@ -396,7 +396,7 @@ def _triton_hpb_fwd_kernel( BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, ): - """out = hr @ orig + hp * (x + bias).""" + """out = hr.T @ orig + hp * (x + bias).""" pid_s = tl.program_id(0) pid_c = tl.program_id(1) offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) @@ -417,8 +417,8 @@ def _triton_hpb_fwd_kernel( out_i = hp_i[:, None] * x_tile for j in tl.static_range(N): - hr_ij = tl.load( - hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + hr_ji = tl.load( + hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, mask=mask_s, other=0.0, ).to(tl.float32) @@ -430,7 +430,7 @@ def _triton_hpb_fwd_kernel( mask=mask_2d, other=0.0, ).to(tl.float32) - out_i += hr_ij[:, None] * orig_j + out_i += hr_ji[:, None] * orig_j tl.store( out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], @@ -505,7 +505,7 @@ def _triton_hpb_bwd_g_x_orig_kernel( BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, ): - """g_x = hp @ go, g_orig = hr.T @ go.""" + """g_x = hp @ go, g_orig = hr @ go.""" pid_s = tl.program_id(0) pid_c = tl.program_id(1) offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) @@ -537,12 +537,12 @@ def _triton_hpb_bwd_g_x_orig_kernel( mask=mask_2d, other=0.0, ).to(tl.float32) - hr_ji = tl.load( - hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, + hr_ij = tl.load( + hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, mask=mask_s, other=0.0, ).to(tl.float32) - g_orig_i += hr_ji[:, None] * go_j + g_orig_i += hr_ij[:, None] * go_j tl.store( g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], g_orig_i.to(g_orig_ptr.dtype.element_ty), @@ -582,7 +582,7 @@ def _triton_hpb_bwd_g_hp_hr_kernel( BLOCK_C: tl.constexpr, BLOCK_S: tl.constexpr, ): - """g_hp = sum_c go*(x+bias), g_hr = go @ orig.T.""" + """g_hp = sum_c go*(x+bias), g_hr = orig @ go.T.""" pid_s = tl.program_id(0) offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) mask_s = offs_s < sb @@ -625,7 +625,7 @@ def _triton_hpb_bwd_g_hp_hr_kernel( ).to(tl.float32) dot_hr = tl.sum(go_i * orig_j, axis=1) g_hr_acc += tl.where( - tl.arange(0, N * N)[None, :] == i * N + j, + tl.arange(0, N * N)[None, :] == j * N + i, dot_hr[:, None], tl.zeros((BLOCK_S, N * N), dtype=tl.float32), ) @@ -1038,7 +1038,8 @@ def _ct_hpb_fwd_kernel( x_exp = ct.expand_dims(x_tile, axis=1) # (TILE_SIZE, 1, TILE_C) out_tile = hp_exp * x_exp # (TILE_SIZE, N, TILE_C) for j in range(N): - hr_col = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) out_tile = out_tile + hr_col * orig_row ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) @@ -1063,7 +1064,8 @@ def _ct_hpb_fwd_bias_kernel( xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # (TILE_SIZE, 1, TILE_C) out_tile = hp_exp * xb_exp # (TILE_SIZE, N, TILE_C) for j in range(N): - hr_col = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) out_tile = out_tile + hr_col * orig_row ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) @@ -1072,7 +1074,7 @@ def _ct_hpb_fwd_bias_kernel( def _ct_hpb_bwd_g_x_orig_kernel( go, hr, hp, g_orig, g_x, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): - """Compute g_x = hp @ go and g_orig = hr.T @ go. + """Compute g_x = hp @ go and g_orig = hr @ go. Grid: (ceil(sb / TILE_SIZE), ceil(C / TILE_C)). 2D grid — no loop, no accumulators. @@ -1091,8 +1093,8 @@ def _ct_hpb_bwd_g_x_orig_kernel( hp_j_exp = ct.expand_dims(hp_j, axis=2) # [TS, 1, 1] go_j = ct.extract(go_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) g_x_tile = g_x_tile + hp_j_exp * go_j - hr_row_j = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) - g_orig_tile = g_orig_tile + ct.reshape(hr_row_j, (TILE_SIZE, N, 1)) * go_j + hr_col_j = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + g_orig_tile = g_orig_tile + hr_col_j * go_j ct.store( g_x, index=(pid, ct_idx), @@ -1104,7 +1106,7 @@ def _ct_hpb_bwd_g_x_orig_kernel( def _ct_hpb_bwd_g_hp_hr_kernel( go, orig, x, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): - """Compute g_hp = sum(go * x) and g_hr = go @ orig.T (no bias). + """Compute g_hp = sum(go * x) and g_hr = orig @ go.T (no bias). Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. """ @@ -1125,7 +1127,7 @@ def _ct_hpb_bwd_g_hp_hr_kernel( ) acc_g_hp = acc_g_hp + ct.sum(go_tile * x_exp, axis=2, keepdims=True) acc_g_hr = acc_g_hr + ct.sum( - ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 ) ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) @@ -1134,7 +1136,7 @@ def _ct_hpb_bwd_g_hp_hr_kernel( def _ct_hpb_bwd_g_hp_hr_bias_kernel( go, orig, x, bias, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt ): - """Compute g_hp = sum(go * (x+bias)) and g_hr = go @ orig.T (with bias). + """Compute g_hp = sum(go * (x+bias)) and g_hr = orig @ go.T (with bias). Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. """ @@ -1156,7 +1158,7 @@ def _ct_hpb_bwd_g_hp_hr_bias_kernel( ) acc_g_hp = acc_g_hp + ct.sum(go_tile * xb_exp, axis=2, keepdims=True) acc_g_hr = acc_g_hr + ct.sum( - ct.expand_dims(go_tile, axis=2) * ct.expand_dims(orig_tile, axis=1), axis=3 + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 ) ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) @@ -2955,8 +2957,8 @@ def _torch_h_post_bda_bwd( hp = h_post.reshape(sb, n) x_flat = x.reshape(sb, C) - g_hr = torch.bmm(go, orig.transpose(1, 2)).view(s, b, n, n) - g_res = torch.bmm(hr.transpose(1, 2), go).view(s, b, n, C) + g_hr = torch.bmm(orig, go.transpose(1, 2)).view(s, b, n, n) + g_res = torch.bmm(hr, go).view(s, b, n, C) g_x = torch.sum(go * hp.unsqueeze(-1), dim=1).view(s, b, C) xb = x_flat if bias is None else x_flat + bias.view(1, C) g_hp = torch.sum(go * xb.unsqueeze(1), dim=2).view(s, b, n) @@ -3216,7 +3218,7 @@ def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: def fused_h_post_bda( h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: - """Fused H_res @ residual + H_post * (x + bias).""" + """Fused H_res.T @ residual + H_post * (x + bias).""" if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) return native_h_post_bda(h_res, original_residual, h_post, x, bias) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 83321b11ab7..e33b62a9d6e 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -69,11 +69,11 @@ def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: def native_h_post_bda( h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: - """Native H_res @ residual + H_post * (x [+ bias]).""" + """Native H_res.T @ residual + H_post * (x [+ bias]).""" s, b, n, C = original_residual.shape h_res_batched = h_res.view(s * b, n, n) residual_batched = original_residual.view(s * b, n, C) - mixed = torch.bmm(h_res_batched, residual_batched).view(s, b, n, C) + mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C) x_expanded = h_post.unsqueeze(-1) * x.unsqueeze(2) if bias is not None: bias_expanded = h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) @@ -153,13 +153,13 @@ class HyperConnectionModule(MegatronModule): Unified mHC (Manifold-Constrained Hyper-Connections) module. Implements the complete mHC propagation: - x_{l+1} = H_res @ x_l + H_post^T @ F(H_pre @ x_l) + x_{l+1} = H_res^T @ x_l + H_post^T @ F(H_pre @ x_l) This module handles: 1. Computing learnable mappings: H_pre, H_post, H_res (with Sinkhorn-Knopp projection) 2. Aggregation: n-stream → 1-stream (H_pre @ x) 3. Expansion: 1-stream → n-stream (H_post^T @ output) - 4. Residual merge: H_res @ x + expanded_output + 4. Residual merge: H_res^T @ x + expanded_output 5. Block-level expand/contract for TransformerBlock boundaries Args: @@ -434,7 +434,7 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: """ Apply H_res to residual using H_res weights. - Computes: H_res @ residual + Computes: H_res.T @ residual Args: h_res: [s, b, n, n] - residual mixing matrix @@ -449,8 +449,8 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: # [s, b, n*C] -> [s, b, n, C] -> [s*b, n, C] residual_batched = residual.view(s, b, n, C).view(s * b, n, C) - # Batch matrix multiply: [s*b, n, n] @ [s*b, n, C] -> [s*b, n, C] - mixed = torch.bmm(h_res_batched, residual_batched) + # Batch matrix multiply: [s*b, n, n].T @ [s*b, n, C] -> [s*b, n, C] + mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched) return mixed.view(s, b, n * C) @@ -611,7 +611,7 @@ def fused_h_res_h_post_bda( Currently implements the operations sequentially using native PyTorch. The computation flow is: - 1. mixed = H_res @ original_residual (apply_h_res) + 1. mixed = H_res.T @ original_residual (apply_h_res) 2. expanded = H_post^T @ layer_output (apply_h_post) 3. output = dropout(expanded + bias) + mixed (bias-dropout-add) @@ -667,7 +667,7 @@ def _fused_h_res_h_post_bda_native( h_res, h_post and bda. When dropout is zero (or inference), uses a single fused/reference kernel - for H_res @ residual + H_post * (x + bias). Falls back to unfused + for H_res.T @ residual + H_post * (x + bias). Falls back to unfused implementation when dropout is needed. Args: diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 20bee2eb812..33d3e93c316 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -91,7 +91,9 @@ def _ref_h_post_bda( h_res: Tensor, orig_res: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: s, b, n, C = orig_res.shape - mixed = torch.bmm(h_res.view(s * b, n, n), orig_res.view(s * b, n, C)).view(s, b, n, C) + h_res_batched = h_res.view(s * b, n, n) + orig_batched = orig_res.view(s * b, n, C) + mixed = torch.bmm(h_res_batched.transpose(1, 2), orig_batched).view(s, b, n, C) x_exp = h_post.unsqueeze(-1) * x.unsqueeze(2) out = x_exp + mixed if bias is not None: @@ -99,6 +101,15 @@ def _ref_h_post_bda( return out +def _h_post_bda_transpose_case(): + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=DTYPE, device=DEVICE) + orig = torch.tensor([[[[10.0, 100.0], [1.0, 2.0]]]], dtype=DTYPE, device=DEVICE) + h_post = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + x = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + expected = torch.tensor([[[[13.0, 106.0], [24.0, 208.0]]]], dtype=DTYPE, device=DEVICE) + return h_res, orig, h_post, x, expected + + def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): proj = torch.matmul(x, weight.t()) norm = x.norm(dim=-1, keepdim=True) @@ -282,6 +293,12 @@ def test_fwd_vs_reference(self, s, b, n, C): class TestNativeHPostBDA: """Tests for native_h_post_bda.""" + def test_forward_uses_h_res_transpose(self): + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = native_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + @pytest.mark.parametrize("with_bias", [True, False]) @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) def test_fwd_bwd_vs_torch_reference(self, s, b, n, C, with_bias): @@ -328,6 +345,14 @@ def _make_inputs(): class TestFusedHPostBDA: """Public fused h_post_bda dispatch/fallback plus numerical correctness.""" + def test_forward_uses_h_res_transpose(self): + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = fused_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + @pytest.mark.flaky_in_dev @_require_cutile @pytest.mark.parametrize("with_bias", [True, False]) diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index 31fcbdff3d1..e05ac5cfd4b 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -46,6 +46,21 @@ def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4 module.cuda() return module + def test_apply_h_res_uses_h_res_transpose(self): + """apply_h_res should compute H_res.T @ residual.""" + module = self._create_hyper_connection_module(hidden_size=4, num_residual_streams=2) + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], device='cuda') + residual = torch.tensor( + [[[10.0, 100.0, 3.0, 4.0, 1.0, 2.0, 5.0, 6.0]]], device='cuda' + ) + expected = torch.tensor( + [[[13.0, 106.0, 18.0, 22.0, 24.0, 208.0, 26.0, 32.0]]], device='cuda' + ) + + mixed = module.apply_h_res(h_res, residual) + + torch.testing.assert_close(mixed, expected, atol=0.0, rtol=0.0) + def test_forward_normal_vs_checkpoint_correctness(self): """ Test that _forward_with_checkpoint produces the same outputs as _forward_normal. From 394c17323b24d152aa8ae517849006429f454f39 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 19 May 2026 14:54:31 +0800 Subject: [PATCH 17/34] Fix Triton mHC LOG2E constexpr --- megatron/core/fusions/fused_mhc_kernels.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 26c56790cd6..e08f07a5945 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -130,6 +130,8 @@ def is_triton_available() -> bool: # ============================================================================ if _TRITON_AVAILABLE: + TLOG2E = tl.constexpr(LOG2E) + # ============================================================================ # Sinkhorn-Knopp # ============================================================================ @@ -154,7 +156,7 @@ def _triton_sinkhorn_fwd_kernel( logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) row_max = tl.max(logits, axis=1) # Subtract row_max before exp to keep the exponent numerically stable. - M = tl.exp2((logits - row_max[:, None]) * LOG2E) + M = tl.exp2((logits - row_max[:, None]) * TLOG2E) tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) for _ in range(NUM_ITERS): From 5b43329e927e0edd4bad21b3a3210440abd914af Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Tue, 19 May 2026 01:33:31 -0700 Subject: [PATCH 18/34] Fix: lint --- megatron/core/fusions/fused_mhc_kernels.py | 5 +---- .../transformer/test_hyper_connection_recompute.py | 4 +--- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index e08f07a5945..2be4447c1a9 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -2864,7 +2864,6 @@ def _cutile_fused_compute_h_proj_rms_bwd( native_sinkhorn, ) - _BACKEND_INFO_LOGGED = False @@ -2901,9 +2900,7 @@ def log_fused_mhc_backend_once() -> None: return _BACKEND_INFO_LOGGED = True log_single_rank( - logger, - logging.INFO, - f"[mHC] fused backend selection: {_mhc_backend_selection()}", + logger, logging.INFO, f"[mHC] fused backend selection: {_mhc_backend_selection()}" ) diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py index e05ac5cfd4b..7c6be437b78 100644 --- a/tests/unit_tests/transformer/test_hyper_connection_recompute.py +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -50,9 +50,7 @@ def test_apply_h_res_uses_h_res_transpose(self): """apply_h_res should compute H_res.T @ residual.""" module = self._create_hyper_connection_module(hidden_size=4, num_residual_streams=2) h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], device='cuda') - residual = torch.tensor( - [[[10.0, 100.0, 3.0, 4.0, 1.0, 2.0, 5.0, 6.0]]], device='cuda' - ) + residual = torch.tensor([[[10.0, 100.0, 3.0, 4.0, 1.0, 2.0, 5.0, 6.0]]], device='cuda') expected = torch.tensor( [[[13.0, 106.0, 18.0, 22.0, 24.0, 208.0, 26.0, 32.0]]], device='cuda' ) From c4259b9e804b16f0d47d7e1b98259a9dfdc6b71b Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Mon, 1 Jun 2026 10:37:59 +0800 Subject: [PATCH 19/34] upd: fix --- megatron/core/fusions/fused_mhc_kernels.py | 96 ++++++++++++++-------- 1 file changed, 60 insertions(+), 36 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 2be4447c1a9..372bc8d4d56 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -32,7 +32,7 @@ def _env_flag(name: str) -> bool: return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") -def _forced_backend() -> str: +def _forced_backend() -> Tuple[str, Optional[Exception]]: value = os.getenv("MHC_FORCE_BACKEND", "auto").strip().lower() value = value.replace("-", "_").replace("+", "_") aliases = { @@ -51,8 +51,10 @@ def _forced_backend() -> str: } if value not in aliases: valid = ", ".join(sorted(aliases)) - raise ValueError(f"Unsupported MHC_FORCE_BACKEND={value!r}; expected one of: {valid}") - return aliases[value] + return "auto", ValueError( + f"Unsupported MHC_FORCE_BACKEND={value!r}; expected one of: {valid}" + ) + return aliases[value], None # --------------------------------------------------------------------------- @@ -87,7 +89,18 @@ def _forced_backend() -> str: pass -_MHC_FORCED_BACKEND = _forced_backend() +_MHC_FORCED_BACKEND, _MHC_BACKEND_VALIDATION_ERROR = _forced_backend() + + +def _record_mhc_backend_validation_error(error: Exception) -> None: + global _MHC_BACKEND_VALIDATION_ERROR + if _MHC_BACKEND_VALIDATION_ERROR is None: + _MHC_BACKEND_VALIDATION_ERROR = error + + +def _raise_mhc_backend_validation_error() -> None: + if _MHC_BACKEND_VALIDATION_ERROR is not None: + raise _MHC_BACKEND_VALIDATION_ERROR if _MHC_FORCED_BACKEND == "native": _TRITON_AVAILABLE = False @@ -95,22 +108,30 @@ def _forced_backend() -> str: _CUTILE_EXPERIMENTAL_AVAILABLE = False elif _MHC_FORCED_BACKEND == "triton": if not _TRITON_AVAILABLE: - raise RuntimeError("MHC_FORCE_BACKEND=triton was requested, but Triton is not available") + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=triton was requested, but Triton is not available") + ) _CUTILE_AVAILABLE = False _CUTILE_EXPERIMENTAL_AVAILABLE = False elif _MHC_FORCED_BACKEND == "cutile": if not _CUTILE_AVAILABLE: - raise RuntimeError("MHC_FORCE_BACKEND=cutile was requested, but cuTile is not available") + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=cutile was requested, but cuTile is not available") + ) _TRITON_AVAILABLE = False if _env_flag("MHC_DISABLE_TRITON"): if _MHC_FORCED_BACKEND == "triton": - raise ValueError("MHC_FORCE_BACKEND=triton conflicts with MHC_DISABLE_TRITON=1") + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=triton conflicts with MHC_DISABLE_TRITON=1") + ) _TRITON_AVAILABLE = False if _env_flag("MHC_DISABLE_CUTILE"): if _MHC_FORCED_BACKEND == "cutile": - raise ValueError("MHC_FORCE_BACKEND=cutile conflicts with MHC_DISABLE_CUTILE=1") + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=cutile conflicts with MHC_DISABLE_CUTILE=1") + ) _CUTILE_AVAILABLE = False _CUTILE_EXPERIMENTAL_AVAILABLE = False @@ -958,16 +979,6 @@ def _ct_h_agg_bwd_kernel(go, x, h_pre, gx, gh, N: ConstInt, TILE_M: ConstInt, TI gh_acc += ct.sum(go_expanded * x_tile, axis=2) ct.store(gh, index=(pid, 0), tile=gh_acc.astype(gh.dtype)) - def _h_agg_autotune_configs(sb, C): - """Generate autotune search space for h_aggregate kernels.""" - for tile_m in (1, 2, 4, 8): - for tile_c in (32, 64, 128, 256, 512, 1024): - if tile_m <= sb and tile_c <= C: - yield {"TILE_M": tile_m, "TILE_C": tile_c} - - _h_agg_fwd_best_cfg: dict = {} - _h_agg_bwd_best_cfg: dict = {} - def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: s, b, n, C = x.shape sb = s * b @@ -976,14 +987,8 @@ def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: x_flat = x.view(sb, n, C) h_flat = h_pre.view(sb, n) - cache_key = (sb, n, C) - cached = _h_agg_fwd_best_cfg.get(cache_key) - # Autotune disabled — causes cudaErrorLaunchFailure during training. - if cached is not None: - tm, tc = cached - else: - tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) ct.launch( stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, (x_flat, h_flat, out, n, tm, tc) ) @@ -1002,14 +1007,8 @@ def _cutile_h_aggregate_bwd( x_flat = x.view(sb, n, C) h_flat = h_pre.view(sb, n) - cache_key = (sb, n, C) - cached = _h_agg_bwd_best_cfg.get(cache_key) - # Autotune disabled — causes cudaErrorLaunchFailure during training. - if cached is not None: - tm, tc = cached - else: - tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) ct.launch( stream, (math.ceil(sb / tm),), @@ -2875,15 +2874,23 @@ def _select_triton_cutile_native(triton_impl) -> str: return "native" -def _mhc_backend_selection() -> str: - """Return a concise description of the selected mHC fused backends.""" +def _mhc_backend_status() -> Tuple[str, bool]: + """Return backend description and whether every backend is native.""" sinkhorn = _select_triton_cutile_native(_get_triton_sinkhorn()) h_aggregate_fwd = _select_triton_cutile_native(_get_triton_h_aggregate_fwd()) h_aggregate_bwd = "cutile" if _CUTILE_AVAILABLE else "native" h_post_bda_fwd = _select_triton_cutile_native(_get_triton_h_post_bda_fwd()) h_post_bda_bwd = _select_triton_cutile_native(_get_triton_h_post_bda_bwd()) proj_rms = "cutile" if _CUTILE_AVAILABLE else "native" - return ( + selected = ( + sinkhorn, + h_aggregate_fwd, + h_aggregate_bwd, + h_post_bda_fwd, + h_post_bda_bwd, + proj_rms, + ) + message = ( f"MHC_FORCE_BACKEND={_MHC_FORCED_BACKEND}; " f"sinkhorn={sinkhorn}; " f"h_aggregate=fwd:{h_aggregate_fwd},bwd:{h_aggregate_bwd}; " @@ -2891,16 +2898,27 @@ def _mhc_backend_selection() -> str: f"proj_rms={proj_rms}; " f"proj_rms_compute_h={proj_rms}" ) + return message, all(backend == "native" for backend in selected) + + +def _mhc_backend_selection() -> str: + """Return a concise description of the selected mHC fused backends.""" + message, _ = _mhc_backend_status() + return message def log_fused_mhc_backend_once() -> None: """Log the fused mHC backend selection once per process.""" + _raise_mhc_backend_validation_error() global _BACKEND_INFO_LOGGED if _BACKEND_INFO_LOGGED: return _BACKEND_INFO_LOGGED = True + backend_selection, all_native = _mhc_backend_status() log_single_rank( - logger, logging.INFO, f"[mHC] fused backend selection: {_mhc_backend_selection()}" + logger, + logging.WARNING if all_native else logging.INFO, + f"[mHC] fused backend selection: {backend_selection}", ) @@ -2940,6 +2958,7 @@ def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tup return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) +@torch.compile def _torch_h_post_bda_bwd( grad_output: Tensor, h_res: Tensor, @@ -3199,6 +3218,7 @@ def backward(ctx, grad_output): def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: """Project logits to a doubly stochastic matrix using Triton, cuTile, then torch.""" + _raise_mhc_backend_validation_error() triton_sinkhorn = _get_triton_sinkhorn() if triton_sinkhorn is not None: return triton_sinkhorn(input_logits, num_iterations, eps) @@ -3209,6 +3229,7 @@ def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: """Weighted n-stream to 1-stream aggregation using Triton/cuTile/torch.""" + _raise_mhc_backend_validation_error() if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: return FusedHAggregate.apply(x, h_pre) return native_h_aggregate(x, h_pre) @@ -3218,6 +3239,7 @@ def fused_h_post_bda( h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] ) -> Tensor: """Fused H_res.T @ residual + H_post * (x + bias).""" + _raise_mhc_backend_validation_error() if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) return native_h_post_bda(h_res, original_residual, h_post, x, bias) @@ -3225,6 +3247,7 @@ def fused_h_post_bda( def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: """Projection + RMS normalization using cuTile, then torch.""" + _raise_mhc_backend_validation_error() if _CUTILE_AVAILABLE: return CutileProjRms.apply(x, weight, eps) return native_proj_rms(x, weight, eps) @@ -3241,6 +3264,7 @@ def fused_proj_rms_compute_h( eps: float = 1e-6, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" + _raise_mhc_backend_validation_error() if _CUTILE_AVAILABLE: return CutileProjRmsComputeH.apply( x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps From c005d27f525f82cf490d702ea05cd41bad4f718b Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Sun, 31 May 2026 19:40:40 -0700 Subject: [PATCH 20/34] upd: fix lint --- megatron/core/fusions/fused_mhc_kernels.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 372bc8d4d56..b35f8b5c50e 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -102,6 +102,7 @@ def _raise_mhc_backend_validation_error() -> None: if _MHC_BACKEND_VALIDATION_ERROR is not None: raise _MHC_BACKEND_VALIDATION_ERROR + if _MHC_FORCED_BACKEND == "native": _TRITON_AVAILABLE = False _CUTILE_AVAILABLE = False From e9041a0de549fa34ce1e05b3c6f141e1c5da8763 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Mon, 1 Jun 2026 12:54:01 +0800 Subject: [PATCH 21/34] fix: handle hybrid mhc tuple and cutile availability --- megatron/core/fusions/fused_mhc_kernels.py | 92 ++++++++++++++++--- megatron/core/models/hybrid/hybrid_block.py | 3 +- .../fusions/test_fused_mhc_kernels.py | 4 +- 3 files changed, 83 insertions(+), 16 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index b35f8b5c50e..48bb7ad3fa0 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -17,6 +17,8 @@ import logging import math import os +import shutil +import subprocess from typing import Optional, Tuple import torch @@ -62,6 +64,8 @@ def _forced_backend() -> Tuple[str, Optional[Exception]]: # --------------------------------------------------------------------------- _CUTILE_AVAILABLE = False _CUTILE_EXPERIMENTAL_AVAILABLE = False +_CUTILE_DEVICE_SUPPORT_CACHE: Optional[bool] = None +_CUTILE_DEVICE_SUPPORT_ERROR: Optional[str] = None try: import cuda.tile as ct @@ -101,6 +105,11 @@ def _record_mhc_backend_validation_error(error: Exception) -> None: def _raise_mhc_backend_validation_error() -> None: if _MHC_BACKEND_VALIDATION_ERROR is not None: raise _MHC_BACKEND_VALIDATION_ERROR + if _MHC_FORCED_BACKEND == "cutile" and not is_cutile_available(): + raise RuntimeError( + "MHC_FORCE_BACKEND=cutile was requested, but cuTile does not support " + f"the current device: {_CUTILE_DEVICE_SUPPORT_ERROR}" + ) if _MHC_FORCED_BACKEND == "native": @@ -139,7 +148,64 @@ def _raise_mhc_backend_validation_error() -> None: def is_cutile_available() -> bool: """Return True if cuTile fused kernels are enabled.""" - return _CUTILE_AVAILABLE + return _CUTILE_AVAILABLE and _cutile_supports_current_device() + + +def _get_tileiras_path() -> Optional[str]: + """Return the tileiras compiler path if it can be found.""" + tileiras = shutil.which("tileiras") + if tileiras is not None: + return tileiras + + cuda_home = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or "/usr/local/cuda" + candidate = os.path.join(cuda_home, "bin", "tileiras") + if os.path.exists(candidate): + return candidate + return None + + +def _cutile_supports_current_device() -> bool: + """Return whether cuTile can compile for the current CUDA device.""" + global _CUTILE_DEVICE_SUPPORT_CACHE, _CUTILE_DEVICE_SUPPORT_ERROR + + if not _CUTILE_AVAILABLE: + return False + if _CUTILE_DEVICE_SUPPORT_CACHE is not None: + return _CUTILE_DEVICE_SUPPORT_CACHE + + if not torch.cuda.is_available(): + _CUTILE_DEVICE_SUPPORT_CACHE = True + return True + + major, minor = torch.cuda.get_device_capability() + arch = f"sm_{major}{minor}" + tileiras = _get_tileiras_path() + if tileiras is None: + _CUTILE_DEVICE_SUPPORT_ERROR = "tileiras compiler was not found" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + try: + result = subprocess.run( + [tileiras, "--gpu-name", arch], + capture_output=True, + check=False, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError) as exc: + _CUTILE_DEVICE_SUPPORT_ERROR = str(exc) + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + output = f"{result.stdout}\n{result.stderr}" + if "Cannot find option named" in output and arch in output: + _CUTILE_DEVICE_SUPPORT_ERROR = output.strip().splitlines()[0] + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + _CUTILE_DEVICE_SUPPORT_CACHE = True + return True def is_triton_available() -> bool: @@ -2870,7 +2936,7 @@ def _cutile_fused_compute_h_proj_rms_bwd( def _select_triton_cutile_native(triton_impl) -> str: if triton_impl is not None: return "triton" - if _CUTILE_AVAILABLE: + if is_cutile_available(): return "cutile" return "native" @@ -2879,10 +2945,10 @@ def _mhc_backend_status() -> Tuple[str, bool]: """Return backend description and whether every backend is native.""" sinkhorn = _select_triton_cutile_native(_get_triton_sinkhorn()) h_aggregate_fwd = _select_triton_cutile_native(_get_triton_h_aggregate_fwd()) - h_aggregate_bwd = "cutile" if _CUTILE_AVAILABLE else "native" + h_aggregate_bwd = "cutile" if is_cutile_available() else "native" h_post_bda_fwd = _select_triton_cutile_native(_get_triton_h_post_bda_fwd()) h_post_bda_bwd = _select_triton_cutile_native(_get_triton_h_post_bda_bwd()) - proj_rms = "cutile" if _CUTILE_AVAILABLE else "native" + proj_rms = "cutile" if is_cutile_available() else "native" selected = ( sinkhorn, h_aggregate_fwd, @@ -3156,7 +3222,7 @@ def forward(ctx, x: Tensor, h_pre: Tensor): triton_fwd = _get_triton_h_aggregate_fwd() if triton_fwd is not None: output = triton_fwd(x, h_pre) - elif _CUTILE_AVAILABLE: + elif is_cutile_available(): output = _cutile_h_aggregate_fwd(x, h_pre) else: output = native_h_aggregate(x, h_pre) @@ -3167,7 +3233,7 @@ def forward(ctx, x: Tensor, h_pre: Tensor): def backward(ctx, grad_output): """Run h_aggregate backward using the best available backend.""" x, h_pre = ctx.saved_tensors - if _CUTILE_AVAILABLE: + if is_cutile_available(): return _cutile_h_aggregate_bwd(grad_output, x, h_pre) return _torch_h_aggregate_bwd(grad_output, x, h_pre) @@ -3188,7 +3254,7 @@ def forward( triton_fwd = _get_triton_h_post_bda_fwd() if triton_fwd is not None: output = triton_fwd(h_res, original_residual, h_post, x, bias) - elif _CUTILE_AVAILABLE: + elif is_cutile_available(): output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) else: output = native_h_post_bda(h_res, original_residual, h_post, x, bias) @@ -3212,7 +3278,7 @@ def backward(ctx, grad_output): triton_bwd = _get_triton_h_post_bda_bwd() if triton_bwd is not None: return triton_bwd(grad_output, h_res, orig_res, h_post, x, bias) - if _CUTILE_AVAILABLE: + if is_cutile_available(): return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) return _torch_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) @@ -3223,7 +3289,7 @@ def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) triton_sinkhorn = _get_triton_sinkhorn() if triton_sinkhorn is not None: return triton_sinkhorn(input_logits, num_iterations, eps) - if _CUTILE_AVAILABLE: + if is_cutile_available(): return CutileSinkhornKnopp.apply(input_logits, num_iterations, eps) return native_sinkhorn(input_logits, num_iterations, eps) @@ -3231,7 +3297,7 @@ def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: """Weighted n-stream to 1-stream aggregation using Triton/cuTile/torch.""" _raise_mhc_backend_validation_error() - if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: + if _TRITON_AVAILABLE or is_cutile_available(): return FusedHAggregate.apply(x, h_pre) return native_h_aggregate(x, h_pre) @@ -3241,7 +3307,7 @@ def fused_h_post_bda( ) -> Tensor: """Fused H_res.T @ residual + H_post * (x + bias).""" _raise_mhc_backend_validation_error() - if _TRITON_AVAILABLE or _CUTILE_AVAILABLE: + if _TRITON_AVAILABLE or is_cutile_available(): return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) return native_h_post_bda(h_res, original_residual, h_post, x, bias) @@ -3249,7 +3315,7 @@ def fused_h_post_bda( def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: """Projection + RMS normalization using cuTile, then torch.""" _raise_mhc_backend_validation_error() - if _CUTILE_AVAILABLE: + if is_cutile_available(): return CutileProjRms.apply(x, weight, eps) return native_proj_rms(x, weight, eps) @@ -3266,7 +3332,7 @@ def fused_proj_rms_compute_h( ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" _raise_mhc_backend_validation_error() - if _CUTILE_AVAILABLE: + if is_cutile_available(): return CutileProjRmsComputeH.apply( x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps ) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index bfdeb6029f1..2b7760b454e 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -144,8 +144,7 @@ def forward( mhc_recompute_manager=None, ) -> Tuple[Tensor, Optional[Tensor]]: """Run the wrapped hybrid layer through one layer-boundary mHC update.""" - residual = hidden_states - aggregated, h_res, h_post = self.hyper_connection( + aggregated, h_res, h_post, residual = self.hyper_connection( hidden_states, mhc_recompute_manager=mhc_recompute_manager ) layer_output, context = self._call_inner_layer( diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 33d3e93c316..220d521ac90 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -23,7 +23,9 @@ native_sinkhorn, ) -_require_cutile = pytest.mark.skipif(not is_cutile_available(), reason="cuTile not installed") +_require_cutile = pytest.mark.skipif( + not is_cutile_available(), reason="cuTile unavailable for current device" +) _require_triton = pytest.mark.skipif(not is_triton_available(), reason="Triton not installed") From c72b3d8ce1ce48df84abafa479daf58a51ea9d68 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Sun, 31 May 2026 22:47:56 -0700 Subject: [PATCH 22/34] upd: fix : lint --- megatron/core/fusions/fused_mhc_kernels.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 48bb7ad3fa0..90031c5efe3 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -187,11 +187,7 @@ def _cutile_supports_current_device() -> bool: try: result = subprocess.run( - [tileiras, "--gpu-name", arch], - capture_output=True, - check=False, - text=True, - timeout=10, + [tileiras, "--gpu-name", arch], capture_output=True, check=False, text=True, timeout=10 ) except (OSError, subprocess.SubprocessError) as exc: _CUTILE_DEVICE_SUPPORT_ERROR = str(exc) From 988183979e0a8fe8b69abbfd252273a9ed984116 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Mon, 1 Jun 2026 16:49:36 +0800 Subject: [PATCH 23/34] test: update hybrid mhc mock tuple --- tests/unit_tests/models/test_hybrid_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index d17d244fa8f..325bcdc09ec 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -161,7 +161,7 @@ def fake_hyper_connection_forward(hidden_states, mhc_recompute_manager=None): aggregated = hidden_states.view(s, b, n, c).mean(dim=2) h_res = torch.empty(s, b, n, n, dtype=hidden_states.dtype) h_post = torch.empty(s, b, n, dtype=hidden_states.dtype) - return aggregated, h_res, h_post + return aggregated, h_res, h_post, hidden_states def fake_fused_h_res_h_post_bda( h_res, From 8cffcbda82aff2547102689c0e9a77cf8ddea91a Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 2 Jun 2026 13:36:49 +0800 Subject: [PATCH 24/34] Update megatron/core/fusions/fused_mhc_kernels.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> --- megatron/core/fusions/fused_mhc_kernels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 90031c5efe3..d52ee262ba2 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1150,8 +1150,8 @@ def _ct_hpb_bwd_g_x_orig_kernel( go_tile = ct.load( go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO ) - g_x_tile = ct.full((TILE_SIZE, 1, TILE_C), 0, dtype=hp.dtype) - g_orig_tile = ct.full((TILE_SIZE, N, TILE_C), 0, dtype=hp.dtype) + g_x_tile = ct.full((TILE_SIZE, 1, TILE_C), 0, dtype=ct.float32) + g_orig_tile = ct.full((TILE_SIZE, N, TILE_C), 0, dtype=ct.float32) for j in range(N): hp_j = ct.extract(hp_tile, (0, j), shape=(TILE_SIZE, 1)) hp_j_exp = ct.expand_dims(hp_j, axis=2) # [TS, 1, 1] From 833598c02f2767bb078633c555314b47c4188b09 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Tue, 2 Jun 2026 13:45:30 +0800 Subject: [PATCH 25/34] Fix native mHC fallback visibility --- megatron/core/fusions/fused_mhc_kernels.py | 12 +++++++++++- megatron/core/transformer/transformer_config.py | 6 +++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index d52ee262ba2..f641c64bd10 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -19,12 +19,13 @@ import os import shutil import subprocess +import warnings from typing import Optional, Tuple import torch from torch import Tensor -from megatron.core._rank_utils import log_single_rank +from megatron.core._rank_utils import log_single_rank, safe_get_rank logger = logging.getLogger(__name__) LOG2E = math.log2(math.e) @@ -2983,6 +2984,14 @@ def log_fused_mhc_backend_once() -> None: logging.WARNING if all_native else logging.INFO, f"[mHC] fused backend selection: {backend_selection}", ) + if all_native and safe_get_rank() == 0: + warnings.warn( + "[mHC] No accelerated mHC backend is available; falling back to native torch " + "implementations. The fallback is functionally equivalent, but may not provide " + "the performance benefits of fused mHC backends.", + UserWarning, + stacklevel=2, + ) def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: @@ -3053,6 +3062,7 @@ def _torch_h_post_bda_bwd( ) +@torch.compile def _torch_proj_rms_compute_h( x: Tensor, weight: Tensor, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index deefd9a6234..a1a09487ed7 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1081,7 +1081,11 @@ class TransformerConfig(ModelParallelConfig): H_aggregate, H_post_bda, ProjRms) with fused/autograd implementations for better performance on supported GPUs. Backend selection is internal and op-specific: Triton for Sinkhorn and H_post_bda backward when available, - cuTile for the remaining fused kernels when available, then torch fallback. + cuTile for the remaining fused kernels when available, then native torch + fallback. If every mHC operation uses the native torch fallback, + use_fused_mhc remains enabled and a rank-0 warning is emitted. The all-native + fallback is functionally equivalent, but may not provide fused backend + performance benefits. """ mhc_recompute_layer_num: Optional[int] = None From 03cfe02f17dec5aeeb350dbf7f3d0f75fa68dd93 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Tue, 2 Jun 2026 01:46:28 -0700 Subject: [PATCH 26/34] fix: golden value for mhc --- .../golden_values_dev_dgx_h100.json | 480 +++++++++--------- 1 file changed, 240 insertions(+), 240 deletions(-) diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json index 51c04e49ae3..b27e30dd333 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json @@ -4,55 +4,55 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 10.86155, - "2": 10.85472, - "3": 10.86697, + "1": 10.86149, + "2": 10.85475, + "3": 10.86695, "4": 10.84628, - "5": 10.8847, - "6": 10.89688, - "7": 10.87274, - "8": 10.86585, - "9": 10.86981, - "10": 10.83765, - "11": 10.89456, - "12": 10.87958, + "5": 10.88468, + "6": 10.89685, + "7": 10.87277, + "8": 10.86588, + "9": 10.86984, + "10": 10.83763, + "11": 10.89461, + "12": 10.87955, "13": 10.87687, - "14": 10.9036, - "15": 10.83116, - "16": 10.83448, - "17": 10.80064, - "18": 10.82072, - "19": 10.81455, - "20": 10.71812, - "21": 10.68626, - "22": 10.53195, - "23": 10.70489, - "24": 10.5856, - "25": 10.51905, - "26": 10.58493, - "27": 10.60103, - "28": 10.53538, - "29": 10.57117, - "30": 10.33243, - "31": 10.05838, - "32": 10.4279, - "33": 10.42024, + "14": 10.90354, + "15": 10.83117, + "16": 10.8345, + "17": 10.80067, + "18": 10.8207, + "19": 10.81458, + "20": 10.71809, + "21": 10.68629, + "22": 10.53196, + "23": 10.70487, + "24": 10.58558, + "25": 10.51904, + "26": 10.58488, + "27": 10.60102, + "28": 10.5354, + "29": 10.57113, + "30": 10.33245, + "31": 10.05837, + "32": 10.42792, + "33": 10.42027, "34": 10.16988, "35": 10.23073, - "36": 10.18762, + "36": 10.18757, "37": 10.31244, - "38": 10.14213, - "39": 10.38141, - "40": 10.04848, - "41": 10.10338, + "38": 10.14208, + "39": 10.38138, + "40": 10.0485, + "41": 10.10336, "42": 10.17158, - "43": 9.78294, + "43": 9.78293, "44": 9.90959, - "45": 9.78513, - "46": 9.76889, - "47": 10.10085, + "45": 9.78518, + "46": 9.76888, + "47": 10.10084, "48": 9.80968, - "49": 9.4878, + "49": 9.48782, "50": 9.86709 } }, @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 1654.0, - "2": 1741.0, - "3": 1678.0, - "4": 1808.0, - "5": 1876.0, - "6": 1898.0, - "7": 1791.0, - "8": 1592.0, - "9": 1864.0, - "10": 1426.0, - "11": 1913.0, - "12": 1700.0, - "13": 1841.0, - "14": 1755.0, - "15": 1891.0, - "16": 1912.0, - "17": 1925.0, - "18": 1645.0, - "19": 1827.0, - "20": 1693.0, - "21": 1814.0, - "22": 1634.0, - "23": 1991.0, - "24": 1687.0, - "25": 1646.0, - "26": 1667.0, - "27": 1786.0, - "28": 2010.0, - "29": 2022.0, - "30": 1897.0, - "31": 1604.0, - "32": 1839.0, - "33": 2184.0, - "34": 2016.0, - "35": 1933.0, - "36": 1882.0, - "37": 2467.0, - "38": 2172.0, - "39": 2422.0, - "40": 2209.0, - "41": 2354.0, - "42": 2327.0, - "43": 1996.0, - "44": 2245.0, - "45": 2156.0, - "46": 2316.0, - "47": 2465.0, - "48": 2450.0, - "49": 2378.0, - "50": 2320.0 + "1": 34399.0, + "2": 5068.0, + "3": 5139.0, + "4": 2133.0, + "5": 34705.0, + "6": 5089.0, + "7": 5475.0, + "8": 5376.0, + "9": 5624.0, + "10": 34196.0, + "11": 1824.0, + "12": 2245.0, + "13": 5078.0, + "14": 1770.0, + "15": 2066.0, + "16": 34671.0, + "17": 4463.0, + "18": 1721.0, + "19": 2101.0, + "20": 1708.0, + "21": 5367.0, + "22": 34462.0, + "23": 2073.0, + "24": 1882.0, + "25": 1990.0, + "26": 4731.0, + "27": 2311.0, + "28": 1957.0, + "29": 5633.0, + "30": 34746.0, + "31": 1810.0, + "32": 34706.0, + "33": 5781.0, + "34": 5295.0, + "35": 4581.0, + "36": 2350.0, + "37": 35156.0, + "38": 2200.0, + "39": 5912.0, + "40": 6668.0, + "41": 5213.0, + "42": 2898.0, + "43": 2068.0, + "44": 35053.0, + "45": 5016.0, + "46": 6505.0, + "47": 2949.0, + "48": 6075.0, + "49": 3332.0, + "50": 35248.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 555865600.0, - "2": 555865600.0, - "3": 555865600.0, - "4": 555865600.0, - "5": 555865600.0, - "6": 555865600.0, - "7": 555865600.0, - "8": 555865600.0, - "9": 555865600.0, - "10": 555865600.0, - "11": 555865600.0, - "12": 555865600.0, - "13": 555865600.0, - "14": 555865600.0, - "15": 555865600.0, - "16": 555865600.0, - "17": 555865600.0, - "18": 555865600.0, - "19": 555865600.0, - "20": 555865600.0, - "21": 555865600.0, - "22": 555865600.0, - "23": 555865600.0, - "24": 555865600.0, - "25": 555865600.0, - "26": 555865600.0, - "27": 555865600.0, - "28": 555865600.0, - "29": 555865600.0, - "30": 555865600.0, - "31": 555865600.0, - "32": 555865600.0, - "33": 555865600.0, - "34": 555865600.0, - "35": 555865600.0, - "36": 555865600.0, - "37": 555865600.0, - "38": 555865600.0, - "39": 555865600.0, - "40": 555865600.0, - "41": 555865600.0, - "42": 555865600.0, - "43": 555865600.0, - "44": 555865600.0, - "45": 555865600.0, - "46": 555865600.0, - "47": 555865600.0, - "48": 555865600.0, - "49": 555865600.0, - "50": 555865600.0 + "1": 556914176.0, + "2": 556914176.0, + "3": 556914176.0, + "4": 556914176.0, + "5": 556914176.0, + "6": 556914176.0, + "7": 556914176.0, + "8": 556914176.0, + "9": 556914176.0, + "10": 557962752.0, + "11": 556914176.0, + "12": 556914176.0, + "13": 556914176.0, + "14": 556914176.0, + "15": 556914176.0, + "16": 556914176.0, + "17": 556914176.0, + "18": 556914176.0, + "19": 556914176.0, + "20": 556914176.0, + "21": 556914176.0, + "22": 556914176.0, + "23": 556914176.0, + "24": 556914176.0, + "25": 556914176.0, + "26": 557962752.0, + "27": 556914176.0, + "28": 556914176.0, + "29": 557962752.0, + "30": 556914176.0, + "31": 556914176.0, + "32": 556914176.0, + "33": 556914176.0, + "34": 557962752.0, + "35": 556914176.0, + "36": 556914176.0, + "37": 556914176.0, + "38": 556914176.0, + "39": 556914176.0, + "40": 557962752.0, + "41": 556914176.0, + "42": 556914176.0, + "43": 556914176.0, + "44": 556914176.0, + "45": 556914176.0, + "46": 556914176.0, + "47": 556914176.0, + "48": 556914176.0, + "49": 556914176.0, + "50": 556914176.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 1746955776.0, - "2": 1934879232.0, - "3": 1934879232.0, - "4": 1934879232.0, - "5": 1934879232.0, - "6": 1934879232.0, - "7": 1934879232.0, - "8": 1934879232.0, - "9": 1934879232.0, - "10": 1934879232.0, - "11": 1934879232.0, - "12": 1934879232.0, - "13": 1934879232.0, - "14": 1934879232.0, - "15": 1934879232.0, - "16": 1934879232.0, - "17": 1934879232.0, - "18": 1934879232.0, - "19": 1934879232.0, - "20": 1934879232.0, - "21": 1934879232.0, - "22": 1934879232.0, - "23": 1934879232.0, - "24": 1934879232.0, - "25": 1934879232.0, - "26": 1935927296.0, - "27": 1935927296.0, - "28": 1935927296.0, - "29": 1935927296.0, - "30": 1935927296.0, - "31": 1935927296.0, - "32": 1935927296.0, - "33": 1935927296.0, - "34": 1935927808.0, - "35": 1935927808.0, - "36": 1935927808.0, - "37": 1935927808.0, - "38": 1935927808.0, - "39": 1935927808.0, - "40": 1935927808.0, - "41": 1935927808.0, - "42": 1935927808.0, - "43": 1935927808.0, - "44": 1935927808.0, - "45": 1935927808.0, - "46": 1935927808.0, - "47": 1935927808.0, - "48": 1935927808.0, - "49": 1935927808.0, - "50": 1935927808.0 + "1": 1746399232.0, + "2": 1937075200.0, + "3": 1937075200.0, + "4": 1937075200.0, + "5": 1937075200.0, + "6": 1937075200.0, + "7": 1937075200.0, + "8": 1937075200.0, + "9": 1937075200.0, + "10": 1937075200.0, + "11": 1937075200.0, + "12": 1937075200.0, + "13": 1937075200.0, + "14": 1937075200.0, + "15": 1937075200.0, + "16": 1937075200.0, + "17": 1937075200.0, + "18": 1937075200.0, + "19": 1937075200.0, + "20": 1937075200.0, + "21": 1937075200.0, + "22": 1937075200.0, + "23": 1937075200.0, + "24": 1937075200.0, + "25": 1937075200.0, + "26": 1937075200.0, + "27": 1937075200.0, + "28": 1937075200.0, + "29": 1937075200.0, + "30": 1937075200.0, + "31": 1937075200.0, + "32": 1937075200.0, + "33": 1937075200.0, + "34": 1937075200.0, + "35": 1937075200.0, + "36": 1937075200.0, + "37": 1937075200.0, + "38": 1937075200.0, + "39": 1937075200.0, + "40": 1937075200.0, + "41": 1937075200.0, + "42": 1937075200.0, + "43": 1937075200.0, + "44": 1937075200.0, + "45": 1937075200.0, + "46": 1937075200.0, + "47": 1937075200.0, + "48": 1937075200.0, + "49": 1937075200.0, + "50": 1937075200.0 } }, "iteration-time": { @@ -233,55 +233,55 @@ "step_interval": 1, "values": { "1": "nan", - "2": 30.16457, - "3": 0.69331, - "4": 0.69236, - "5": 0.69011, - "6": 0.69431, - "7": 0.68739, - "8": 0.68382, - "9": 0.68467, - "10": 0.68772, - "11": 0.68456, - "12": 0.68193, - "13": 0.69128, - "14": 0.68205, - "15": 0.68208, - "16": 0.68686, - "17": 0.68599, - "18": 0.68567, - "19": 0.69039, - "20": 0.68754, - "21": 0.69693, - "22": 0.69584, - "23": 0.69108, - "24": 0.68558, - "25": 0.68807, - "26": 0.71618, - "27": 0.69064, - "28": 0.7033, - "29": 0.70268, - "30": 0.69756, - "31": 0.70299, - "32": 0.69313, - "33": 0.69455, - "34": 0.70319, - "35": 0.68861, - "36": 0.68994, - "37": 0.6851, - "38": 0.68942, - "39": 0.68364, - "40": 0.67981, - "41": 0.67851, - "42": 0.68629, - "43": 0.69028, - "44": 0.68641, - "45": 0.68513, - "46": 0.68077, - "47": 0.68391, - "48": 0.68376, - "49": 0.68312, - "50": 0.68105 + "2": 23.68682, + "3": 0.49278, + "4": 0.48539, + "5": 0.48882, + "6": 0.48719, + "7": 0.48515, + "8": 0.48616, + "9": 0.48409, + "10": 0.48496, + "11": 0.48688, + "12": 0.48809, + "13": 0.48532, + "14": 0.48548, + "15": 0.48618, + "16": 0.48449, + "17": 0.48545, + "18": 0.4894, + "19": 0.48704, + "20": 0.48687, + "21": 0.48792, + "22": 0.50768, + "23": 0.48897, + "24": 0.48564, + "25": 0.48418, + "26": 0.49874, + "27": 0.4869, + "28": 0.48559, + "29": 0.48507, + "30": 0.48655, + "31": 0.49426, + "32": 0.49577, + "33": 0.48641, + "34": 0.4864, + "35": 0.48522, + "36": 0.48214, + "37": 0.4844, + "38": 0.48496, + "39": 0.48722, + "40": 0.48728, + "41": 0.48777, + "42": 0.48522, + "43": 0.48574, + "44": 0.48593, + "45": 0.48608, + "46": 0.48732, + "47": 0.48832, + "48": 0.48904, + "49": 0.48735, + "50": 0.48544 } } } \ No newline at end of file From 53e2d36cea6d117f8f67732c5c37a2ff1904e10e Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Wed, 3 Jun 2026 14:56:10 +0800 Subject: [PATCH 27/34] Align mHC compute_h eps and Sinkhorn semantics --- megatron/core/fusions/fused_mhc_kernels.py | 116 ++++++++++++++---- megatron/core/transformer/hyper_connection.py | 27 ++-- .../fusions/test_fused_mhc_kernels.py | 16 ++- 3 files changed, 122 insertions(+), 37 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index f641c64bd10..97bc39ab07e 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -244,11 +244,16 @@ def _triton_sinkhorn_fwd_kernel( M = tl.exp2((logits - row_max[:, None]) * TLOG2E) tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) - for _ in range(NUM_ITERS): + if NUM_ITERS > 0: row_sum = tl.sum(M, axis=1) - M = M / (row_sum[:, None] + eps) + M = M / row_sum[:, None] + eps col_sum = tl.sum(M, axis=0) M = M / (col_sum[None, :] + eps) + for _ in range(NUM_ITERS - 1): + row_sum = tl.sum(M, axis=1) + M = M / (row_sum[:, None] + eps) + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) @@ -287,7 +292,10 @@ def _triton_sinkhorn_bwd_kernel( row_sum = tl.sum(M, axis=1) tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) - M = M / (row_sum[:, None] + eps) + if t == 0: + M = M / row_sum[:, None] + eps + else: + M = M / (row_sum[:, None] + eps) ws_off = M_ws_base + (2 * t + 1) * HC * HC tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) @@ -315,8 +323,12 @@ def _triton_sinkhorn_bwd_kernel( ).to(tl.float32) row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) - grad = grad / (row_s[:, None] + eps) - row_corr = tl.sum(grad * M, axis=1) + if t == 0: + grad = grad / row_s[:, None] + row_corr = tl.sum(grad * (M - eps), axis=1) + else: + grad = grad / (row_s[:, None] + eps) + row_corr = tl.sum(grad * M, axis=1) grad = grad - row_corr[:, None] M = tl.load( ws_M_ptr + M_ws_base + (2 * t) * HC * HC + offs_r[:, None] * HC + offs_c[None, :] @@ -364,7 +376,7 @@ class TritonFusedSinkhorn(torch.autograd.Function): def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): """Run Triton Sinkhorn forward and save initial matrix for backward.""" out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) - ctx.save_for_backward(M_init) + ctx.save_for_backward(input_logits if num_iterations == 0 else M_init) ctx.num_iterations = num_iterations ctx.eps = eps return out @@ -372,7 +384,15 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output: Tensor): """Run Triton Sinkhorn backward.""" - (M_init,) = ctx.saved_tensors + (saved,) = ctx.saved_tensors + if ctx.num_iterations == 0: + with torch.enable_grad(): + logits = saved.detach().requires_grad_(True) + row_max = logits.max(dim=-1, keepdim=True).values + M = torch.exp(logits - row_max) + M.backward(grad_output) + return logits.grad, None, None + M_init = saved grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None @@ -835,11 +855,16 @@ def _ct_sinkhorn_fwd_kernel( index=(pid, 0, 0), tile=ct.reshape(M.astype(M_init_out.dtype), (TILE_SIZE, HC, HC)), ) - for _ in range(NUM_ITERS): + if NUM_ITERS > 0: row_sum = ct.sum(M, axis=2, keepdims=True) - M = M / (row_sum + eps) + M = M / row_sum + eps col_sum = ct.sum(M, axis=1, keepdims=True) M = M / (col_sum + eps) + for _ in range(NUM_ITERS - 1): + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / (row_sum + eps) + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) ct.store(out, index=(pid, 0, 0), tile=ct.reshape(M.astype(out.dtype), (TILE_SIZE, HC, HC))) @ct.kernel @@ -864,7 +889,10 @@ def _ct_sinkhorn_bwd_kernel( ct.store(ws_M, index=(M_base + 2 * t, 0, 0), tile=M) row_sum = ct.sum(M, axis=2, keepdims=True) ct.store(ws_rs, index=(v_base + t, 0, 0), tile=row_sum) - M = M / (row_sum + eps) + if t == 0: + M = M / row_sum + eps + else: + M = M / (row_sum + eps) ct.store(ws_M, index=(M_base + 2 * t + 1, 0, 0), tile=M) col_sum = ct.sum(M, axis=1, keepdims=True) ct.store(ws_cs, index=(v_base + t, 0, 0), tile=col_sum) @@ -879,8 +907,12 @@ def _ct_sinkhorn_bwd_kernel( grad = grad - col_corr M = ct.load(ws_M, index=(M_base + 2 * t + 1, 0, 0), shape=(TILE_SIZE, HC, HC)) row_s = ct.load(ws_rs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, HC, 1)) - grad = grad / (row_s + eps) - row_corr = ct.sum(grad * M, axis=2, keepdims=True) + if t == 0: + grad = grad / row_s + row_corr = ct.sum(grad * (M - eps), axis=2, keepdims=True) + else: + grad = grad / (row_s + eps) + row_corr = ct.sum(grad * M, axis=2, keepdims=True) grad = grad - row_corr M = ct.load(ws_M, index=(M_base + 2 * t, 0, 0), shape=(TILE_SIZE, HC, HC)) grad = grad * M @@ -1469,6 +1501,7 @@ def _ct_reduce_compute_h_kernel( K: int, n: ConstInt, eps: float, + compute_h_eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, SPLIT_K: ConstInt, @@ -1526,8 +1559,8 @@ def _ct_reduce_compute_h_kernel( h_pre_linear = pre_accum * alpha_pre * inv_r_eps + bias_pre h_post_linear = post_accum * alpha_post * inv_r_eps + bias_post - h_pre = _ct_sigmoid(h_pre_linear) - h_post = _ct_sigmoid(h_post_linear) * 2.0 + h_pre = _ct_sigmoid(h_pre_linear) + compute_h_eps + h_post = (_ct_sigmoid(h_post_linear) + compute_h_eps) * 2.0 ct.store(H_PRE, index=(bid_m, 0), tile=h_pre.astype(H_PRE.dtype)) ct.store(H_POST, index=(bid_m, 0), tile=h_post.astype(H_POST.dtype)) @@ -1872,6 +1905,7 @@ def _cutile_reduce_compute_h( N: int, K: int, eps: float, + compute_h_eps: float, _proj_tile_m: int, tile_n: int, split_k: int, @@ -1918,6 +1952,7 @@ def _make_args(tm): K, n, eps, + compute_h_eps, tm, tile_n, split_k, @@ -1956,6 +1991,7 @@ def _cutile_proj_rms_compute_h_fwd( alpha_res: Tensor, n: int, eps: float, + compute_h_eps: float, ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Fused proj_rms + compute_h forward. @@ -2101,6 +2137,7 @@ def _cutile_proj_rms_compute_h_fwd( N, K, eps, + compute_h_eps, tm, TILE_N, split_k, @@ -2235,6 +2272,7 @@ def _ct_fused_grad_h_proj_kernel( N: int, n: ConstInt, eps: float, + compute_h_eps: float, TILE_SIZE_M: ConstInt, TILE_SIZE_N: ConstInt, HAS_GRAD_H_PRE: ConstInt, @@ -2278,7 +2316,8 @@ def _ct_fused_grad_h_proj_kernel( PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) proj_pre = ct.astype(proj_pre, ct.float32) - grad_h_pre = gy_pre * h_pre * (1.0 - h_pre) + sigmoid_pre = h_pre - compute_h_eps + grad_h_pre = gy_pre * sigmoid_pre * (1.0 - sigmoid_pre) grad_proj_pre = grad_h_pre * alpha_pre * inv_r_eps grad_r_from_h += ct.sum( grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True @@ -2301,8 +2340,8 @@ def _ct_fused_grad_h_proj_kernel( PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) proj_post = ct.astype(proj_post, ct.float32) - half_h_post = h_post * 0.5 - grad_h_post = gy_post * half_h_post * (1.0 - half_h_post) * 2.0 + sigmoid_post = h_post * 0.5 - compute_h_eps + grad_h_post = gy_post * sigmoid_post * (1.0 - sigmoid_post) * 2.0 grad_proj_post = grad_h_post * alpha_post * inv_r_eps grad_r_from_h += ct.sum( grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True @@ -2674,6 +2713,7 @@ def _cutile_fused_compute_h_proj_rms_bwd( bias: Tensor, n: int, eps: float, + compute_h_eps: float, ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: """Fused compute_h + proj_rms backward. @@ -2733,6 +2773,7 @@ def _cutile_fused_compute_h_proj_rms_bwd( N, n, eps, + compute_h_eps, tile_m_precomp, TILE_N, int(has_grad_h_pre), @@ -3072,6 +3113,7 @@ def _torch_proj_rms_compute_h( bias: Tensor, n: int, eps: float, + compute_h_eps: float = 1e-6, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: proj = torch.matmul(x, weight.t()) r = x.norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) @@ -3080,8 +3122,8 @@ def _torch_proj_rms_compute_h( dim=-1, ) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) - h_pre = h[..., :n].sigmoid() - h_post = h[..., n : 2 * n].sigmoid() * 2 + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = (h[..., n : 2 * n].sigmoid() + compute_h_eps) * 2 h_res = h[..., 2 * n :] return h_pre, h_post, h_res, r @@ -3095,7 +3137,7 @@ class CutileSinkhornKnopp(torch.autograd.Function): def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): """Run cuTile Sinkhorn forward and save initial matrix for backward.""" output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) - ctx.save_for_backward(M_init) + ctx.save_for_backward(input_logits if num_iterations == 0 else M_init) ctx.num_iterations = num_iterations ctx.eps = eps return output @@ -3103,7 +3145,15 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output): """Run cuTile Sinkhorn backward.""" - (M_init,) = ctx.saved_tensors + (saved,) = ctx.saved_tensors + if ctx.num_iterations == 0: + with torch.enable_grad(): + logits = saved.detach().requires_grad_(True) + row_max = logits.max(dim=-1, keepdim=True).values + M = torch.exp(logits - row_max) + M.backward(grad_output) + return logits.grad, None, None + M_init = saved grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None @@ -3155,10 +3205,11 @@ def forward( bias: Tensor, n: int, eps: float = 1e-6, + compute_h_eps: float = 1e-6, ): """Run fused cuTile projection, RMS normalization, and compute_h forward.""" h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( - x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps + x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, compute_h_eps ) ctx.save_for_backward( x, @@ -3175,6 +3226,7 @@ def forward( ) ctx.n = n ctx.eps = eps + ctx.compute_h_eps = compute_h_eps return h_pre, h_post, h_res, r @staticmethod @@ -3213,10 +3265,21 @@ def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): bias_param, ctx.n, ctx.eps, + ctx.compute_h_eps, ) ) - return grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None + return ( + grad_x, + grad_weight, + grad_ap, + grad_apo, + grad_ar, + grad_bias, + None, + None, + None, + ) class FusedHAggregate(torch.autograd.Function): @@ -3335,11 +3398,14 @@ def fused_proj_rms_compute_h( bias: Tensor, n: int, eps: float = 1e-6, + compute_h_eps: float = 1e-6, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" _raise_mhc_backend_validation_error() if is_cutile_available(): return CutileProjRmsComputeH.apply( - x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps ) - return _torch_proj_rms_compute_h(x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps) + return _torch_proj_rms_compute_h( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index e33b62a9d6e..e1a2bd8331a 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -15,14 +15,22 @@ if TYPE_CHECKING: from megatron.core.tensor_parallel.random import CheckpointManager +_MHC_SINKHORN_EPS = 1e-6 +_MHC_COMPUTE_H_EPS = 1e-6 + @torch.compile def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: row_max = input_logits.max(dim=-1, keepdim=True).values M = torch.exp(input_logits - row_max) - for _ in range(num_iterations): - M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps) - M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps) + if num_iterations == 0: + return M + + M = M / M.sum(dim=-1, keepdim=True) + eps + M = M / (M.sum(dim=-2, keepdim=True) + eps) + for _ in range(num_iterations - 1): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) return M @@ -174,6 +182,8 @@ def __init__(self, config: TransformerConfig, layer_number: int): self.n = config.num_residual_streams self.hidden_size = config.hidden_size self.sinkhorn_iterations = config.mhc_sinkhorn_iterations + self.sinkhorn_eps = _MHC_SINKHORN_EPS + self.compute_h_eps = _MHC_COMPUTE_H_EPS # Projection weights for dynamic mappings # Input: [s, b, n*C] -> Output: n^2 + 2n values per token @@ -277,10 +287,10 @@ def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: ) h = r * proj * alpha_ + self.bias # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre) - h_pre = h[..., : self.n].sigmoid() # [s, b, n] + h_pre = h[..., : self.n].sigmoid() + self.compute_h_eps # [s, b, n] - # H_post = 2σ(α_post * (θ_post @ x̃) + b_post) - h_post = h[..., self.n : 2 * self.n].sigmoid() * 2 # [s, b, n] + # H_post = 2(σ(α_post * (θ_post @ x̃) + b_post) + eps) + h_post = (h[..., self.n : 2 * self.n].sigmoid() + self.compute_h_eps) * 2 h_res = h[..., 2 * self.n :] return h_pre, h_post, h_res @@ -314,6 +324,7 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: self.bias, self.n, self.norm_eps, + self.compute_h_eps, ) h_pre = h_pre.view(s, b, self.n) h_post = h_post.view(s, b, self.n) @@ -326,7 +337,9 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: h_pre, h_post, h_res = self._compute_h(proj, r) h_res = h_res.view(s, b, self.n, self.n) - h_res = self._sinkhorn_op(h_res, self.sinkhorn_iterations, self.norm_eps) # [s, b, n, n] + h_res = self._sinkhorn_op( + h_res, self.sinkhorn_iterations, self.sinkhorn_eps + ) # [s, b, n, n] return h_pre, h_post, h_res diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 220d521ac90..2c5662e8dc8 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -79,9 +79,14 @@ def _info(): def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: row_max = logits.max(dim=-1, keepdim=True).values M = torch.exp(logits - row_max) - for _ in range(num_iters): - M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps) - M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps) + if num_iters == 0: + return M + + M = M / M.sum(dim=-1, keepdim=True) + eps + M = M / (M.sum(dim=-2, keepdim=True) + eps) + for _ in range(num_iters - 1): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) return M @@ -129,6 +134,7 @@ def _ref_proj_rms_compute_h( bias: Tensor, n: int, eps: float = 1e-6, + compute_h_eps: float = 1e-6, ): """Reference: fused proj_rms + compute_h.""" proj = torch.matmul(x, weight.t()) @@ -138,8 +144,8 @@ def _ref_proj_rms_compute_h( N = proj.shape[-1] alpha = torch.cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n)]) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) - h_pre = h[..., :n].sigmoid() - h_post = h[..., n : 2 * n].sigmoid() * 2 + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = (h[..., n : 2 * n].sigmoid() + compute_h_eps) * 2 h_res = h[..., 2 * n :] return h_pre, h_post, h_res, r From dafa5096f6f3bf0c32d8eb61e2ec2207152e681f Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Wed, 3 Jun 2026 16:01:02 +0800 Subject: [PATCH 28/34] Fix mHC fused backend availability and precision --- megatron/core/fusions/fused_mhc_kernels.py | 54 +++++++++++++------ pyproject.toml | 1 + .../fusions/test_fused_mhc_kernels.py | 17 +++++- 3 files changed, 55 insertions(+), 17 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 97bc39ab07e..e057ede2112 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -175,8 +175,9 @@ def _cutile_supports_current_device() -> bool: return _CUTILE_DEVICE_SUPPORT_CACHE if not torch.cuda.is_available(): - _CUTILE_DEVICE_SUPPORT_CACHE = True - return True + _CUTILE_DEVICE_SUPPORT_ERROR = "CUDA is not available" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False major, minor = torch.cuda.get_device_capability() arch = f"sm_{major}{minor}" @@ -339,7 +340,7 @@ def _triton_sinkhorn_bwd_kernel( tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) def _triton_sinkhorn_fwd( - input_logits: Tensor, num_iterations: int, eps: float = 1e-8 + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 ) -> Tuple[Tensor, Tensor]: original_shape = input_logits.shape hc = original_shape[-1] @@ -352,7 +353,7 @@ def _triton_sinkhorn_fwd( return out.view(original_shape), M_init.view(original_shape) def _triton_sinkhorn_bwd( - grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 ) -> Tensor: original_shape = grad_output.shape hc = original_shape[-1] @@ -744,6 +745,7 @@ def _triton_hpb_bwd_g_hp_hr_kernel( mask=mask_s[:, None], ) + # N is expected to stay small for mHC, so this simple extraction is acceptable. nn_offs = tl.arange(0, N * N) for i in tl.static_range(N): for j in tl.static_range(N): @@ -831,6 +833,22 @@ def _triton_h_post_bda_bwd( g_bias, ) +_TRITON_IMPLS = ( + { + "sinkhorn": triton_fused_sinkhorn, + "h_aggregate_fwd": _triton_h_aggregate_fwd, + "h_post_bda_fwd": _triton_h_post_bda_fwd, + "h_post_bda_bwd": _triton_h_post_bda_bwd, + } + if _TRITON_AVAILABLE + else { + "sinkhorn": None, + "h_aggregate_fwd": None, + "h_post_bda_fwd": None, + "h_post_bda_bwd": None, + } +) + # ============================================================================ # CuTile implementations (only defined when cuda.tile is available) @@ -928,7 +946,7 @@ def _sinkhorn_autotune_tile_sizes(N_batch): _sinkhorn_bwd_best_cfg: dict = {} def _cutile_sinkhorn_fwd( - input_logits: Tensor, num_iterations: int, eps: float = 1e-8 + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 ) -> Tuple[Tensor, Tensor]: original_shape = input_logits.shape hc = original_shape[-1] @@ -975,7 +993,7 @@ def _cutile_sinkhorn_fwd( return out.view(original_shape), M_init.view(original_shape) def _cutile_sinkhorn_bwd( - grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 ) -> Tensor: original_shape = grad_output.shape hc = original_shape[-1] @@ -1778,8 +1796,8 @@ def _cutile_proj_rms_fwd( _ct_proj_rms_fwd_kernel, (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), ) - proj = proj.view(split_k, M, N).sum(dim=0).to(dtype=x.dtype) - norm = norm.view(split_k, M, 1).sum(dim=0).to(dtype=x.dtype) + proj = proj.view(split_k, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(split_k, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) else: # Autotune on first call for this shape. from types import SimpleNamespace @@ -1799,8 +1817,8 @@ def _cutile_proj_rms_fwd( _ct_proj_rms_fwd_kernel, (x, weight, proj, norm, r, M, N, K, eps, tm, tn, tk, split_k), ) - proj = proj.view(split_k, M, N).sum(dim=0) - norm = norm.view(split_k, M, 1).sum(dim=0) + proj = proj.view(split_k, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(split_k, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) else: mx_split_k = max(cfg.SPLIT_K for cfg in configs) proj = torch.empty(mx_split_k * M, N, dtype=x.dtype, device=dev) @@ -1861,8 +1879,12 @@ def _cutile_proj_rms_fwd( ), ) - proj = proj.view(best.SPLIT_K, M, N).sum(dim=0) - norm = norm.view(best.SPLIT_K, M, 1).sum(dim=0) + proj = ( + proj.view(best.SPLIT_K, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + ) + norm = ( + norm.view(best.SPLIT_K, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + ) norm = torch.sqrt(norm) r = 1.0 / (norm / math.sqrt(K) + eps) return proj, norm, r @@ -3043,25 +3065,25 @@ def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: def _get_triton_sinkhorn(): if not _TRITON_AVAILABLE: return None - return globals().get("triton_fused_sinkhorn") + return _TRITON_IMPLS["sinkhorn"] def _get_triton_h_aggregate_fwd(): if not _TRITON_AVAILABLE: return None - return globals().get("_triton_h_aggregate_fwd") + return _TRITON_IMPLS["h_aggregate_fwd"] def _get_triton_h_post_bda_fwd(): if not _TRITON_AVAILABLE: return None - return globals().get("_triton_h_post_bda_fwd") + return _TRITON_IMPLS["h_post_bda_fwd"] def _get_triton_h_post_bda_bwd(): if not _TRITON_AVAILABLE: return None - return globals().get("_triton_h_post_bda_bwd") + return _TRITON_IMPLS["h_post_bda_bwd"] def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tuple[Tensor, Tensor]: diff --git a/pyproject.toml b/pyproject.toml index 2fb61c6f36c..467e7c82e28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,6 +236,7 @@ markers = [ "internal: mark a test as a test to private/internal functions.", "flaky: mark flaky tests for LTS environment", "flaky_in_dev: mark flaky tests for DEV environment", + "no_cuda_required: mark tests that should run without CUDA", ] [tool.coverage.run] diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 2c5662e8dc8..059e0a49faa 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -15,6 +15,7 @@ import torch from torch import Tensor +import megatron.core.fusions.fused_mhc_kernels as mhc_kernels from megatron.core.fusions.fused_mhc_kernels import is_cutile_available, is_triton_available from megatron.core.transformer.hyper_connection import ( native_h_aggregate, @@ -30,7 +31,9 @@ @pytest.fixture(autouse=True) -def _skip_without_cuda(): +def _skip_without_cuda(request): + if request.node.get_closest_marker("no_cuda_required"): + return if not torch.cuda.is_available(): pytest.skip("CUDA not available") @@ -71,6 +74,18 @@ def _info(): print(f"\n [backend: {backend}]") +@pytest.mark.no_cuda_required +def test_cutile_unavailable_without_cuda(monkeypatch): + monkeypatch.setattr(mhc_kernels, "_CUTILE_AVAILABLE", True) + monkeypatch.setattr(mhc_kernels, "_CUTILE_DEVICE_SUPPORT_CACHE", None) + monkeypatch.setattr(mhc_kernels, "_CUTILE_DEVICE_SUPPORT_ERROR", None) + monkeypatch.setattr(mhc_kernels.torch.cuda, "is_available", lambda: False) + + assert not mhc_kernels.is_cutile_available() + assert mhc_kernels._CUTILE_DEVICE_SUPPORT_CACHE is False + assert mhc_kernels._CUTILE_DEVICE_SUPPORT_ERROR == "CUDA is not available" + + # ============================================================================ # Pure-PyTorch differentiable references (used by both fwd AND bwd tests) # ============================================================================ From 395d7a21a5b72ec38dfd0f1eee3459fd32bc27ac Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Wed, 3 Jun 2026 16:03:48 +0800 Subject: [PATCH 29/34] Remove no-CUDA mHC test marker --- pyproject.toml | 1 - .../fusions/test_fused_mhc_kernels.py | 17 +---------------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 467e7c82e28..2fb61c6f36c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,7 +236,6 @@ markers = [ "internal: mark a test as a test to private/internal functions.", "flaky: mark flaky tests for LTS environment", "flaky_in_dev: mark flaky tests for DEV environment", - "no_cuda_required: mark tests that should run without CUDA", ] [tool.coverage.run] diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 059e0a49faa..2c5662e8dc8 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -15,7 +15,6 @@ import torch from torch import Tensor -import megatron.core.fusions.fused_mhc_kernels as mhc_kernels from megatron.core.fusions.fused_mhc_kernels import is_cutile_available, is_triton_available from megatron.core.transformer.hyper_connection import ( native_h_aggregate, @@ -31,9 +30,7 @@ @pytest.fixture(autouse=True) -def _skip_without_cuda(request): - if request.node.get_closest_marker("no_cuda_required"): - return +def _skip_without_cuda(): if not torch.cuda.is_available(): pytest.skip("CUDA not available") @@ -74,18 +71,6 @@ def _info(): print(f"\n [backend: {backend}]") -@pytest.mark.no_cuda_required -def test_cutile_unavailable_without_cuda(monkeypatch): - monkeypatch.setattr(mhc_kernels, "_CUTILE_AVAILABLE", True) - monkeypatch.setattr(mhc_kernels, "_CUTILE_DEVICE_SUPPORT_CACHE", None) - monkeypatch.setattr(mhc_kernels, "_CUTILE_DEVICE_SUPPORT_ERROR", None) - monkeypatch.setattr(mhc_kernels.torch.cuda, "is_available", lambda: False) - - assert not mhc_kernels.is_cutile_available() - assert mhc_kernels._CUTILE_DEVICE_SUPPORT_CACHE is False - assert mhc_kernels._CUTILE_DEVICE_SUPPORT_ERROR == "CUDA is not available" - - # ============================================================================ # Pure-PyTorch differentiable references (used by both fwd AND bwd tests) # ============================================================================ From 512e56e3f0d0c0baf6ef1a3d3cefabfaa28f927f Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Wed, 3 Jun 2026 16:33:55 +0800 Subject: [PATCH 30/34] Align cuTile proj RMS eps defaults --- megatron/core/fusions/fused_mhc_kernels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index e057ede2112..362b2a76492 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -1766,7 +1766,7 @@ def _default_proj_rms_fwd_config(M: int, K: int, TILE_N: int): _proj_rms_fwd_best_cfg: dict = {} def _cutile_proj_rms_fwd( - x: Tensor, weight: Tensor, eps: float = 1e-8 + x: Tensor, weight: Tensor, eps: float = 1e-6 ) -> Tuple[Tensor, Tensor, Tensor]: M, K = x.shape N = weight.shape[0] @@ -2183,7 +2183,7 @@ def _cutile_proj_rms_bwd( x: Tensor, weight: Tensor, norm: Tensor, - eps: float = 1e-8, + eps: float = 1e-6, ) -> Tuple[Tensor, Tensor]: M, K = x.shape N = weight.shape[0] From 7a5ae5467971c7e9dced4139f635583da6dc4f37 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Wed, 3 Jun 2026 01:58:22 -0700 Subject: [PATCH 31/34] fix: lint and golden values --- megatron/core/fusions/fused_mhc_kernels.py | 28 +- .../golden_values_dev_dgx_h100.json | 490 +++++++++--------- 2 files changed, 250 insertions(+), 268 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 362b2a76492..8c0f25ebedc 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -833,6 +833,7 @@ def _triton_h_post_bda_bwd( g_bias, ) + _TRITON_IMPLS = ( { "sinkhorn": triton_fused_sinkhorn, @@ -841,12 +842,7 @@ def _triton_h_post_bda_bwd( "h_post_bda_bwd": _triton_h_post_bda_bwd, } if _TRITON_AVAILABLE - else { - "sinkhorn": None, - "h_aggregate_fwd": None, - "h_post_bda_fwd": None, - "h_post_bda_bwd": None, - } + else {"sinkhorn": None, "h_aggregate_fwd": None, "h_post_bda_fwd": None, "h_post_bda_bwd": None} ) @@ -1879,12 +1875,8 @@ def _cutile_proj_rms_fwd( ), ) - proj = ( - proj.view(best.SPLIT_K, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) - ) - norm = ( - norm.view(best.SPLIT_K, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) - ) + proj = proj.view(best.SPLIT_K, M, N).to(torch.float32).sum(dim=0).to(dtype=x.dtype) + norm = norm.view(best.SPLIT_K, M, 1).to(torch.float32).sum(dim=0).to(dtype=x.dtype) norm = torch.sqrt(norm) r = 1.0 / (norm / math.sqrt(K) + eps) return proj, norm, r @@ -3291,17 +3283,7 @@ def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): ) ) - return ( - grad_x, - grad_weight, - grad_ap, - grad_apo, - grad_ar, - grad_bias, - None, - None, - None, - ) + return (grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None, None) class FusedHAggregate(torch.autograd.Function): diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json index b27e30dd333..1742b322e7f 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 10.86149, - "2": 10.85475, - "3": 10.86695, - "4": 10.84628, + "1": 10.8615, + "2": 10.85472, + "3": 10.8669, + "4": 10.84631, "5": 10.88468, - "6": 10.89685, - "7": 10.87277, - "8": 10.86588, - "9": 10.86984, + "6": 10.89684, + "7": 10.8728, + "8": 10.86586, + "9": 10.86986, "10": 10.83763, - "11": 10.89461, + "11": 10.89456, "12": 10.87955, - "13": 10.87687, - "14": 10.90354, - "15": 10.83117, - "16": 10.8345, - "17": 10.80067, - "18": 10.8207, + "13": 10.87685, + "14": 10.90355, + "15": 10.83119, + "16": 10.83448, + "17": 10.80062, + "18": 10.82067, "19": 10.81458, - "20": 10.71809, - "21": 10.68629, - "22": 10.53196, - "23": 10.70487, - "24": 10.58558, - "25": 10.51904, - "26": 10.58488, - "27": 10.60102, - "28": 10.5354, - "29": 10.57113, - "30": 10.33245, - "31": 10.05837, - "32": 10.42792, - "33": 10.42027, - "34": 10.16988, - "35": 10.23073, - "36": 10.18757, - "37": 10.31244, - "38": 10.14208, - "39": 10.38138, - "40": 10.0485, - "41": 10.10336, - "42": 10.17158, - "43": 9.78293, - "44": 9.90959, - "45": 9.78518, - "46": 9.76888, - "47": 10.10084, - "48": 9.80968, - "49": 9.48782, - "50": 9.86709 + "20": 10.71813, + "21": 10.68631, + "22": 10.53201, + "23": 10.70485, + "24": 10.58556, + "25": 10.51905, + "26": 10.58493, + "27": 10.60097, + "28": 10.53539, + "29": 10.57115, + "30": 10.33246, + "31": 10.05838, + "32": 10.42787, + "33": 10.42026, + "34": 10.16987, + "35": 10.23069, + "36": 10.18758, + "37": 10.31242, + "38": 10.14211, + "39": 10.38137, + "40": 10.04848, + "41": 10.10339, + "42": 10.1716, + "43": 9.78295, + "44": 9.90964, + "45": 9.78514, + "46": 9.76887, + "47": 10.10085, + "48": 9.80969, + "49": 9.48775, + "50": 9.8671 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 34399.0, - "2": 5068.0, - "3": 5139.0, - "4": 2133.0, - "5": 34705.0, - "6": 5089.0, - "7": 5475.0, - "8": 5376.0, - "9": 5624.0, - "10": 34196.0, - "11": 1824.0, - "12": 2245.0, - "13": 5078.0, - "14": 1770.0, - "15": 2066.0, - "16": 34671.0, - "17": 4463.0, - "18": 1721.0, - "19": 2101.0, - "20": 1708.0, - "21": 5367.0, - "22": 34462.0, - "23": 2073.0, - "24": 1882.0, - "25": 1990.0, - "26": 4731.0, - "27": 2311.0, - "28": 1957.0, - "29": 5633.0, - "30": 34746.0, - "31": 1810.0, - "32": 34706.0, - "33": 5781.0, - "34": 5295.0, - "35": 4581.0, - "36": 2350.0, - "37": 35156.0, - "38": 2200.0, - "39": 5912.0, - "40": 6668.0, - "41": 5213.0, - "42": 2898.0, - "43": 2068.0, - "44": 35053.0, - "45": 5016.0, - "46": 6505.0, - "47": 2949.0, - "48": 6075.0, - "49": 3332.0, - "50": 35248.0 + "1": 34445.0, + "2": 5092.0, + "3": 5082.0, + "4": 2234.0, + "5": 34694.0, + "6": 5124.0, + "7": 5424.0, + "8": 5377.0, + "9": 5661.0, + "10": 34264.0, + "11": 1794.0, + "12": 2253.0, + "13": 5143.0, + "14": 1701.0, + "15": 1985.0, + "16": 34635.0, + "17": 4485.0, + "18": 1835.0, + "19": 2091.0, + "20": 1737.0, + "21": 5432.0, + "22": 34529.0, + "23": 2165.0, + "24": 5614.0, + "25": 1986.0, + "26": 4677.0, + "27": 2303.0, + "28": 1987.0, + "29": 5692.0, + "30": 34700.0, + "31": 1806.0, + "32": 34651.0, + "33": 5871.0, + "34": 5319.0, + "35": 4587.0, + "36": 2273.0, + "37": 35183.0, + "38": 2106.0, + "39": 5924.0, + "40": 6636.0, + "41": 5250.0, + "42": 2906.0, + "43": 2093.0, + "44": 34915.0, + "45": 4917.0, + "46": 6584.0, + "47": 2997.0, + "48": 6084.0, + "49": 3289.0, + "50": 35194.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 556914176.0, - "2": 556914176.0, - "3": 556914176.0, - "4": 556914176.0, - "5": 556914176.0, - "6": 556914176.0, - "7": 556914176.0, - "8": 556914176.0, - "9": 556914176.0, - "10": 557962752.0, - "11": 556914176.0, - "12": 556914176.0, - "13": 556914176.0, - "14": 556914176.0, - "15": 556914176.0, - "16": 556914176.0, - "17": 556914176.0, - "18": 556914176.0, - "19": 556914176.0, - "20": 556914176.0, - "21": 556914176.0, - "22": 556914176.0, - "23": 556914176.0, - "24": 556914176.0, - "25": 556914176.0, - "26": 557962752.0, - "27": 556914176.0, - "28": 556914176.0, - "29": 557962752.0, - "30": 556914176.0, - "31": 556914176.0, - "32": 556914176.0, - "33": 556914176.0, - "34": 557962752.0, - "35": 556914176.0, - "36": 556914176.0, - "37": 556914176.0, - "38": 556914176.0, - "39": 556914176.0, - "40": 557962752.0, - "41": 556914176.0, - "42": 556914176.0, - "43": 556914176.0, - "44": 556914176.0, - "45": 556914176.0, - "46": 556914176.0, - "47": 556914176.0, - "48": 556914176.0, - "49": 556914176.0, - "50": 556914176.0 + "1": 554030592.0, + "2": 554030592.0, + "3": 554030592.0, + "4": 554030592.0, + "5": 554030592.0, + "6": 554030592.0, + "7": 554030592.0, + "8": 554030592.0, + "9": 554030592.0, + "10": 554030592.0, + "11": 555079168.0, + "12": 554030592.0, + "13": 554030592.0, + "14": 554030592.0, + "15": 554030592.0, + "16": 554030592.0, + "17": 554030592.0, + "18": 554030592.0, + "19": 554030592.0, + "20": 554030592.0, + "21": 554030592.0, + "22": 554030592.0, + "23": 554030592.0, + "24": 554030592.0, + "25": 554030592.0, + "26": 554817024.0, + "27": 554030592.0, + "28": 554030592.0, + "29": 554817024.0, + "30": 554030592.0, + "31": 554030592.0, + "32": 555079168.0, + "33": 554030592.0, + "34": 554030592.0, + "35": 555079168.0, + "36": 554030592.0, + "37": 554030592.0, + "38": 554030592.0, + "39": 554030592.0, + "40": 554030592.0, + "41": 554030592.0, + "42": 554030592.0, + "43": 554030592.0, + "44": 554030592.0, + "45": 555079168.0, + "46": 554030592.0, + "47": 554030592.0, + "48": 554030592.0, + "49": 554030592.0, + "50": 554030592.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 1746399232.0, - "2": 1937075200.0, - "3": 1937075200.0, - "4": 1937075200.0, - "5": 1937075200.0, - "6": 1937075200.0, - "7": 1937075200.0, - "8": 1937075200.0, - "9": 1937075200.0, - "10": 1937075200.0, - "11": 1937075200.0, - "12": 1937075200.0, - "13": 1937075200.0, - "14": 1937075200.0, - "15": 1937075200.0, - "16": 1937075200.0, - "17": 1937075200.0, - "18": 1937075200.0, - "19": 1937075200.0, - "20": 1937075200.0, - "21": 1937075200.0, - "22": 1937075200.0, - "23": 1937075200.0, - "24": 1937075200.0, - "25": 1937075200.0, - "26": 1937075200.0, - "27": 1937075200.0, - "28": 1937075200.0, - "29": 1937075200.0, - "30": 1937075200.0, - "31": 1937075200.0, - "32": 1937075200.0, - "33": 1937075200.0, - "34": 1937075200.0, - "35": 1937075200.0, - "36": 1937075200.0, - "37": 1937075200.0, - "38": 1937075200.0, - "39": 1937075200.0, - "40": 1937075200.0, - "41": 1937075200.0, - "42": 1937075200.0, - "43": 1937075200.0, - "44": 1937075200.0, - "45": 1937075200.0, - "46": 1937075200.0, - "47": 1937075200.0, - "48": 1937075200.0, - "49": 1937075200.0, - "50": 1937075200.0 + "1": 1746497024.0, + "2": 1934944256.0, + "3": 1935075840.0, + "4": 1935075840.0, + "5": 1935075840.0, + "6": 1935075840.0, + "7": 1935075840.0, + "8": 1935075840.0, + "9": 1935075840.0, + "10": 1935075840.0, + "11": 1935075840.0, + "12": 1935075840.0, + "13": 1935075840.0, + "14": 1935075840.0, + "15": 1935075840.0, + "16": 1935075840.0, + "17": 1935075840.0, + "18": 1935075840.0, + "19": 1935075840.0, + "20": 1935075840.0, + "21": 1935075840.0, + "22": 1935075840.0, + "23": 1935075840.0, + "24": 1935075840.0, + "25": 1935075840.0, + "26": 1936123392.0, + "27": 1936124416.0, + "28": 1936124416.0, + "29": 1936124416.0, + "30": 1936124416.0, + "31": 1936124416.0, + "32": 1936124416.0, + "33": 1936124416.0, + "34": 1936124416.0, + "35": 1936124416.0, + "36": 1936124416.0, + "37": 1936124416.0, + "38": 1936124416.0, + "39": 1936124416.0, + "40": 1936124416.0, + "41": 1936124416.0, + "42": 1936124416.0, + "43": 1936124416.0, + "44": 1936124416.0, + "45": 1936124416.0, + "46": 1936124416.0, + "47": 1936124416.0, + "48": 1936124416.0, + "49": 1936124416.0, + "50": 1936124416.0 } }, "iteration-time": { @@ -233,55 +233,55 @@ "step_interval": 1, "values": { "1": "nan", - "2": 23.68682, - "3": 0.49278, - "4": 0.48539, - "5": 0.48882, - "6": 0.48719, - "7": 0.48515, - "8": 0.48616, - "9": 0.48409, - "10": 0.48496, - "11": 0.48688, - "12": 0.48809, - "13": 0.48532, - "14": 0.48548, - "15": 0.48618, - "16": 0.48449, - "17": 0.48545, - "18": 0.4894, - "19": 0.48704, - "20": 0.48687, - "21": 0.48792, - "22": 0.50768, - "23": 0.48897, - "24": 0.48564, - "25": 0.48418, - "26": 0.49874, - "27": 0.4869, - "28": 0.48559, - "29": 0.48507, - "30": 0.48655, - "31": 0.49426, - "32": 0.49577, - "33": 0.48641, - "34": 0.4864, - "35": 0.48522, - "36": 0.48214, - "37": 0.4844, - "38": 0.48496, - "39": 0.48722, - "40": 0.48728, - "41": 0.48777, - "42": 0.48522, - "43": 0.48574, - "44": 0.48593, - "45": 0.48608, - "46": 0.48732, - "47": 0.48832, - "48": 0.48904, - "49": 0.48735, - "50": 0.48544 + "2": 23.10433, + "3": 0.49351, + "4": 0.48583, + "5": 0.48743, + "6": 0.48248, + "7": 0.48132, + "8": 0.47946, + "9": 0.47876, + "10": 0.48555, + "11": 0.48606, + "12": 0.4821, + "13": 0.48438, + "14": 0.47915, + "15": 0.48041, + "16": 0.48223, + "17": 0.48385, + "18": 0.4823, + "19": 0.47983, + "20": 0.47979, + "21": 0.48223, + "22": 0.48145, + "23": 0.4834, + "24": 0.47861, + "25": 0.4789, + "26": 0.49661, + "27": 0.48275, + "28": 0.48647, + "29": 0.48495, + "30": 0.48443, + "31": 0.48439, + "32": 0.48536, + "33": 0.4851, + "34": 0.48436, + "35": 0.48538, + "36": 0.48395, + "37": 0.48157, + "38": 0.4833, + "39": 0.48478, + "40": 0.4854, + "41": 0.48518, + "42": 0.48467, + "43": 0.48414, + "44": 0.48508, + "45": 0.48453, + "46": 0.48611, + "47": 0.48532, + "48": 0.48346, + "49": 0.48455, + "50": 0.48453 } } } \ No newline at end of file From 86e18218339c048b300086ad09fc7e04f57a41e8 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 4 Jun 2026 09:39:55 +0800 Subject: [PATCH 32/34] Refine mHC Sinkhorn and h_post eps semantics --- megatron/core/fusions/fused_mhc_kernels.py | 56 +++++++------------ megatron/core/transformer/hyper_connection.py | 7 +-- .../fusions/test_fused_mhc_kernels.py | 5 +- 3 files changed, 22 insertions(+), 46 deletions(-) diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py index 8c0f25ebedc..92c436690a6 100644 --- a/megatron/core/fusions/fused_mhc_kernels.py +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -245,16 +245,15 @@ def _triton_sinkhorn_fwd_kernel( M = tl.exp2((logits - row_max[:, None]) * TLOG2E) tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) - if NUM_ITERS > 0: + row_sum = tl.sum(M, axis=1) + M = M / row_sum[:, None] + eps + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + for _ in range(NUM_ITERS - 1): row_sum = tl.sum(M, axis=1) - M = M / row_sum[:, None] + eps + M = M / (row_sum[:, None] + eps) col_sum = tl.sum(M, axis=0) M = M / (col_sum[None, :] + eps) - for _ in range(NUM_ITERS - 1): - row_sum = tl.sum(M, axis=1) - M = M / (row_sum[:, None] + eps) - col_sum = tl.sum(M, axis=0) - M = M / (col_sum[None, :] + eps) tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) @@ -377,7 +376,7 @@ class TritonFusedSinkhorn(torch.autograd.Function): def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): """Run Triton Sinkhorn forward and save initial matrix for backward.""" out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) - ctx.save_for_backward(input_logits if num_iterations == 0 else M_init) + ctx.save_for_backward(M_init) ctx.num_iterations = num_iterations ctx.eps = eps return out @@ -385,15 +384,7 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output: Tensor): """Run Triton Sinkhorn backward.""" - (saved,) = ctx.saved_tensors - if ctx.num_iterations == 0: - with torch.enable_grad(): - logits = saved.detach().requires_grad_(True) - row_max = logits.max(dim=-1, keepdim=True).values - M = torch.exp(logits - row_max) - M.backward(grad_output) - return logits.grad, None, None - M_init = saved + (M_init,) = ctx.saved_tensors grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None @@ -869,16 +860,15 @@ def _ct_sinkhorn_fwd_kernel( index=(pid, 0, 0), tile=ct.reshape(M.astype(M_init_out.dtype), (TILE_SIZE, HC, HC)), ) - if NUM_ITERS > 0: + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / row_sum + eps + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + for _ in range(NUM_ITERS - 1): row_sum = ct.sum(M, axis=2, keepdims=True) - M = M / row_sum + eps + M = M / (row_sum + eps) col_sum = ct.sum(M, axis=1, keepdims=True) M = M / (col_sum + eps) - for _ in range(NUM_ITERS - 1): - row_sum = ct.sum(M, axis=2, keepdims=True) - M = M / (row_sum + eps) - col_sum = ct.sum(M, axis=1, keepdims=True) - M = M / (col_sum + eps) ct.store(out, index=(pid, 0, 0), tile=ct.reshape(M.astype(out.dtype), (TILE_SIZE, HC, HC))) @ct.kernel @@ -1574,7 +1564,7 @@ def _ct_reduce_compute_h_kernel( h_pre_linear = pre_accum * alpha_pre * inv_r_eps + bias_pre h_post_linear = post_accum * alpha_post * inv_r_eps + bias_post h_pre = _ct_sigmoid(h_pre_linear) + compute_h_eps - h_post = (_ct_sigmoid(h_post_linear) + compute_h_eps) * 2.0 + h_post = _ct_sigmoid(h_post_linear) * 2.0 ct.store(H_PRE, index=(bid_m, 0), tile=h_pre.astype(H_PRE.dtype)) ct.store(H_POST, index=(bid_m, 0), tile=h_post.astype(H_POST.dtype)) @@ -2354,7 +2344,7 @@ def _ct_fused_grad_h_proj_kernel( PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO ) proj_post = ct.astype(proj_post, ct.float32) - sigmoid_post = h_post * 0.5 - compute_h_eps + sigmoid_post = h_post * 0.5 grad_h_post = gy_post * sigmoid_post * (1.0 - sigmoid_post) * 2.0 grad_proj_post = grad_h_post * alpha_post * inv_r_eps grad_r_from_h += ct.sum( @@ -3137,7 +3127,7 @@ def _torch_proj_rms_compute_h( ) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) h_pre = h[..., :n].sigmoid() + compute_h_eps - h_post = (h[..., n : 2 * n].sigmoid() + compute_h_eps) * 2 + h_post = h[..., n : 2 * n].sigmoid() * 2 h_res = h[..., 2 * n :] return h_pre, h_post, h_res, r @@ -3151,7 +3141,7 @@ class CutileSinkhornKnopp(torch.autograd.Function): def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): """Run cuTile Sinkhorn forward and save initial matrix for backward.""" output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) - ctx.save_for_backward(input_logits if num_iterations == 0 else M_init) + ctx.save_for_backward(M_init) ctx.num_iterations = num_iterations ctx.eps = eps return output @@ -3159,15 +3149,7 @@ def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): @staticmethod def backward(ctx, grad_output): """Run cuTile Sinkhorn backward.""" - (saved,) = ctx.saved_tensors - if ctx.num_iterations == 0: - with torch.enable_grad(): - logits = saved.detach().requires_grad_(True) - row_max = logits.max(dim=-1, keepdim=True).values - M = torch.exp(logits - row_max) - M.backward(grad_output) - return logits.grad, None, None - M_init = saved + (M_init,) = ctx.saved_tensors grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) return grad_input, None, None diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index e1a2bd8331a..94b759ca23c 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -23,9 +23,6 @@ def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: row_max = input_logits.max(dim=-1, keepdim=True).values M = torch.exp(input_logits - row_max) - if num_iterations == 0: - return M - M = M / M.sum(dim=-1, keepdim=True) + eps M = M / (M.sum(dim=-2, keepdim=True) + eps) for _ in range(num_iterations - 1): @@ -289,8 +286,8 @@ def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre) h_pre = h[..., : self.n].sigmoid() + self.compute_h_eps # [s, b, n] - # H_post = 2(σ(α_post * (θ_post @ x̃) + b_post) + eps) - h_post = (h[..., self.n : 2 * self.n].sigmoid() + self.compute_h_eps) * 2 + # H_post = 2σ(α_post * (θ_post @ x̃) + b_post) + h_post = h[..., self.n : 2 * self.n].sigmoid() * 2 h_res = h[..., 2 * self.n :] return h_pre, h_post, h_res diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 2c5662e8dc8..3e17c18664d 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -79,9 +79,6 @@ def _info(): def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: row_max = logits.max(dim=-1, keepdim=True).values M = torch.exp(logits - row_max) - if num_iters == 0: - return M - M = M / M.sum(dim=-1, keepdim=True) + eps M = M / (M.sum(dim=-2, keepdim=True) + eps) for _ in range(num_iters - 1): @@ -145,7 +142,7 @@ def _ref_proj_rms_compute_h( alpha = torch.cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n)]) h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) h_pre = h[..., :n].sigmoid() + compute_h_eps - h_post = (h[..., n : 2 * n].sigmoid() + compute_h_eps) * 2 + h_post = h[..., n : 2 * n].sigmoid() * 2 h_res = h[..., 2 * n :] return h_pre, h_post, h_res, r From 961bbca701e548ccd908ea9280a97633918b2a91 Mon Sep 17 00:00:00 2001 From: jingqiny-99 Date: Thu, 4 Jun 2026 09:58:30 +0800 Subject: [PATCH 33/34] Use softmax for native Sinkhorn first row step --- megatron/core/transformer/hyper_connection.py | 4 +--- tests/unit_tests/fusions/test_fused_mhc_kernels.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 94b759ca23c..3bb7813fcd7 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -21,9 +21,7 @@ @torch.compile def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: - row_max = input_logits.max(dim=-1, keepdim=True).values - M = torch.exp(input_logits - row_max) - M = M / M.sum(dim=-1, keepdim=True) + eps + M = input_logits.softmax(dim=-1) + eps M = M / (M.sum(dim=-2, keepdim=True) + eps) for _ in range(num_iterations - 1): M = M / (M.sum(dim=-1, keepdim=True) + eps) diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py index 3e17c18664d..fffe2dc36bb 100644 --- a/tests/unit_tests/fusions/test_fused_mhc_kernels.py +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -77,9 +77,7 @@ def _info(): def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: - row_max = logits.max(dim=-1, keepdim=True).values - M = torch.exp(logits - row_max) - M = M / M.sum(dim=-1, keepdim=True) + eps + M = logits.softmax(dim=-1) + eps M = M / (M.sum(dim=-2, keepdim=True) + eps) for _ in range(num_iters - 1): M = M / (M.sum(dim=-1, keepdim=True) + eps) From e9687d72a5b1797df9af1ad182b48030644dbfe0 Mon Sep 17 00:00:00 2001 From: Jingqin Yang Date: Wed, 10 Jun 2026 23:09:34 -0700 Subject: [PATCH 34/34] fix: golden values again --- .../golden_values_dev_dgx_h100.json | 288 +++++++++--------- 1 file changed, 144 insertions(+), 144 deletions(-) diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json index 1742b322e7f..6ea673bd273 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json @@ -5,55 +5,55 @@ "step_interval": 1, "values": { "1": 10.8615, - "2": 10.85472, - "3": 10.8669, - "4": 10.84631, - "5": 10.88468, - "6": 10.89684, - "7": 10.8728, - "8": 10.86586, - "9": 10.86986, + "2": 10.85474, + "3": 10.86694, + "4": 10.84627, + "5": 10.88469, + "6": 10.89688, + "7": 10.87279, + "8": 10.86582, + "9": 10.86985, "10": 10.83763, - "11": 10.89456, + "11": 10.8946, "12": 10.87955, - "13": 10.87685, - "14": 10.90355, - "15": 10.83119, - "16": 10.83448, - "17": 10.80062, - "18": 10.82067, - "19": 10.81458, - "20": 10.71813, - "21": 10.68631, - "22": 10.53201, - "23": 10.70485, - "24": 10.58556, - "25": 10.51905, - "26": 10.58493, - "27": 10.60097, - "28": 10.53539, - "29": 10.57115, - "30": 10.33246, - "31": 10.05838, - "32": 10.42787, - "33": 10.42026, - "34": 10.16987, - "35": 10.23069, - "36": 10.18758, - "37": 10.31242, - "38": 10.14211, - "39": 10.38137, + "13": 10.87688, + "14": 10.90357, + "15": 10.83118, + "16": 10.83452, + "17": 10.80061, + "18": 10.8207, + "19": 10.81454, + "20": 10.71814, + "21": 10.68632, + "22": 10.53196, + "23": 10.70484, + "24": 10.58558, + "25": 10.51904, + "26": 10.58492, + "27": 10.60104, + "28": 10.5354, + "29": 10.57116, + "30": 10.33245, + "31": 10.05835, + "32": 10.42789, + "33": 10.42025, + "34": 10.16981, + "35": 10.23073, + "36": 10.18759, + "37": 10.31244, + "38": 10.1421, + "39": 10.38138, "40": 10.04848, "41": 10.10339, - "42": 10.1716, - "43": 9.78295, - "44": 9.90964, - "45": 9.78514, - "46": 9.76887, - "47": 10.10085, - "48": 9.80969, - "49": 9.48775, - "50": 9.8671 + "42": 10.17156, + "43": 9.78293, + "44": 9.90959, + "45": 9.78517, + "46": 9.76885, + "47": 10.10086, + "48": 9.80966, + "49": 9.48779, + "50": 9.86706 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 34445.0, - "2": 5092.0, - "3": 5082.0, - "4": 2234.0, - "5": 34694.0, - "6": 5124.0, - "7": 5424.0, - "8": 5377.0, - "9": 5661.0, - "10": 34264.0, - "11": 1794.0, - "12": 2253.0, - "13": 5143.0, - "14": 1701.0, - "15": 1985.0, - "16": 34635.0, - "17": 4485.0, - "18": 1835.0, - "19": 2091.0, - "20": 1737.0, - "21": 5432.0, - "22": 34529.0, - "23": 2165.0, - "24": 5614.0, - "25": 1986.0, - "26": 4677.0, - "27": 2303.0, - "28": 1987.0, - "29": 5692.0, - "30": 34700.0, - "31": 1806.0, - "32": 34651.0, - "33": 5871.0, - "34": 5319.0, - "35": 4587.0, - "36": 2273.0, - "37": 35183.0, - "38": 2106.0, - "39": 5924.0, - "40": 6636.0, - "41": 5250.0, - "42": 2906.0, - "43": 2093.0, - "44": 34915.0, - "45": 4917.0, - "46": 6584.0, - "47": 2997.0, - "48": 6084.0, - "49": 3289.0, - "50": 35194.0 + "1": 2145.0, + "2": 34499.0, + "3": 4946.0, + "4": 2054.0, + "5": 2182.0, + "6": 5738.0, + "7": 2027.0, + "8": 4745.0, + "9": 2308.0, + "10": 5889.0, + "11": 5460.0, + "12": 5084.0, + "13": 2207.0, + "14": 5705.0, + "15": 2440.0, + "16": 5360.0, + "17": 2395.0, + "18": 4330.0, + "19": 34574.0, + "20": 5921.0, + "21": 4368.0, + "22": 34392.0, + "23": 34883.0, + "24": 1610.0, + "25": 4786.0, + "26": 4652.0, + "27": 5558.0, + "28": 2821.0, + "29": 2256.0, + "30": 5895.0, + "31": 4542.0, + "32": 34648.0, + "33": 5116.0, + "34": 2009.0, + "35": 2309.0, + "36": 4918.0, + "37": 5902.0, + "38": 2241.0, + "39": 2447.0, + "40": 34981.0, + "41": 5104.0, + "42": 5731.0, + "43": 5320.0, + "44": 2591.0, + "45": 2242.0, + "46": 5263.0, + "47": 5652.0, + "48": 35221.0, + "49": 2518.0, + "50": 2612.0 } }, "mem-allocated-bytes": { @@ -233,55 +233,55 @@ "step_interval": 1, "values": { "1": "nan", - "2": 23.10433, - "3": 0.49351, - "4": 0.48583, - "5": 0.48743, - "6": 0.48248, - "7": 0.48132, - "8": 0.47946, - "9": 0.47876, - "10": 0.48555, - "11": 0.48606, - "12": 0.4821, - "13": 0.48438, - "14": 0.47915, - "15": 0.48041, - "16": 0.48223, - "17": 0.48385, - "18": 0.4823, - "19": 0.47983, - "20": 0.47979, - "21": 0.48223, - "22": 0.48145, - "23": 0.4834, - "24": 0.47861, - "25": 0.4789, - "26": 0.49661, - "27": 0.48275, - "28": 0.48647, - "29": 0.48495, - "30": 0.48443, - "31": 0.48439, - "32": 0.48536, - "33": 0.4851, - "34": 0.48436, - "35": 0.48538, - "36": 0.48395, - "37": 0.48157, - "38": 0.4833, - "39": 0.48478, - "40": 0.4854, - "41": 0.48518, - "42": 0.48467, - "43": 0.48414, - "44": 0.48508, - "45": 0.48453, - "46": 0.48611, - "47": 0.48532, - "48": 0.48346, - "49": 0.48455, - "50": 0.48453 + "2": 23.09864, + "3": 0.46809, + "4": 0.45676, + "5": 0.45281, + "6": 0.45141, + "7": 0.45108, + "8": 0.45268, + "9": 0.45264, + "10": 0.45588, + "11": 0.45475, + "12": 0.45436, + "13": 0.45256, + "14": 0.45413, + "15": 0.45289, + "16": 0.45389, + "17": 0.45246, + "18": 0.456, + "19": 0.45408, + "20": 0.4536, + "21": 0.45464, + "22": 0.45355, + "23": 0.45318, + "24": 0.45254, + "25": 0.45466, + "26": 0.46781, + "27": 0.45087, + "28": 0.45103, + "29": 0.45285, + "30": 0.45365, + "31": 0.45193, + "32": 0.45232, + "33": 0.45259, + "34": 0.45235, + "35": 0.45154, + "36": 0.45743, + "37": 0.46097, + "38": 0.46155, + "39": 0.45845, + "40": 0.45893, + "41": 0.45848, + "42": 0.45956, + "43": 0.45884, + "44": 0.45764, + "45": 0.45891, + "46": 0.45815, + "47": 0.45661, + "48": 0.45943, + "49": 0.45719, + "50": 0.45865 } } } \ No newline at end of file