Test spec of Puzzle 9 seems to mismatch the problem setting?

[nil0x9]

In the given template code, there is no B1 in signature of kernel jit function flashatt_kernel, which seems to mismatch the problem description? Should i add another parameter to the function?

@triton.jit
def flashatt_kernel(q_ptr, k_ptr, v_ptr, z_ptr, N0, T, B0: tl.constexpr):
    return

test(flashatt_kernel, flashatt_spec, B={"B0":200},
     nelem={"N0": 200, "T": 200})

the problem setting states as follows:

A scalar version of FlashAttention.

Uses zero programs. Block size B0 represents k of length N0.
Block size B0 represents q of length N0. Block size B0 represents v of length N0.
Sequence length is T. Process it B1 < T elements at a time.

$$z_{i} = \sum_{j} \text{softmax}(q_1 k_1, \ldots, q_T k_T)j v{j} \text{ for } i = 1\ldots N_0$$

This can be done in 1 loop using a similar trick from the last puzzle.

[nil0x9]

Btw, B={"B0":200} seems to be problematic too, for we usually use tl.arange(0, B0) in the kernel to calculate offsets and tl.arange only accepts ranges of powers of 2.

[SiriusNEO]

Hi @nil0x9 , I also find there is something wrong with Puzzle 9. I fixed the description here: https://github.com/SiriusNEO/Triton-Puzzles-Lite/blob/main/puzzles_ans.py#L587-L600. And I also think B0: 200 is problematic. You can use B0: 64 to test it. Hoping this can help you!