Skip to content

Conversation

YangKai0616
Copy link

Currently running pytest -rA gpt_oss/triton/attention.py will report the following error:

PASSED attention.py::test_eq[0-None-0.125-64-8-8-128-1-1]
PASSED attention.py::test_eq[0-None-0.125-64-8-8-128-1-2]
PASSED attention.py::test_eq[0-None-0.125-64-8-8-128-128-1]
PASSED attention.py::test_eq[0-None-0.125-64-8-8-128-128-2]
PASSED attention.py::test_eq[0-None-0.125-64-8-8-32-1-1]
PASSED attention.py::test_eq[0-None-0.125-64-8-8-32-1-2]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-128-1-1]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-128-1-2]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-128-128-1]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-128-128-2]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-32-1-1]
PASSED attention.py::test_eq[0-128-0.125-64-8-8-32-1-2]
SKIPPED [8] attention.py:217: too many queries
FAILED attention.py::test_eq[5-None-0.125-64-8-8-128-1-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-None-0.125-64-8-8-128-1-2] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-None-0.125-64-8-8-128-128-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-None-0.125-64-8-8-128-128-2] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-None-0.125-64-8-8-32-1-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-None-0.125-64-8-8-32-1-2] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-128-1-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-128-1-2] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-128-128-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-128-128-2] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-32-1-1] - AssertionError: Tensor-likes are not close!
FAILED attention.py::test_eq[5-128-0.125-64-8-8-32-1-2] - AssertionError: Tensor-likes are not close!

In the _attn_fwd kernel, there is a logical error in the calculation of the loop boundary lo: when start_q > 0 (e.g. start_q = 5), lo is incorrectly restricted to start_q (i.e., 5). This causes key positions 0 to start_q-1 (i.e., 0-4) to be skipped. However, these positions are valid under the causal mask (since query position start_q = 5 can attend to historical keys 0-4). In the reference implementation, these positions are correctly included. Therefore, regardless of whether BANDWIDTH is enabled, the loop starting point lo should begin from 0 (the absolute start of the key sequence) rather than start_q.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants