Update hstu and fix ragged attn (#59)

Summary: Update HSTU ragged attention kernel with code change. Pull Request resolved: #59 Test Plan: OSS CI Reviewed By: manman-ren Differential Revision: D66257228 Pulled By: xuzhao9 fbshipit-source-id: b68272edbb55979bb30bf601f987330de7b76707
pytorch-labs · Nov 21, 2024 · 9f7e919 · 9f7e919
1 parent 17b38a4
commit 9f7e919
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 7 deletions.
diff --git 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
diff --git 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
             )
+4 −2		.gitignore
+4 −4		generative_recommenders/data/eval.py
+0 −42		generative_recommenders/modeling/ndp_module.py
+5 −151		generative_recommenders/modeling/sequential/autoregressive_losses.py
+12 −13		generative_recommenders/modeling/sequential/encoder_utils.py
+5 −5		generative_recommenders/modeling/sequential/hstu.py
+194 −0		generative_recommenders/modeling/sequential/losses/sampled_softmax.py
+4 −4		generative_recommenders/modeling/sequential/sasrec.py
+0 −50		generative_recommenders/modeling/similarity/mol.py
+26 −51		generative_recommenders/modeling/similarity_module.py
+3 −1		generative_recommenders/modeling/similarity_utils.py
+55 −0		generative_recommenders/ops/pytorch/layer_norm.py
+8 −1		generative_recommenders/ops/triton/triton_hstu_linear.py
+58 −70		generative_recommenders/ops/triton/triton_ragged_hstu_attention.py
+24 −23		generative_recommenders/rails/similarities/dot_product_similarity_fn.py
+53 −0		generative_recommenders/rails/similarities/module.py
+0 −1		generative_recommenders/trainer/data_loader.py
+10 −2		generative_recommenders/trainer/train.py
+1 −0		main.py