Đây là một chủ đề khá nâng cao, nằm ở giao điểm của AI Compiler, GPU Programming, PyTorch, Triton và Machine Learning Systems. Tuy nhiên, nếu học đúng lộ trình thì Helion lại khá dễ hiểu vì nó được thiết kế để giấu đi sự phức tạp của Triton.

Tôi sẽ dạy theo kiểu từ nền tảng đến chuyên sâu.

Bài 1. Helion ra đời để giải quyết vấn đề gì?

Muốn hiểu Helion trước tiên phải hiểu GPU Kernel.

Giả sử ta có phép cộng hai vector.

C = A + B

Trong PyTorch chỉ cần:

C = A + B

PyTorch sẽ gọi kernel CUDA đã được viết sẵn.

Nếu bạn muốn tự viết kernel GPU thì có 3 lựa chọn.

Mức	Công nghệ	Độ khó
Cao nhất	CUDA C++	Rất khó
Trung bình	Triton	Khó
Cao hơn PyTorch	Helion	Dễ hơn Triton

Có thể hình dung:

PyTorch
    ↓
Helion
    ↓
Triton
    ↓
CUDA
    ↓
GPU

Helion thực chất không chạy trực tiếp trên GPU.

Nó sinh ra Triton.

Triton lại sinh ra CUDA/PTX.

Vì sao Triton vẫn còn khó?

Ví dụ viết Matrix Multiplication bằng Triton.

Bạn phải tự tính:

block size
số warp
pid
stride
pointer
masking
indexing

Ví dụ:

pid = tl.program_id(0)

pid_m = ...
pid_n = ...

offs_m = ...
offs_n = ...

ptr = base + ...

Khoảng 60% code chỉ để tính địa chỉ bộ nhớ.

Không phải tính toán.

Helion muốn bạn không phải làm việc đó.

Ý tưởng của Helion

Thay vì bảo GPU:

“Hãy đọc địa chỉ số 10032”

Helion cho bạn nói:

“Lấy tile này.”

Ví dụ

for tile_m, tile_n in hl.tile([m, n]):

đọc giống Python.

Nhưng Helion sẽ tự biến thành

program_id

↓

tile index

↓

pointer

↓

memory address

↓

GPU

Helion là gì?

Tài liệu viết:

Python Embedded DSL

DSL nghĩa là:

Domain Specific Language.

Tức là

Một ngôn ngữ nhỏ nằm trong Python.

Giống như:

SQL nằm trong Java

TensorFlow Graph nằm trong Python

JAX DSL

Helion DSL

Bạn vẫn viết Python.

Chỉ một số lệnh

hl.tile()

hl.zeros()

hl.range()

mang ý nghĩa đặc biệt.

Một kernel Helion gồm hai phần

Ví dụ:

@helion.kernel()
def matmul(...):

    ...

Trong kernel có hai vùng.

Phần 1

m, k = x.size()

out = torch.empty(...)

Đây là Python bình thường.

Chạy trên CPU.

Phần 2

for tile in hl.tile(...):

Đây không còn là Python nữa.

Helion sẽ dịch đoạn này thành Triton.

Nói cách khác

CPU code

↓

Helion boundary

↓

GPU code

hl.tile() là gì?

Đây là khái niệm quan trọng nhất.

Giả sử ma trận

1024 × 1024

GPU không xử lý toàn bộ.

Nó chia thành các ô.

Ví dụ

□□□□□□□□

□□□□□□□□

□□□□□□□□

□□□□□□□□

Mỗi ô gọi là

Tile

Ví dụ

64×64

thì

1024

↓

64

↓

16 tile

Theo chiều ngang.

Chiều dọc cũng vậy.

Tổng cộng

16 ×16 =256 tile

GPU sẽ chạy song song 256 tile.

Trong Helion

for tile_m, tile_n in hl.tile([m, n]):

không phải vòng lặp Python.

Mà nghĩa là

“Hãy chia ma trận thành các tile.”

Ví dụ Matrix Multiplication

Helion:

for tile_m, tile_n in hl.tile([m, n]):

    acc = hl.zeros(...)

    for tile_k in hl.tile(k):

        acc = torch.addmm(...)

Nếu viết dưới dạng sơ đồ.

        M

┌────┬────┐
│    │    │
├────┼────┤
│    │    │
└────┴────┘
      ↑

GPU xử lý từng tile

Bên trong mỗi tile

for tile_k

lại chia tiếp chiều K.

Giống GEMM chuẩn.

Điều đặc biệt của Helion

Trong Triton bạn phải tự chọn

Tile = 16

hay

Tile =32

hay

Tile =64

Helion không.

Bạn chỉ viết

hl.tile()

Nó tự thử.

Autotuning là gì?

Đây là trái tim của Helion.

Giả sử có các lựa chọn:

Tile

Warp

Indexing

pointer

block_ptr

descriptor

PID

flat

blocked

persistent

Có hàng nghìn khả năng.

Ví dụ

64
×

8 warp

×

block_ptr

×

persistent

Hay

32

×

4 warp

×

pointer

×

flat

Helion sẽ chạy thử tất cả.

Đo thời gian.

Lấy cái nhanh nhất.

Đây chính là

Autotuning

Tại sao mất tới 10 phút?

Vì Helion có thể thử

cấu hình.

Ví dụ log

Generation 1

↓

Generation 2

↓

Generation 3

...

↓

Generation 20

Đây giống thuật toán tiến hóa.

Nó giữ cấu hình tốt.

Loại cấu hình chậm.

Sinh cấu hình mới.

Differential Evolution

Trong log có

DifferentialEvolutionSearch

Đây là thuật toán tối ưu.

Không thử ngẫu nhiên.

Mà giống tiến hóa.

Config A

Config B

↓

Lai

↓

Config C

↓

Đột biến

↓

Config D

Sau nhiều thế hệ.

Nó hội tụ.

Config là gì?

Đây là cấu hình GPU.

Ví dụ

Config(

block_sizes

num_warps

num_stages

indexing
)

Đây không phải code.

Mà là cách GPU chạy.

Ví dụ

block_size

↓

64×64

Hay

warp

↓

8

Tại sao nên lưu Config?

Sau khi autotune.

Helion in ra

Config(...)

Bạn copy.

@helion.kernel(
config=...
)

Lần sau

Không cần thử nữa.

Khởi động ngay.

Settings khác Config như thế nào?

Đây là điểm nhiều người mới học dễ nhầm lẫn.

Config	Settings
Quyết định cách GPU thực thi	Quyết định cách Helion biên dịch
Ảnh hưởng hiệu năng	Ảnh hưởng quá trình phát triển
Được autotune	Không autotune
block size, num_warps, indexing…	print_output_code, autotune_effort, logging…

Ví dụ:

@helion.kernel(
    autotune_effort="none",
    print_output_code=True
)

Đây là Settings: dùng để phát triển và gỡ lỗi.

Còn:

@helion.kernel(
    config=helion.Config(
        block_sizes=[64,64,64],
        num_warps=8
    )
)

Đây là Config: quyết định hiệu năng của kernel trên GPU.

Quá trình biên dịch Helion

Có thể hình dung toàn bộ quá trình như sau:

Mã Helion
      │
      ▼
Phân tích Python AST
      │
      ▼
Sinh Triton
      │
      ▼
TorchInductor
      │
      ▼
PTX/CUDA
      │
      ▼
GPU NVIDIA

Helion đóng vai trò là một lớp trừu tượng cao hơn Triton. Bạn mô tả bài toán ở mức “xử lý theo tile”, còn Helion tự sinh mã Triton và tự tối ưu.

Vì sao Helion đáng chú ý?

Nếu nhìn theo mức độ trừu tượng:

Công nghệ	Bạn phải quan tâm
CUDA	Thread, Block, Shared Memory, Warp, Pointer, Synchronization
Triton	Tile, Pointer, Stride, PID, Mask
Helion	Chỉ cần Tile và các phép toán PyTorch
PyTorch	Chỉ cần gọi API có sẵn

Xu hướng của Helion tương tự như lịch sử phát triển của lập trình: thay vì yêu cầu lập trình viên điều khiển phần cứng ở mức thấp, nó chuyển gánh nặng tối ưu sang trình biên dịch và bộ autotuner.

Lộ trình học Helion tôi đề xuất

Để hiểu Helion một cách bài bản, nên học theo thứ tự sau:

Kiến trúc GPU: Thread, Warp, Block, Grid, Shared Memory.
CUDA cơ bản: hiểu cách GPU thực thi kernel.
PyTorch Tensor: broadcasting, views, reductions, torch.addmm, torch.sum,…
Triton: PID, tile, pointer, masking, block pointers.
Helion: DSL, hl.tile(), autotuning, Config, Settings.
TorchInductor và PyTorch Compiler: để hiểu Helion tích hợp vào hệ sinh thái PyTorch như thế nào.

Với nền tảng này, bạn sẽ không chỉ biết cách viết kernel bằng Helion, mà còn hiểu vì sao Helion có thể sinh ra các kernel Triton hiệu năng cao và cách đánh giá, tinh chỉnh chúng khi cần.

%pip install "torch==2.9.*" --index-url https://download.pytorch.org/whl/cu126
%pip install helion

"""
Helion Softmax Kernel Examples
==============================
This example demonstrates multiple Helion kernel implementations of the softmax function,
including a simple wrapper around PyTorch's softmax, and a numerically optimized two-pass version.
The example also includes a check function to compare these kernels against PyTorch's
built-in softmax for correctness.
"""

# %%
from __future__ import annotations
import torch
import helion
from helion._testing import run_example
import helion.language as hl


# %%
@helion.kernel(autotune_effort="quick")
def softmax(x: torch.Tensor) -> torch.Tensor:
    """
    Simple Helion kernel wrapping PyTorch's softmax function.
    Args:
        x (torch.Tensor): Input tensor of shape [n, m].
    Returns:
        torch.Tensor: Softmax output tensor of the same shape.
    """
    n, _m = x.size()
    out = torch.empty_like(x)
    for tile_n in hl.tile(n):
        out[tile_n, :] = torch.nn.functional.softmax(x[tile_n, :], dim=1)
    return out


# %%
def check(m: int, n: int) -> None:
    """
    Runs correctness checks comparing Helion softmax kernels against PyTorch's softmax.
    Args:
        m (int): Number of rows in input tensor.
        n (int): Number of columns in input tensor.
    """
    x = torch.randn([m, n], device="cuda", dtype=torch.float16)
    run_example(softmax, lambda x: torch.nn.functional.softmax(x, dim=1), (x,))


# %%
def main() -> None:
    """
    Main function to run the softmax kernel correctness check with example input size.
    """
    check(4096, 2560)


# %%
if __name__ == "__main__":
    main()

"""
Helion Softmax Kernel Examples
==============================
This example demonstrates multiple Helion kernel implementations of the softmax function,
including a simple wrapper around PyTorch's softmax, and a numerically optimized two-pass version.
The example also includes a check function to compare these kernels against PyTorch's
built-in softmax for correctness.
"""

# %%
from __future__ import annotations
import torch
import helion
from helion._testing import run_example
import helion.language as hl


# %%
@helion.kernel(autotune_effort="quick")
def softmax_two_pass(x: torch.Tensor) -> torch.Tensor:
    """
    Numerically optimized Helion kernel performing softmax in two passes.
    Args:
        x (torch.Tensor): Input tensor of shape [m, n].
    Returns:
        torch.Tensor: Softmax output tensor of the same shape.
    """
    m, n = x.size()
    out = torch.empty_like(x)
    block_size_m = hl.register_block_size(m)
    block_size_n = hl.register_block_size(n)
    for tile_m in hl.tile(m, block_size=block_size_m):
        mi = hl.full([tile_m], float("-inf"), dtype=torch.float32)
        di = hl.zeros([tile_m], dtype=torch.float32)
        for tile_n in hl.tile(n, block_size=block_size_n):
            values = x[tile_m, tile_n]
            local_amax = torch.amax(values, dim=1)
            mi_next = torch.maximum(mi, local_amax)
            di = di * torch.exp(mi - mi_next) + torch.exp(
                values - mi_next[:, None]
            ).sum(dim=1)
            mi = mi_next
        for tile_n in hl.tile(n, block_size=block_size_n):
            values = x[tile_m, tile_n]
            out[tile_m, tile_n] = torch.exp(values - mi[:, None]) / di[:, None]
    return out


# %%
def check(m: int, n: int) -> None:
    """
    Runs correctness checks comparing Helion softmax kernels against PyTorch's softmax.
    Args:
        m (int): Number of rows in input tensor.
        n (int): Number of columns in input tensor.
    """
    x = torch.randn([m, n], device="cuda", dtype=torch.float16)
    run_example(softmax_two_pass, lambda x: torch.nn.functional.softmax(x, dim=1), (x,))


# %%
def main() -> None:
    """
    Main function to run the softmax kernel correctness check with example input size.
    """
    check(4096, 2560)


# %%
if __name__ == "__main__":
    main()

Đoạn mã gồm hai cách viết Softmax bằng Helion:

Kernel	Cách làm
`softmax`	Gọi trực tiếp `torch.nn.functional.softmax` trên từng nhóm hàng
`softmax_two_pass`	Tự triển khai thuật toán Softmax ổn định số học theo hai lượt

Cả hai đều được Helion biên dịch thành một GPU kernel Triton.

1. Softmax là gì?

Với một hàng dữ liệu:

[
x = [x_1,x_2,\ldots,x_n]
]

Softmax được tính như sau:

[
\operatorname{softmax}(x_i)

\frac{e^{x_i}}{\sum_j e^{x_j}}
]

Ví dụ:

[
x=[1,2,3]
]

Kết quả gần đúng:

[
[0.09,\ 0.24,\ 0.67]
]

Tổng các phần tử luôn bằng 1.

Trong đoạn mã, Softmax được tính theo chiều:

dim=1

Nghĩa là tính Softmax trên từng hàng của tensor hai chiều.

2. Cài đặt thư viện

%pip install "torch==2.9.*" --index-url https://download.pytorch.org/whl/cu126
%pip install helion

Dòng đầu cài PyTorch 2.9 có hỗ trợ CUDA 12.6.

Dòng thứ hai cài Helion.

Cú pháp %pip thường được dùng trong:

Jupyter Notebook
Google Colab
IPython

Nếu chạy trong file Python thông thường thì cài bằng terminal:

pip install "torch==2.9.*" --index-url https://download.pytorch.org/whl/cu126
pip install helion

3. Các thư viện được import

from __future__ import annotations
import torch
import helion
from helion._testing import run_example
import helion.language as hl

Thành phần	Vai trò
`torch`	Tạo tensor và sử dụng các phép toán PyTorch
`helion`	Định nghĩa và biên dịch Helion kernel
`run_example`	Kiểm tra kết quả và có thể benchmark kernel
`helion.language as hl`	Cung cấp các toán tử DSL như `hl.tile`, `hl.zeros`
`annotations`	Trì hoãn việc đánh giá type hint

Lưu ý rằng helion._testing có dấu gạch dưới. Đây thường là API nội bộ phục vụ kiểm thử, không nhất thiết ổn định như API công khai.

4. Kernel Softmax thứ nhất

@helion.kernel(autotune_effort="quick")
def softmax(x: torch.Tensor) -> torch.Tensor:

Decorator:

@helion.kernel(...)

biến hàm Python thành một Helion kernel.

Tham số:

autotune_effort="quick"

yêu cầu Helion thực hiện autotuning ở mức nhanh.

Có thể hiểu tương đối như sau:

Mức autotune	Đặc điểm
`none`	Không tìm cấu hình tối ưu
`quick`	Thử một số cấu hình, thời gian ngắn hơn
Mức cao hơn	Tìm kiếm rộng hơn, mất nhiều thời gian hơn

Autotuner có thể lựa chọn:

số hàng trong mỗi tile;
kích thước block;
số warp;
cách đánh địa chỉ bộ nhớ;
chiến lược ánh xạ Program ID;
cách thực hiện reduction.

4.1 Lấy kích thước tensor

n, _m = x.size()

Nếu x có shape:

[4096, 2560]

thì:

n = 4096
_m = 2560

Tuy nhiên tên biến ở đây hơi dễ gây nhầm.

Theo cách đặt tên phổ biến, nên viết:

m, n = x.size()

Trong kernel hiện tại:

n thực chất là số hàng;
_m thực chất là số cột.

Dấu _ ở đầu _m cho biết biến này không được sử dụng tiếp.

Viết rõ hơn có thể là:

num_rows, _num_cols = x.size()

4.2 Tạo tensor đầu ra

out = torch.empty_like(x)

Tạo tensor out:

cùng kích thước với x;
cùng kiểu dữ liệu;
cùng thiết bị;
chưa được khởi tạo giá trị.

Nếu x có:

shape  = [4096, 2560]
dtype  = float16
device = cuda

thì out cũng có các thuộc tính tương tự.

Dùng empty_like nhanh hơn việc khởi tạo toàn bộ tensor bằng 0 vì kernel sẽ ghi đè tất cả các giá trị.

4.3 Chia hàng thành các tile

for tile_n in hl.tile(n):

Đây không phải vòng lặp Python thông thường.

Nó có nghĩa:

Chia chiều gồm n hàng thành các tile và ánh xạ các tile đó lên các chương trình GPU.

Giả sử có 4096 hàng và autotuner chọn tile 64 hàng:

[
4096 / 64 = 64\text{ tile}
]

GPU có thể xử lý các tile này song song.

tile_n không phải một số nguyên đơn lẻ. Nó đại diện cho một tập chỉ số hàng, chẳng hạn:

tile 0: hàng 0 đến 63
tile 1: hàng 64 đến 127
tile 2: hàng 128 đến 191
...

4.4 Tính Softmax trên các hàng thuộc tile

out[tile_n, :] = torch.nn.functional.softmax(
    x[tile_n, :],
    dim=1
)

Trong đó:

x[tile_n, :]

có nghĩa:

lấy các hàng thuộc tile_n;
lấy toàn bộ cột.

Ví dụ tile chứa 64 hàng, tensor tạm có thể có shape:

[64, 2560]

Sau đó:

softmax(..., dim=1)

tính Softmax trên 2560 phần tử của từng hàng.

Kết quả được ghi vào đúng vị trí:

out[tile_n, :]

Có thể hình dung:

Tensor x: 4096 × 2560

┌──────────────────────────┐
│ tile hàng số 1           │ → softmax từng hàng
├──────────────────────────┤
│ tile hàng số 2           │ → softmax từng hàng
├──────────────────────────┤
│ tile hàng số 3           │ → softmax từng hàng
├──────────────────────────┤
│ ...                      │
└──────────────────────────┘

4.5 Trả kết quả

return out

Phần bên ngoài vòng lặp như:

x.size()
torch.empty_like()
return

được dùng để chuẩn bị kernel.

Phần bên trong:

for tile_n in hl.tile(n):

được Helion chuyển thành mã GPU.

5. Hàm kiểm tra kernel thứ nhất

def check(m: int, n: int) -> None:

Hàm nhận:

m: số hàng;
n: số cột.

5.1 Tạo dữ liệu kiểm thử

x = torch.randn(
    [m, n],
    device="cuda",
    dtype=torch.float16
)

Ví dụ:

check(4096, 2560)

sẽ tạo tensor:

shape  = [4096, 2560]
device = GPU
dtype  = float16

torch.randn tạo số ngẫu nhiên theo phân phối chuẩn.

5.2 So sánh Helion với PyTorch

run_example(
    softmax,
    lambda x: torch.nn.functional.softmax(x, dim=1),
    (x,)
)

Ba tham số chính:

Tham số	Ý nghĩa
`softmax`	Kernel Helion cần kiểm tra
`lambda x: ...`	Hàm PyTorch dùng làm kết quả tham chiếu
`(x,)`	Tuple chứa đối số truyền vào

Dấu phẩy trong:

(x,)

rất quan trọng. Đây là tuple có một phần tử.

Không có dấu phẩy:

(x)

thì chỉ là biểu thức x, không phải tuple.

Về bản chất, hàm kiểm tra làm gần giống:

actual = softmax(x)
expected = torch.nn.functional.softmax(x, dim=1)

assert actual gần bằng expected

6. Hàm `main`

def main() -> None:
    check(4096, 2560)

Chạy kiểm thử với tensor có:

4096 hàng;
2560 cột.

Tổng số phần tử:

[
4096 \times 2560 = 10{,}485{,}760
]

Vì kiểu float16 dùng 2 byte mỗi phần tử, riêng một tensor chiếm khoảng:

[
10{,}485{,}760 \times 2

20{,}971{,}520\text{ byte}
]

Tức khoảng 20 MiB.

6.1 Kiểm tra chương trình có được chạy trực tiếp không

if __name__ == "__main__":
    main()

Nếu file được chạy trực tiếp:

python softmax_example.py

thì main() được gọi.

Nếu file được import từ module khác thì main() không tự chạy.

Trong Jupyter Notebook, đoạn này thường vẫn chạy vì __name__ thường bằng "__main__".

7. Kernel thứ hai: `softmax_two_pass`

Kernel thứ hai tự triển khai Softmax thay vì gọi trực tiếp:

torch.nn.functional.softmax

Mục tiêu là xử lý trường hợp số cột lớn mà không nhất thiết phải tải toàn bộ một hàng vào register cùng lúc.

@helion.kernel(autotune_effort="quick")
def softmax_two_pass(x: torch.Tensor) -> torch.Tensor:

8. Tại sao cần Softmax ổn định số học?

Công thức trực tiếp:

[
\frac{e^{x_i}}{\sum_j e^{x_j}}
]

có thể bị tràn số.

Ví dụ:

[
e^{1000}
]

là một số quá lớn, dễ trở thành vô cực:

inf

Softmax ổn định được tính bằng cách trừ giá trị lớn nhất:

[
\operatorname{softmax}(x_i)

\frac{e^{x_i-m}}{\sum_j e^{x_j-m}}
]

với:

[
m=\max_j x_j
]

Vì:

[
x_i-m\leq 0
]

nên:

[
e^{x_i-m}\leq 1
]

Giảm đáng kể nguy cơ overflow.

9. Chuẩn bị trong `softmax_two_pass`

m, n = x.size()
out = torch.empty_like(x)

Ở kernel này tên biến đúng với quy ước:

Biến	Ý nghĩa
`m`	Số hàng
`n`	Số cột

9.1 Đăng ký block size

block_size_m = hl.register_block_size(m)
block_size_n = hl.register_block_size(n)

Helion được yêu cầu tạo các tham số block size cho hai chiều:

chiều hàng m;
chiều cột n.

Các giá trị này có thể được autotuner lựa chọn.

Có thể hình dung:

block_size_m ∈ {1, 2, 4, 8, ...}
block_size_n ∈ {32, 64, 128, 256, ...}

Tập giá trị thực tế phụ thuộc Helion, phần cứng và không gian tìm kiếm.

Khác với:

hl.tile(m)

để Helion tự xử lý hoàn toàn, ở đây code lấy block size ra thành biến nhằm dùng lại cùng một kích thước trong nhiều vòng lặp.

Đặc biệt, block_size_n được dùng ở cả lượt thứ nhất và lượt thứ hai.

10. Vòng lặp theo tile hàng

for tile_m in hl.tile(m, block_size=block_size_m):

Kernel xử lý một nhóm hàng tại một thời điểm.

Giả sử:

m = 4096
block_size_m = 8

thì mỗi chương trình GPU xử lý 8 hàng.

Số chương trình cần thiết gần bằng:

[
4096/8=512
]

Các chương trình này có thể chạy song song.

11. Hai biến trạng thái `mi` và `di`

mi = hl.full(
    [tile_m],
    float("-inf"),
    dtype=torch.float32
)

di = hl.zeros(
    [tile_m],
    dtype=torch.float32
)

Mỗi hàng đang được xử lý có hai giá trị trạng thái:

Biến	Ý nghĩa
`mi`	Giá trị lớn nhất đã thấy trong các block cột trước
`di`	Tổng mũ đã được hiệu chỉnh theo giá trị lớn nhất hiện tại

mi được khởi tạo bằng:

[
-\infty
]

vì chưa đọc giá trị nào.

di được khởi tạo bằng:

[
0
]

Cả hai dùng float32 dù đầu vào là float16.

Điều này rất quan trọng để tăng độ chính xác khi:

tìm max;
tính hàm mũ;
cộng tổng nhiều phần tử.

12. Lượt thứ nhất: tính max và mẫu số

for tile_n in hl.tile(n, block_size=block_size_n):

Chiều cột được chia thành các block.

Giả sử:

n = 2560
block_size_n = 256

thì mỗi hàng có 10 block cột:

12.1 Đọc một block dữ liệu

values = x[tile_m, tile_n]

Nếu:

tile_m có 8 hàng
tile_n có 256 cột

thì values có shape logic:

[8, 256]

12.2 Tìm max cục bộ

local_amax = torch.amax(values, dim=1)

Tìm giá trị lớn nhất trên từng hàng trong block cột hiện tại.

Nếu:

values shape = [8, 256]

thì:

local_amax shape = [8]

Mỗi hàng có một max cục bộ.

12.3 Cập nhật max toàn cục

mi_next = torch.maximum(mi, local_amax)

So sánh:

mi: max của tất cả block đã xử lý;
local_amax: max của block hiện tại.

Kết quả:

[
m_{\text{next}}

\max(m_{\text{old}},m_{\text{local}})
]

13. Công thức cập nhật `di`

Đây là dòng quan trọng nhất:

di = di * torch.exp(mi - mi_next) + torch.exp(
    values - mi_next[:, None]
).sum(dim=1)

Để hiểu, trước tiên hãy xét mục tiêu.

Sau khi đọc tất cả các cột, ta muốn:

[
d_i = \sum_j e^{x_{ij}-m_i}
]

trong đó:

[
m_i=\max_j x_{ij}
]

Nhưng dữ liệu được đọc theo từng block. Khi đọc block mới, giá trị max có thể thay đổi.

13.1 Phần mới

torch.exp(values - mi_next[:, None]).sum(dim=1)

Tính:

[
\sum_{x\in\text{block mới}}e^{x-m_{\text{next}}}
]

mi_next có shape:

[tile_m]

Còn values có shape:

[tile_m, tile_n]

Cần thêm một chiều:

mi_next[:, None]

để có shape:

[tile_m, 1]

Sau đó broadcasting:

values              [tile_m, tile_n]
mi_next[:, None]    [tile_m,      1]
                    ─────────────────
kết quả             [tile_m, tile_n]

13.2 Hiệu chỉnh phần tổng cũ

di * torch.exp(mi - mi_next)

Trước đó, di được biểu diễn theo max cũ:

[
d_{\text{old}}

\sum_{\text{cũ}}e^{x-m_{\text{old}}}
]

Nhưng giờ max mới là (m_{\text{next}}).

Ta cần đổi hệ quy chiếu:

[
e^{x-m_{\text{next}}}

e^{x-m_{\text{old}}}
\cdot
e^{m_{\text{old}}-m_{\text{next}}}
]

Do đó:

[
d_{\text{old adjusted}}

d_{\text{old}}
\cdot
e^{m_{\text{old}}-m_{\text{next}}}
]

Đó chính là:

di * torch.exp(mi - mi_next)

13.3 Công thức đầy đủ

[
d_{\text{next}}

d_{\text{old}}
e^{m_{\text{old}}-m_{\text{next}}}
+
\sum_{\text{block mới}}
e^{x-m_{\text{next}}}
]

Đây thường được gọi là online softmax recurrence.

Nó cho phép tính max và tổng mũ trong cùng một lượt đọc dữ liệu, thay vì phải:

đọc toàn bộ để tìm max;
đọc lại để tính tổng.

Tuy nhiên toàn bộ kernel vẫn được gọi là “two-pass” vì còn một lượt thứ hai để ghi kết quả cuối cùng.

13.4 Cập nhật max

mi = mi_next

Sau khi cập nhật tổng, lưu max mới để xử lý block tiếp theo.

14. Ví dụ số cho công thức cập nhật

Giả sử một hàng được chia làm hai block.

Block thứ nhất:

[
[1,2]
]

Ta có:

[
m_{\text{old}}=2
]

[
d_{\text{old}}

e^{1-2}+e^{2-2}

e^{-1}+1
]

Xử lý block thứ hai:

[
[3,4]
]

Max mới:

[
m_{\text{next}}=4
]

Phần tổng cũ phải được chuyển từ hệ quy chiếu 2 sang hệ quy chiếu 4:

[
d_{\text{old adjusted}}

d_{\text{old}}e^{2-4}
]

Phần block mới:

[
e^{3-4}+e^{4-4}
]

Tổng cuối:

[
d_{\text{next}}

(e^{-1}+1)e^{-2}+e^{-1}+1
]

Rút gọn:

[
d_{\text{next}}

e^{-3}+e^{-2}+e^{-1}+1
]

Đúng bằng:

[
\sum_{x\in[1,2,3,4]}e^{x-4}
]

15. Kết thúc lượt thứ nhất

Sau vòng lặp đầu:

for tile_n in hl.tile(...):

ta có, với mỗi hàng:

mi = giá trị lớn nhất của toàn bộ hàng
di = tổng exp(x - mi) của toàn bộ hàng

Tức là đã có đầy đủ:

số trừ ổn định số học;
mẫu số của Softmax.

Nhưng chưa ghi kết quả đầu ra.

16. Lượt thứ hai: tính kết quả Softmax

for tile_n in hl.tile(n, block_size=block_size_n):

Kernel duyệt lại từng block cột.

values = x[tile_m, tile_n]

Đọc lại dữ liệu đầu vào.

Sau đó:

out[tile_m, tile_n] = (
    torch.exp(values - mi[:, None])
    / di[:, None]
)

Đây chính là công thức:

[
\operatorname{softmax}(x)

\frac{e^{x-m}}{d}
]

với:

[
d=\sum_j e^{x_j-m}
]

Broadcasting:

values          [block_rows, block_cols]
mi[:, None]     [block_rows,          1]
di[:, None]     [block_rows,          1]

Mỗi hàng dùng một giá trị max và một mẫu số riêng.

17. Vì sao gọi là hai lượt?

Lượt 1

Đọc x
→ tìm max
→ tính tổng exp đã chuẩn hóa

Lượt 2

Đọc lại x
→ tính exp(x - max) / tổng
→ ghi out

Sơ đồ:

                 Lượt 1
x ───────────────────────────────────┐
                                    │
       max từng hàng + mẫu số        │
                                    ▼
                               mi và di
                                    │
                 Lượt 2             │
x ───────────────────────────────────┤
                                    ▼
                         exp(x-mi) / di
                                    │
                                    ▼
                                   out

18. So sánh hai kernel

Tiêu chí	`softmax`	`softmax_two_pass`
Mức trừu tượng	Cao	Thấp hơn
Cách viết	Gọi PyTorch Softmax	Tự triển khai công thức
Chia tile theo hàng	Có	Có
Chia tile theo cột	Helion tự xử lý	Viết rõ trong code
Ổn định số học	Do PyTorch/Inductor xử lý	Thể hiện rõ bằng `mi`, `di`
Hỗ trợ hàng rất rộng	Phụ thuộc compiler	Chủ động loop theo block cột
Dễ đọc	Dễ hơn	Khó hơn
Kiểm soát thuật toán	Ít hơn	Nhiều hơn

19. Điểm đáng chú ý về GPU

Một hàng có 2560 phần tử float16.

Dung lượng thô:

[
2560\times2=5120\text{ byte}
]

Nhưng trong quá trình tính Softmax, nhiều giá trị thường được chuyển sang float32.

Nếu cố giữ toàn bộ dữ liệu trung gian của nhiều hàng trong register, áp lực register có thể cao.

Kernel hai lượt cho phép xử lý chiều cột theo từng phần:

256 cột
→ 256 cột
→ 256 cột
→ ...

Thay vì bắt buộc xử lý toàn bộ 2560 cột cùng lúc.

Lợi ích tiềm năng:

giảm register pressure;
hỗ trợ chiều reduction lớn;
cho autotuner lựa chọn kích thước block;
phù hợp hơn với nhiều kích thước tensor.

Đổi lại, đầu vào phải được đọc hai lần.

20. Một Helion kernel vẫn chỉ là một GPU kernel

Mặc dù trong softmax_two_pass có ba vòng lặp:

for tile_m in hl.tile(...):
    for tile_n in hl.tile(...):
        ...
    for tile_n in hl.tile(...):
        ...

không có nghĩa là tạo ba lần gọi kernel GPU.

Ý nghĩa gần đúng:

Vòng lặp	Cách Helion ánh xạ
`tile_m` bên ngoài	Ánh xạ lên grid/program GPU
`tile_n` lượt 1	Vòng lặp nội bộ trong mỗi program
`tile_n` lượt 2	Vòng lặp nội bộ thứ hai trong cùng program

Một program GPU xử lý một tile hàng:

Program 0 → một nhóm hàng
Program 1 → nhóm hàng tiếp theo
Program 2 → nhóm hàng tiếp theo
...

Trong mỗi program:

lặp qua các block cột lần 1
lặp qua các block cột lần 2

21. Mã nguồn có phần bị lặp

Trong nội dung bạn gửi, phần sau xuất hiện hai lần:

from __future__ import annotations
import torch
import helion
from helion._testing import run_example
import helion.language as hl

Tiêu đề chuỗi tài liệu cũng bị lặp:

"""
Helion Softmax Kernel Examples
...
"""

Đây không phải lỗi nghiêm trọng, nhưng có thể gộp thành một file rõ ràng hơn.

22. Phiên bản tổ chức lại dễ đọc hơn

from __future__ import annotations

import torch
import helion
import helion.language as hl
from helion._testing import run_example


@helion.kernel(autotune_effort="quick")
def softmax(x: torch.Tensor) -> torch.Tensor:
    num_rows, _num_cols = x.size()
    out = torch.empty_like(x)

    for row_tile in hl.tile(num_rows):
        out[row_tile, :] = torch.nn.functional.softmax(
            x[row_tile, :],
            dim=1,
        )

    return out


@helion.kernel(autotune_effort="quick")
def softmax_two_pass(x: torch.Tensor) -> torch.Tensor:
    num_rows, num_cols = x.size()
    out = torch.empty_like(x)

    row_block_size = hl.register_block_size(num_rows)
    col_block_size = hl.register_block_size(num_cols)

    for row_tile in hl.tile(
        num_rows,
        block_size=row_block_size,
    ):
        running_max = hl.full(
            [row_tile],
            float("-inf"),
            dtype=torch.float32,
        )

        running_sum = hl.zeros(
            [row_tile],
            dtype=torch.float32,
        )

        # Lượt 1: tìm max và mẫu số.
        for col_tile in hl.tile(
            num_cols,
            block_size=col_block_size,
        ):
            values = x[row_tile, col_tile]

            local_max = torch.amax(values, dim=1)
            next_max = torch.maximum(running_max, local_max)

            running_sum = (
                running_sum
                * torch.exp(running_max - next_max)
                + torch.exp(
                    values - next_max[:, None]
                ).sum(dim=1)
            )

            running_max = next_max

        # Lượt 2: chuẩn hóa và ghi đầu ra.
        for col_tile in hl.tile(
            num_cols,
            block_size=col_block_size,
        ):
            values = x[row_tile, col_tile]

            out[row_tile, col_tile] = (
                torch.exp(values - running_max[:, None])
                / running_sum[:, None]
            )

    return out


def check(kernel, rows: int, cols: int) -> None:
    x = torch.randn(
        [rows, cols],
        device="cuda",
        dtype=torch.float16,
    )

    run_example(
        kernel,
        lambda value: torch.nn.functional.softmax(value, dim=1),
        (x,),
    )


def main() -> None:
    check(softmax, 4096, 2560)
    check(softmax_two_pass, 4096, 2560)


if __name__ == "__main__":
    main()

23. Tóm tắt luồng thực thi

Kernel đơn giản

Input x
  ↓
Chia theo nhóm hàng
  ↓
PyTorch softmax trên từng nhóm hàng
  ↓
Output

Kernel hai lượt

Input x
  ↓
Chia theo nhóm hàng
  ↓
Với mỗi nhóm hàng:
  ├─ Chia chiều cột thành block
  ├─ Lượt 1: tính running max và running sum
  ├─ Lượt 2: tính xác suất Softmax
  └─ Ghi output

Ý tưởng quan trọng nhất của softmax_two_pass là công thức cập nhật online:

[
m_{\text{new}}

\max(m_{\text{old}},m_{\text{block}})
]

[
d_{\text{new}}

d_{\text{old}}e^{m_{\text{old}}-m_{\text{new}}}
+
\sum_{\text{block}}e^{x-m_{\text{new}}}
]

Công thức này giúp tính Softmax ổn định số học khi chiều cột lớn mà không cần giữ toàn bộ hàng trong register cùng lúc.

Helion language

Bài 1. Helion ra đời để giải quyết vấn đề gì?

Vì sao Triton vẫn còn khó?

Ý tưởng của Helion

Helion là gì?

Một kernel Helion gồm hai phần

Phần 1

Phần 2

hl.tile() là gì?

Ví dụ Matrix Multiplication

Điều đặc biệt của Helion

Autotuning là gì?

Tại sao mất tới 10 phút?

Differential Evolution

Config là gì?

Tại sao nên lưu Config?

Settings khác Config như thế nào?

Quá trình biên dịch Helion

Vì sao Helion đáng chú ý?

Lộ trình học Helion tôi đề xuất

1. Softmax là gì?

[\operatorname{softmax}(x_i)

2. Cài đặt thư viện

3. Các thư viện được import

4. Kernel Softmax thứ nhất

4.1 Lấy kích thước tensor

4.2 Tạo tensor đầu ra

4.3 Chia hàng thành các tile

4.4 Tính Softmax trên các hàng thuộc tile

4.5 Trả kết quả

5. Hàm kiểm tra kernel thứ nhất

5.1 Tạo dữ liệu kiểm thử

5.2 So sánh Helion với PyTorch

6. Hàm main

[10{,}485{,}760 \times 2

6.1 Kiểm tra chương trình có được chạy trực tiếp không

7. Kernel thứ hai: softmax_two_pass

8. Tại sao cần Softmax ổn định số học?

[\operatorname{softmax}(x_i)

9. Chuẩn bị trong softmax_two_pass

9.1 Đăng ký block size

10. Vòng lặp theo tile hàng

11. Hai biến trạng thái mi và di

12. Lượt thứ nhất: tính max và mẫu số

12.1 Đọc một block dữ liệu

12.2 Tìm max cục bộ

12.3 Cập nhật max toàn cục

[m_{\text{next}}

13. Công thức cập nhật di

13.1 Phần mới

13.2 Hiệu chỉnh phần tổng cũ

[d_{\text{old}}

[e^{x-m_{\text{next}}}

[d_{\text{old adjusted}}

13.3 Công thức đầy đủ

[d_{\text{next}}

13.4 Cập nhật max

14. Ví dụ số cho công thức cập nhật

[d_{\text{old}}

e^{1-2}+e^{2-2}

[d_{\text{old adjusted}}

[d_{\text{next}}

[d_{\text{next}}

15. Kết thúc lượt thứ nhất

16. Lượt thứ hai: tính kết quả Softmax

[\operatorname{softmax}(x)

17. Vì sao gọi là hai lượt?

Lượt 1

Lượt 2

18. So sánh hai kernel

19. Điểm đáng chú ý về GPU

20. Một Helion kernel vẫn chỉ là một GPU kernel

21. Mã nguồn có phần bị lặp

22. Phiên bản tổ chức lại dễ đọc hơn

23. Tóm tắt luồng thực thi

Kernel đơn giản

Kernel hai lượt

[m_{\text{new}}

[d_{\text{new}}

Bình luận

[
\operatorname{softmax}(x_i)

6. Hàm `main`

[
10{,}485{,}760 \times 2

7. Kernel thứ hai: `softmax_two_pass`

[
\operatorname{softmax}(x_i)

9. Chuẩn bị trong `softmax_two_pass`

11. Hai biến trạng thái `mi` và `di`

[
m_{\text{next}}

13. Công thức cập nhật `di`

[
d_{\text{old}}

[
e^{x-m_{\text{next}}}

[
d_{\text{old adjusted}}

[
d_{\text{next}}

[
d_{\text{old}}

[
d_{\text{old adjusted}}

[
d_{\text{next}}

[
d_{\text{next}}

[
\operatorname{softmax}(x)

[
m_{\text{new}}

[
d_{\text{new}}