fix: attention refactor and base doc #216

gaetansnl · 2022-12-14T11:10:31Z

Linear is partially refactored because input name is fixed by early_config_prune
I avoided important changes except refactor to linear:

Removed

rematerialize rm and rn to save registers
rm = m_block_idx * BLOCK_M + tl.arange(0, BLOCK_M)
rn = n_block_idx * BLOCK_N + tl.arange(0, BLOCK_N)

used variable initialized with tl.max_contiguous(tl.multiple_of everywhere

fixes #229

gaetansnl · 2023-01-04T09:41:51Z

full test is passing, i suggest re-run on 3090 just in case

jonathlela

It's a bit hard to review but the new naming convention make the code easier to read.
There is still some omissions in the whole renaming. I've pointed out 3-4 of them in the feedback, not sure if there is others remaining.

jonathlela · 2023-01-10T13:51:34Z

src/kernl/implementations/README.md

+Suffixes list:
+- `_idx` integer representing an index
+- `_size` integer representing a size
+- `_off` integer reoresenting an offset


representing

jonathlela · 2023-01-10T13:57:52Z

src/kernl/implementations/README.md

+## Dimension naming
+- Dimension is singular
+- If dimension follows variable name from a formula. You can use this name. Example MNK for matmul
+- Use `col` or `row` singular if you don't have a new for the last two dimensions


don't have a name ?

jonathlela · 2023-01-10T14:25:35Z

src/kernl/implementations/layer_norm.py

+    output_row_stride,
+    output_col_stride,
+    a_row_stride,
+    a_col_stride,
    N,


should we use n_size and block_n_size ?

jonathlela · 2023-01-10T14:29:12Z

src/kernl/implementations/linear_layer.py

-    pid_n = (pid % width) // (group_size)
+    group_idx = program_idx // width
+    group_size = min(grid_m - group_idx * GROUP_M, GROUP_M)
+    m_block_idx = group_idx * GROUP_M + (program_idx % group_size)


should it be block_m_idx and block_n_idx ?

jonathlela · 2023-01-10T14:32:40Z

src/kernl/implementations/attention_skinny.py

-    BLOCK_N: tl.constexpr,
+    BLOCK_M_SIZE: tl.constexpr,
+    BLOCK_DHEAD_SIZE: tl.constexpr,
+    BLOCK_N_SIZE: tl.constexpr,
 ):
    # Index of the block on M axis (M axis is the rows of matrix K)
    n_block_idx = tl.program_id(0)


Should it be block_n_idx and block_m_idx ?

jonathlela · 2023-01-10T14:34:05Z

src/kernl/implementations/attention_skinny.py

-    ptrs_q = Q + offs_q
-    ptrs_k = K + offs_k
-    ptrs_v = V + offs_v
+    ptrs_q = q_ptr + offs_q


q_ptrs, k_ptrs and v_ptrs ?

pommedeterresautee

Tried to review but it's really hard.
Did it passed tests? (full pytest)
Added few remarks on top of those of @jonathlela

pommedeterresautee · 2023-01-11T15:02:24Z

src/kernl/implementations/README.md

@@ -0,0 +1,41 @@
+# Naming conventions


README in a more visible location like contribution folder? With a good title / name so we know its purpose.

pommedeterresautee · 2023-01-11T15:04:46Z

src/kernl/implementations/attention.py

@@ -186,21 +186,21 @@ def _fwd_kernel(
                   ┌────────────┐
               │   │            │
    M Dimension│   ├────────────┤     ┌───┐
-      size_m   │   │            │     │   │ BLOCK_M
+      m_size   │   │            │     │   │ BLOCK_M


BLOCK_M_SIZE ? same fo N

gaetansnl · 2023-01-11T15:21:13Z

@pommedeterresautee no got an error at the end, I will try to reproduce

white-gorilla

The Markown is optimized but the lack of line breaks makes it less readable in .md for maintainers and not readable once converted to HTML on the site.

Don't hesitate to put line breaks after the titles and especially before the lists (otherwise they will be displayed inline).
Tip: stay consistent on punctuation, especially the . at the end of the line and the : (either you put them everywhere or you don't).

Otherwise it's fine.

pommedeterresautee · 2023-01-17T15:43:36Z

❯ pytest test/test_torchdynamo.py 
...
============================================================================================== 270 passed in 3488.03s (0:58:08) ==============================================================================================

pommedeterresautee

quite hard to review line by line but from what I checked, it's ok.
plus test pass

white-gorilla · 2023-01-17T16:05:00Z

@gaetansnl Still missing a line break.

fix: attention refactor and base doc

741c3e1

github-actions bot added the fix hurrah, bug fixed! label Dec 14, 2022

gaetansnl added 6 commits December 15, 2022 09:32

Merge branch 'main' into fix/refactor-kernels

ac0f064

Merge branch 'main' into fix/refactor-kernels

0ea61f1

fix: refactor attention_skinny

7f2d793

fix: refactor vec_mat

25cecde

fix: refactor bacthed_matmul

54e405c

fix: refactor layernorm

1ece296

github-actions bot added fix hurrah, bug fixed! and removed fix hurrah, bug fixed! labels Jan 3, 2023

gaetansnl added 2 commits January 3, 2023 16:24

fix: refactor linear_layer.py

bf1af34

fix: reformat

5e9f74b

gaetansnl requested a review from pommedeterresautee January 3, 2023 15:55

gaetansnl marked this pull request as ready for review January 3, 2023 15:55

fix: line length

b53bb3f

github-actions bot added fix hurrah, bug fixed! and removed fix hurrah, bug fixed! labels Jan 4, 2023

fix: plural

3234dc6

github-actions bot added fix hurrah, bug fixed! and removed fix hurrah, bug fixed! labels Jan 4, 2023

fix: refactor

3cf033b

jonathlela approved these changes Jan 10, 2023

View reviewed changes

gaetansnl added 2 commits January 11, 2023 08:45

Merge branch 'main' into fix/refactor-kernels

22c6312

fix: fix

a35a129

pommedeterresautee reviewed Jan 11, 2023

View reviewed changes

Merge branch 'main' into fix/refactor-kernels

0ce8f1c

gaetansnl added 3 commits January 17, 2023 09:05

Merge branch 'main' into fix/refactor-kernels

1f40dee

fix: naming in attention comment

98a3d97

fix: move docs

438b0c2

gaetansnl requested review from white-gorilla and pommedeterresautee January 17, 2023 08:30

white-gorilla requested changes Jan 17, 2023

View reviewed changes

fix: fix doc formatting

6d60895

gaetansnl requested a review from white-gorilla January 17, 2023 15:06

pommedeterresautee approved these changes Jan 17, 2023

View reviewed changes

white-gorilla approved these changes Jan 17, 2023

View reviewed changes

gaetansnl and others added 2 commits January 17, 2023 17:36

fix: fix doc formatting

8965fc9

Merge branch 'main' into fix/refactor-kernels

791be20

gaetansnl merged commit 46e8a58 into main Jan 17, 2023

gaetansnl deleted the fix/refactor-kernels branch January 17, 2023 16:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: attention refactor and base doc #216

fix: attention refactor and base doc #216

gaetansnl commented Dec 14, 2022 •

edited

Loading

gaetansnl commented Jan 4, 2023 •

edited

Loading

jonathlela left a comment

jonathlela Jan 10, 2023

gaetansnl Jan 11, 2023

jonathlela Jan 10, 2023

gaetansnl Jan 11, 2023

jonathlela Jan 10, 2023

gaetansnl Jan 11, 2023

jonathlela Jan 10, 2023

gaetansnl Jan 11, 2023

jonathlela Jan 10, 2023

gaetansnl Jan 11, 2023

jonathlela Jan 10, 2023

pommedeterresautee Jan 11, 2023

gaetansnl Jan 11, 2023

pommedeterresautee left a comment

pommedeterresautee Jan 11, 2023

pommedeterresautee Jan 11, 2023

gaetansnl commented Jan 11, 2023

white-gorilla left a comment

pommedeterresautee commented Jan 17, 2023

pommedeterresautee left a comment

white-gorilla commented Jan 17, 2023

fix: attention refactor and base doc #216

fix: attention refactor and base doc #216

Conversation

gaetansnl commented Dec 14, 2022 • edited Loading

gaetansnl commented Jan 4, 2023 • edited Loading

jonathlela left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pommedeterresautee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gaetansnl commented Jan 11, 2023

white-gorilla left a comment

Choose a reason for hiding this comment

pommedeterresautee commented Jan 17, 2023

pommedeterresautee left a comment

Choose a reason for hiding this comment

white-gorilla commented Jan 17, 2023

gaetansnl commented Dec 14, 2022 •

edited

Loading

gaetansnl commented Jan 4, 2023 •

edited

Loading