[WIP] mHC: Manifold-constrained Hyper Connection by anhminhnguyenhoang · Pull Request #1859 · ROCm/aiter

anhminhnguyenhoang · 2026-01-16T15:15:28Z

Co-authors: @waqahmed-amd-fi @anhminhnguyenhoang

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…ection mapping

… kernel (#1877) * Refactor mHC kernel and wrapper to implement equations 14-18 with fused kernel * improve comments * Enhance documentation for mhc function: clarify equations, input/output shapes, and activation details * Enhance documentation in test_mhc.py: clarify equations, input/output shapes, and activation details for mHC kernel tests * Add _sinkhorn_knopp_log_domain_kernel to the fusion module * Add logging and sync Sinkhorn-Knopp function for doubly stochastic matrices * sync log-domain Sinkhorn-Knopp kernel for doubly stochastic matrix projection * Improve logging in mhc function to include all alpha parameters

github-actions · 2026-01-20T13:04:12Z

aiter/ops/triton/fusions/mhc.py

+    else:
+        assert out.shape == (M, N), f"Output shape mismatch: expected ({M}, {N}), got {out.shape}"
+        assert out.dtype == x.dtype, f"Output dtype mismatch: expected {x.dtype}, got {out.dtype}"
+        assert out.device == x.device, f"Output device mismatch"


⚠️ [ruff] <F541> _{reported by reviewdog 🐶}
f-string without any placeholders

Suggested change

assert out.device == x.device, f"Output device mismatch"

assert out.device == x.device, "Output device mismatch"

…plified comments for sinkhorn-knopp impl

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/fusions/test_mhc.py

+
+    # Res-stream: no constraints (identity activation)
+    # Just verify it exists
+    assert out_res.shape == (M, n_squared), f"Res-stream shape mismatch"


⚠️ [ruff] <F541> _{reported by reviewdog 🐶}
f-string without any placeholders

Suggested change

assert out_res.shape == (M, n_squared), f"Res-stream shape mismatch"

assert out_res.shape == (M, n_squared), "Res-stream shape mismatch"

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/__init__.py

 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

+from .mhc_ref import *


⚠️ [ruff] <F403> _{reported by reviewdog 🐶}
from .mhc_ref import * used; unable to detect undefined names

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/__init__.py

 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

+from .mhc_ref import *
 from .mla_decode_ref import *


⚠️ [ruff] <F403> _{reported by reviewdog 🐶}
from .mla_decode_ref import * used; unable to detect undefined names

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/__init__.py


+from .mhc_ref import *
 from .mla_decode_ref import *
 from .mla_extend_ref import *


⚠️ [ruff] <F403> _{reported by reviewdog 🐶}
from .mla_extend_ref import * used; unable to detect undefined names

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/__init__.py

+from .mhc_ref import *
 from .mla_decode_ref import *
 from .mla_extend_ref import *
 from .rotary_embedding import *


⚠️ [ruff] <F403> _{reported by reviewdog 🐶}
from .rotary_embedding import * used; unable to detect undefined names

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/mhc_ref.py

+        - H^res: [2n:2n+n²] residual connection (identity) (n² elements)
+    """
+    x_f32 = x.to(torch.float32)
+    nC = x.shape[1]


⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable nC is assigned to but never used

Suggested change

nC = x.shape[1]

x.shape[1]

github-actions · 2026-01-20T15:09:09Z

op_tests/triton_tests/utils/mhc_ref.py

+    H_tilde = x_norm @ phi_f32
+
+    # Split into three streams
+    n_squared = n * n


⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable n_squared is assigned to but never used

Suggested change

n_squared = n * n

n * n

* Refactor mHC kernel and wrapper to implement equations 14-18 with fused kernel * improve comments * Enhance documentation for mhc function: clarify equations, input/output shapes, and activation details * Enhance documentation in test_mhc.py: clarify equations, input/output shapes, and activation details for mHC kernel tests * Add _sinkhorn_knopp_log_domain_kernel to the fusion module * Add logging and sync Sinkhorn-Knopp function for doubly stochastic matrices * sync log-domain Sinkhorn-Knopp kernel for doubly stochastic matrix projection * Improve logging in mhc function to include all alpha parameters * Fix H dimensions * Refactor mHC function to return separate output tensors for pre, post, and residual streams * Refactor mhc_torch to return separate output tensors for pre, post, and residual streams * Adjust tolerance for is_doubly_stochastic assertion in test_sk_matrix_sizes for bfloat16 precision

…ke H_res doubly stochastic

anhminhnguyenhoang

Looks good, I would personally clean up the comments as they look a bit redundant

anhminhnguyenhoang · 2026-01-21T09:45:49Z

op_tests/triton_tests/fusions/test_mhc.py

        H_res_torch.to(torch.float32),
-        atol=1e-2,
-        rtol=1e-2,
+        atol=5e-2,


Did you run into test failure because of this for similar tests that you need to relax the tolerance?

Yes, mainly because of sinkhorn which is an iterative process and returns higher differences due you only 10 iterations. May be we can try 20 for better results?

…o pre, post, and residual streams; update tests accordingly.

…ccuracy of assertions.

- Update benchmark script to use dynamic configurations

…dd log for reference

…multiple sizes

…erformance

…d parameters tuning for better performance on MI355X (gfx950)

This reverts commit 33ded07.

Changes: - Fix BLOCK_N calculation: Use max(n_factorial, n_squared) instead of conditionally selecting based on hres_op flag. This ensures BLOCK_N is always large enough for both intermediate computation (n! values for softmax over permutations) and final output (n² doubly stochastic matrix). - Rename n_factorial to n_res_expected for clarity, since this variable holds factorial(n) in lite mode but n_squared in sinkhorn mode. The new name better reflects its purpose as the expected dimension of the residual stream output. - Remove redundant conditional in acc_res_cols assignment since the mode-specific logic is already handled in n_res_expected definition.

add initial implementation of projection mapping

c35c8ac

anhminhnguyenhoang assigned anhminhnguyenhoang and waqahmed-amd-fi Jan 16, 2026

anhminhnguyenhoang requested a review from a team January 16, 2026 15:15

anhminhnguyenhoang marked this pull request as draft January 16, 2026 15:15

anhminhnguyenhoang changed the title ~~mHC: Manifold-constrained Hyper Connection~~ [WIP] mHC: Manifold-constrained Hyper Connection Jan 16, 2026

anhminhnguyenhoang requested review from Chi-Chu319, hellozhuo-amd and juuso-oskari January 16, 2026 15:20

waqahmed-amd-fi and others added 5 commits January 19, 2026 04:57

Refactor mHC kernel and wrapper to include sigmoid activation in proj…

2a8325c

…ection mapping

Add Sinkhorn-Knopp log-domain kernel implementation

f11244d

clean up sinkhorn-knopp tests

87e5839

review invalid test case

152bb21

github-actions bot reviewed Jan 20, 2026

View reviewed changes

waqahmed-amd-fi and others added 3 commits January 20, 2026 08:35

Fix H dims

36e36d6

fix test_mhc_output_range

6156c13

Refactor test cases in mHC and Sinkhorn-Knopp implementations and sim…

80b8d34

…plified comments for sinkhorn-knopp impl

github-actions bot reviewed Jan 20, 2026

View reviewed changes

waqahmed-amd-fi and others added 3 commits January 21, 2026 10:29

optimization to loads x_tile once, reducing memory bandwidth

d952829

Update mHC implementation to apply Sinkhorn-Knopp (Equation 19) to ma…

30741ca

…ke H_res doubly stochastic

anhminhnguyenhoang commented Jan 21, 2026

View reviewed changes

waqahmed-amd-fi and others added 7 commits January 21, 2026 06:08

Refactor mHC implementation to separate projection (phi) matrices int…

686711f

…o pre, post, and residual streams; update tests accordingly.

Enhance mHC fused kernel to implement stream-aware processing

831572b

Refactor mHC implementation

ad198b3

Adjust tolerance levels in mHC tests based on input size to improve a…

05810eb

…ccuracy of assertions.

Add benchmark scripts for mHC kernel performance evaluation

7d333a8

add modes to bench

46d023a

- Add naive configs for fused mHC and Sinkhorn-Knopp kernels

e8f4464

- Update benchmark script to use dynamic configurations

waqahmed-amd-fi and others added 30 commits January 30, 2026 03:21

update comments and replace if-else with assert check

57c75ac

revised _mhc_lite_fused_split_kernel kernel

68c76b0

revised _mhc_lite_fused_reduce_kernel

604426d

add mhc-lite bench mode

e52ebdc

integrate mhc-lite into mhc_fused

e70f3dd

update config loading for mode

41b4908

remove duplicate implementation of non-iterative sinkhorn-knopp and a…

73e4541

…dd log for reference

add configuration files for gfx950 MHC FUSED and SINKHORN variants

827cda9

add tuned configs for mhc-lite for mi300x

5ddbbff

add tuned configs for mhc-lite for mi355x

27cb8e8

update MHC FUSED Lite configurations for improved performance across …

ac98e74

…multiple sizes

allow configurable BLOCK_N size for fused_mhc in lite mode

7ff6954

optimize softmax computation in _compute_hres_mhc_lite for improved p…

5e399d6

…erformance

add BLOCK_N configuration for MHC FUSED Lite across multiple sizes an…

33f0b65

…d parameters tuning for better performance on MI355X (gfx950)

update NUM_KSPLIT in gfx950 sinkhorn mode to bring back perf

d0bfb9c

Add HRES_OP to configuration and update benchmark naming

0f50f0f

refactor for optimizations

0371721

fix broken tests

a403562

Add hres_op tests

c1811d6

Improve BLOCK_N loading

068766c

Add tuned block configs for gfx942

4ad48ed

disable hres=0 algorithm for improved perf

33ded07

Revert "disable hres=0 algorithm for improved perf"

2c28b18

This reverts commit 33ded07.

add bug fix for modifying cached config

38d417d

Updated BLOCK_N values to block configs and config loading behaviour

8370fe4

Tuned configuration files for improved performance

83b840d

add further tuned configs

20d58fc

fix tests

3cb1301

add further tuned configs 2

0420e60

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] mHC: Manifold-constrained Hyper Connection#1859

[WIP] mHC: Manifold-constrained Hyper Connection#1859
anhminhnguyenhoang wants to merge 68 commits intomainfrom
feat/mhc-deepseek

anhminhnguyenhoang commented Jan 16, 2026 •

edited

Loading

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

github-actions bot Jan 20, 2026

Uh oh!

anhminhnguyenhoang left a comment

Uh oh!

anhminhnguyenhoang Jan 21, 2026

Uh oh!

waqahmed-amd-fi Jan 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	assert out.device == x.device, f"Output device mismatch"
	assert out.device == x.device, "Output device mismatch"

	assert out_res.shape == (M, n_squared), f"Res-stream shape mismatch"
	assert out_res.shape == (M, n_squared), "Res-stream shape mismatch"

Conversation

anhminhnguyenhoang commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

anhminhnguyenhoang left a comment

Choose a reason for hiding this comment

Uh oh!

anhminhnguyenhoang Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

waqahmed-amd-fi Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anhminhnguyenhoang commented Jan 16, 2026 •

edited

Loading