From 18e6b461fcee89c3b162eecfda099fae810a48ec Mon Sep 17 00:00:00 2001 From: Manman Ren Date: Tue, 26 Nov 2024 09:33:11 -0800 Subject: [PATCH] fix Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags: --- tritonbench/operators/flash_attention/operator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tritonbench/operators/flash_attention/operator.py b/tritonbench/operators/flash_attention/operator.py index 518516c6..16f90814 100644 --- a/tritonbench/operators/flash_attention/operator.py +++ b/tritonbench/operators/flash_attention/operator.py @@ -475,7 +475,7 @@ def get_input_iter(self) -> Generator: H = self.H def get_ctx_vals(): - for i in range(self.SEQ_LEN, 11): + for i in range(self.SEQ_LEN, self.SEQ_LEN+1): N_CTX = 2**i # BATCH = 16384 // N_CTX # H = 2048 // D_HEAD