From 39dc4dfc7e5a7d82a8b5f3b8abe7102cf0f2fe14 Mon Sep 17 00:00:00 2001 From: Gene Der Su Date: Thu, 15 Jan 2026 10:08:02 -0800 Subject: [PATCH 1/3] [Aiter] upgrade aiter to be compatible with latest TE Signed-off-by: Gene Der Su --- .../pytorch/kernels/attention/attention_csrc_impl.py | 9 ++++++++- setup.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/primus_turbo/pytorch/kernels/attention/attention_csrc_impl.py b/primus_turbo/pytorch/kernels/attention/attention_csrc_impl.py index 09dda0a12..42a4f2baa 100644 --- a/primus_turbo/pytorch/kernels/attention/attention_csrc_impl.py +++ b/primus_turbo/pytorch/kernels/attention/attention_csrc_impl.py @@ -32,8 +32,11 @@ def attention_aiter_csrc_forward_impl( alibi_slopes: Optional[torch.Tensor], return_lse: bool, return_softmax: bool, + sink_size: int = 0, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward( q, k, @@ -43,8 +46,12 @@ def attention_aiter_csrc_forward_impl( causal, window_size_left, window_size_right, + sink_size, bias, alibi_slopes, + q_descale, + k_descale, + v_descale, return_lse, return_softmax, ) diff --git a/setup.py b/setup.py index 3aaa08fed..5084eee68 100644 --- a/setup.py +++ b/setup.py @@ -386,7 +386,7 @@ def build_jax_extension(): if torch_ext is not None and not is_package_installed("aiter"): print("[Primus-Turbo Setup] aiter not found, will be installed automatically.") install_requires.append( - "aiter @ git+https://github.com/ROCm/aiter.git@a25cb58ee4ee97cb5ad4b426b648a8e66103d706" + "aiter @ git+https://github.com/ROCm/aiter.git@ed4bdeac27b930f8efb1976af14db88ba3d648ad" ) else: print("[Primus-Turbo Setup] Skipping aiter installation.") From 72559954bdc94f63b0700e7be1e8159275a071b7 Mon Sep 17 00:00:00 2001 From: Gene Der Su Date: Thu, 15 Jan 2026 11:14:51 -0800 Subject: [PATCH 2/3] fix egg name Signed-off-by: Gene Der Su --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 5084eee68..2913b4346 100644 --- a/setup.py +++ b/setup.py @@ -386,7 +386,7 @@ def build_jax_extension(): if torch_ext is not None and not is_package_installed("aiter"): print("[Primus-Turbo Setup] aiter not found, will be installed automatically.") install_requires.append( - "aiter @ git+https://github.com/ROCm/aiter.git@ed4bdeac27b930f8efb1976af14db88ba3d648ad" + "amd-aiter @ git+https://github.com/ROCm/aiter.git@ed4bdeac27b930f8efb1976af14db88ba3d648ad" ) else: print("[Primus-Turbo Setup] Skipping aiter installation.") From 944e242018507d68e0c4da5052ac10feef485a1c Mon Sep 17 00:00:00 2001 From: Gene Der Su Date: Fri, 16 Jan 2026 15:33:40 -0800 Subject: [PATCH 3/3] update amd-aiter package check Signed-off-by: Gene Der Su --- setup.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index 2913b4346..135f0b0a8 100644 --- a/setup.py +++ b/setup.py @@ -382,15 +382,14 @@ def build_jax_extension(): "hip-python", ] - # Conditionally add aiter if torch_ext is being built and aiter is not already installed - if torch_ext is not None and not is_package_installed("aiter"): - print("[Primus-Turbo Setup] aiter not found, will be installed automatically.") + # Conditionally add amd-aiter if torch_ext is being built and amd-aiter is not already installed + if torch_ext is not None and not is_package_installed("amd-aiter"): + print("[Primus-Turbo Setup] amd-aiter not found, will be installed automatically.") install_requires.append( "amd-aiter @ git+https://github.com/ROCm/aiter.git@ed4bdeac27b930f8efb1976af14db88ba3d648ad" ) else: - print("[Primus-Turbo Setup] Skipping aiter installation.") - + print("[Primus-Turbo Setup] Skipping amd-aiter installation.") if BUILD_JAX: entry_points["jax_plugins"] = ["primus_turbo = primus_turbo.jax"]