From 9f7e9194c71e031b160440d8eb6fe6449c1310c4 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Wed, 20 Nov 2024 16:14:03 -0800 Subject: [PATCH] Update hstu and fix ragged attn (#59) Summary: Update HSTU ragged attention kernel with code change. Pull Request resolved: https://github.com/pytorch-labs/tritonbench/pull/59 Test Plan: OSS CI Reviewed By: manman-ren Differential Revision: D66257228 Pulled By: xuzhao9 fbshipit-source-id: b68272edbb55979bb30bf601f987330de7b76707 --- .github/workflows/pr.yaml | 5 ----- submodules/generative-recommenders | 2 +- tritonbench/operators/ragged_attention/hstu.py | 2 +- 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index b587d0d7..eb796c12 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -1,11 +1,6 @@ name: TritonBench PR Test on: pull_request: - paths: - - .ci/* - - test/test_gpu/* - - tritonbench/* - - .github/workflows/pr.yaml push: branches: - main diff --git a/submodules/generative-recommenders b/submodules/generative-recommenders index 53321455..77d8b947 160000 --- a/submodules/generative-recommenders +++ b/submodules/generative-recommenders @@ -1 +1 @@ -Subproject commit 53321455e93bef51b2622cd97ae6b28ca3c6fa68 +Subproject commit 77d8b9474e9a1acc7a892f0caf9fcec16d287aa0 diff --git a/tritonbench/operators/ragged_attention/hstu.py b/tritonbench/operators/ragged_attention/hstu.py index 896c5ee6..bfedd369 100644 --- a/tritonbench/operators/ragged_attention/hstu.py +++ b/tritonbench/operators/ragged_attention/hstu.py @@ -170,7 +170,7 @@ def forward( kwargs["num_targets"], kwargs["ATTN_BIAS_TYPE"], # relative_bias_type kwargs["MAX_ATTN_LEN"], # max_attn_len - kwargs["contextual_seq_len"], # contextual_seq_len + kwargs["CONTEXTUAL_SEQ_LEN"], # contextual_seq_len kwargs["sort_by_length_indices"], # sort_by_length )