Helion — PyTorch-native 한 ML 커널 DSL

2024–2025 사이 PyTorch 진영의 큰 흐름 하나는 — 커널 작성을 사용자에게서 컴파일러에게로 다시 끌어오는 방향이다. Helion 은 그 흐름의 가장 새로운 한 점. Triton 이 “파이썬으로 GPU 커널을 짜자” 였다면, Helion 은 “PyTorch 코드처럼 보이는 커널 DSL” 이 출발점이다.

강의는 PyTorch Labs 의 세 사람 — torch.compile / Inductor 의 핵심을 짠 Jason Ansel, Triton 통합을 끌고 온 Oguz Ulgen, 그리고 Will Feng — 이 함께 발표한 형태다. 자막이 실패해 화면에서 직접 코드를 본 디테일은 잡히지 않지만, 외부에 공개된 Helion 의 설계 문서와 PyTorch 컴파일러 진영의 일관된 이야기 흐름에서 다음 세 줄을 읽어낼 수 있다.

1. Triton 은 “파이썬으로 짜는 SASS” 에 가깝다 — 표현력은 강력하지만, tl.load / tl.store / BLOCK_SIZE / mask 같은 GPU-네이티브 개념을 매번 손으로 다뤄야 한다.
2. PyTorch 사용자는 그렇게 짜고 싶지 않다 — 텐서 슬라이싱, 브로드캐스팅, 같은 indexing 의미가 커널 안에서도 그대로 작동했으면 한다.
3. 그렇다고 torch.compile 만으로는 부족하다 — 컴파일러가 자동으로 잘 만드는 것 이상이 필요한 자리(특수한 attention 변형, 새 양자화 커널, custom fused op)는 항상 남는다.

Helion 은 이 셋 사이의 빈 자리를 메우려 한다 — “직접 커널을 짜되, PyTorch 의 의미가 그대로 살아 있고, 결국에는 Triton 으로 lowering 되어 같은 백엔드 위에서 돈다”. 새 컴파일러를 처음부터 짜는 게 아니라, Triton 위에 얹는 더 친숙한 layer 라는 위치가 핵심이다.

Triton 으로 커널을 한 100번 짠 사람이라면 같은 패턴이 반복된다는 걸 안다 — pid 계산, offsets 만들기, mask 처리, BLOCK_SIZE 상수 정의. 코드의 의미는 “이 텐서의 이 슬라이스를 그 텐서의 저 슬라이스에 더한다” 인데, 표현은 그것보다 한 층 아래에 머문다.

Helion 팀의 진단 — “이 다섯이 매번 등장한다는 건 DSL 한 단계 더 위 에서 자동화될 수 있는 표시” 라는 것. 그래서 Helion 은 Triton 의 lowering 백엔드를 그대로 두되, 사용자에게 보이는 표면을 PyTorch 텐서 의미로 가져온다.

Helion 의 가장 큰 차별점은 “DSL 안의 변수가 PyTorch 텐서의 의미를 따른다” 는 것. 같은 사용자가 같은 인덱싱을 두 군데서 다르게 외울 필요가 없다는 게 핵심.

# row-wise softmax (대략)
@triton.jit
def softmax_kernel(x_ptr, y_ptr, n_cols,
                   stride_xm, stride_ym,
                   BLOCK: tl.constexpr):
    pid    = tl.program_id(0)
    cols   = tl.arange(0, BLOCK)
    mask   = cols < n_cols
    x_row  = x_ptr + pid * stride_xm
    x      = tl.load(x_row + cols, mask=mask,
                     other=-float("inf"))
    x_max  = tl.max(x, axis=0)
    x      = x - x_max
    e      = tl.exp(x)
    z      = tl.sum(e, axis=0)
    y      = e / z
    y_row  = y_ptr + pid * stride_ym
    tl.store(y_row + cols, y, mask=mask)

# 강의 시연을 외부 자료로 재구성 — 확인 필요
@helion.kernel
def softmax(x: Tensor) -> Tensor:
    # tile 단위 loop — 행 분할은 컴파일러에게 맡김
    for tile_m in hl.tile(x.size(0)):
        row    = x[tile_m, :]            # PyTorch slice
        x_max  = row.amax(dim=-1, keepdim=True)
        e      = (row - x_max).exp()
        out    = e / e.sum(dim=-1, keepdim=True)
        y[tile_m, :] = out
    return y

두 코드가 같은 일을 한다 — 그러나 표현이 “다른 언어” 다. Helion 의 코드 안에서 row.amax(dim=-1) 는 PyTorch 의 그것과 같은 의미다. tl.max 와 의미가 다른 별도의 함수가 아니다. 컴파일러가 알아서 tile 안에서의 reduction 으로 lowering 한다.

이 디자인이 사용자에게 주는 효과는 두 가지로 정리된다.

• 학습 곡선 단축 — 이미 PyTorch 를 아는 사용자는 “tile loop 한 줄” 만 새로 배우면 된다.
• 검증의 단순화 — 같은 코드를 numerical 비교용 PyTorch reference 와 그대로 실행할 수 있을 가능성(원리적으로). 강의에서 demo 형태로 보여줬을 가능성이 큼 — 확인 필요.

Helion 코드가 GPU 위에서 어떻게 도는지 한 번에 그려두는 게 빠르다. 자기 컴파일러를 다시 짜는 대신, 기존 인프라(Triton, Inductor, MLIR) 를 받아쓰는 형태가 일관된다.

이 그림에서 짚어야 하는 디자인 결정 세 가지.

1. L1 의 “텐서 의미 그대로” 가 가장 비싸다 — Helion IR 가 단순히 Triton AST 의 wrapper 였다면 의미가 없다. 슬라이스 / 브로드캐스트 / dtype 추론이 IR 수준에서 살아 있어야 L2 의 search 가 의미 있는 변형을 만들 수 있다.
2. L2 가 사용자 코드에 손 안 대고 launch shape 을 결정한다 — 같은 Helion 코드가 hidden_dim=64 와 4096 에서 다른 BLOCK 으로 lowering 된다.
3. L3 의 출력은 “읽을 수 있는 Triton” — 디버깅과 후속 최적화를 위해, Helion 이 만든 Triton 코드를 사람이 펼쳐 볼 수 있어야 한다는 디자인. TORCH_LOGS=output_code 의 사상과 같다.

강의의 마지막 데모로 등장했을 가능성이 가장 높은 예시 — FlashAttention-스타일 의 tiled softmax + masked attention. Triton 진영의 거의 모든 새 DSL 이 이걸 첫 데모로 쓰는 이유는 — 충분히 복잡해서 표현력의 차이가 드러난다.

@triton.jit
def attn_fwd(Q, K, V, O, M, ...,
             BLOCK_M: tl.constexpr,
             BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m*BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)

    q = tl.load(Q + offs_m[:,None]*sQm + ...,
                mask=...)
    m_i = tl.full([BLOCK_M], -float("inf"))
    l_i = tl.zeros([BLOCK_M])
    acc = tl.zeros([BLOCK_M, D])

    for start_n in range(0, N, BLOCK_N):
        k = tl.load(K + ..., mask=...)
        v = tl.load(V + ..., mask=...)
        s = tl.dot(q, tl.trans(k))
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp(m_i - m_new)
        p = tl.exp(s - m_new[:,None])
        l_i = l_i*alpha + tl.sum(p, axis=1)
        acc = acc*alpha[:,None] + tl.dot(p, v)
        m_i = m_new

    o = acc / l_i[:,None]
    tl.store(O + ..., o, mask=...)

@helion.kernel
def attn_fwd(Q, K, V):
    O = torch.empty_like(Q)
    for tile_m in hl.tile(Q.size(0)):
        q   = Q[tile_m, :]
        m_i = torch.full([q.size(0)], -inf)
        l_i = torch.zeros_like(m_i)
        acc = torch.zeros_like(q)

        for tile_n in hl.tile(K.size(0)):
            k = K[tile_n, :]
            v = V[tile_n, :]
            s = q @ k.transpose(-1, -2)

            # online softmax — PyTorch op 그대로
            m_new = torch.maximum(m_i, s.amax(-1))
            alpha = (m_i - m_new).exp()
            p     = (s - m_new[..., None]).exp()
            l_i   = l_i*alpha + p.sum(-1)
            acc   = acc*alpha[..., None] + p @ v
            m_i   = m_new

        O[tile_m, :] = acc / l_i[..., None]
    return O

두 코드가 의도하는 알고리즘은 같다 — online softmax 로 score matrix 를 통째로 들지 않으면서 attention 을 계산. 차이는 코드의 밀도다. Helion 쪽에서 tl.load/tl.store/mask 의 boilerplate 가 사라진 자리에 알고리즘 자체가 더 잘 보인다.

Triton 의 @triton.autotune(configs=[...]) 를 한 번이라도 손으로 깔아본 사람은 안다 — configs 가 점점 길어지고, 새 모양이 등장할 때마다 다시 손을 봐야 한다. Helion 은 이 자리를 컴파일러 안으로 끌어들인다.

실전 사용에서 Helion 의 autotune 이 의미를 갖는 자리는 shape 가 데이터에 따라 달라지는 워크로드다 — variable-length attention, KV cache 의 동적 크기, batch 별 다른 hidden_dim. Triton 의 정적 configs 에서 이런 워크로드가 가장 어렵게 다뤄지던 자리.

L001 의 의사결정 사다리(torch → Triton → CUDA) 에 한 칸이 추가된다 — torch → Helion → Triton → CUDA. 각 칸의 trigger 가 무엇인지 정리.

강의에서 팀이 명시적으로 강조한 입장 — “Helion 은 Triton 을 안 짜기 위해서가 아니라, Triton 을 더 잘 짤 수 있는 자리를 분명히 하기 위해서다”. 확인 필요 — 실제 발화는 자막이 없어 정확하지 않지만, 팀의 외부 발표에서 일관되게 등장하는 입장.

새 DSL 의 미래는 사용자 수가 결정한다. Helion 의 초기 채택 자리는 — 외부 자료에서 확인 가능한 한도 안에서 — 두 곳으로 모인다.

• PyTorch 의 내부 op 구현 — torch.compile 의 Inductor 가 만들어내는 fused 커널과 별개로, “손으로 잘 짜야 하는” op 들(특수 attention 변형, rotary embedding fused, RMSNorm fused 등) 의 차세대 구현체로 Helion 이 후보.
• 외부 모델 코드의 hot path — vLLM, SGLang 같은 추론 엔진의 attention 변형, 또는 학습 측의 LIGER 류 라이브러리에서 — 확인 필요. 강의 발표 시점이 비교적 새로워 publicly 안 보이는 채택이 더 많을 수 있다.

실용적 관점 — 새 사용자가 Helion 을 처음 만질 때의 진입점은 같은 함수를 PyTorch eager / torch.compile / Helion 세 가지로 짜고 timing 비교다. 커널 작성 학습 곡선의 처음을 그래프 한 장으로 바로 잡을 수 있는 비교.

강의의 Q&A 에서 등장했을 법한 — 그리고 외부 자료에서 추정 가능한 — 향후 방향을 묶어둔다.

1. backward pass 자동화 — forward 만 짠 Helion 커널에서 backward 를 자동 생성하는 길. PyTorch 의 autograd 와 어떻게 합쳐질지가 design space.
2. multi-GPU primitives — Triton 자체가 single-GPU 모델인 것처럼, Helion 도 그렇다. L087 NVSHMEM 의 영역과 어떻게 맞물릴지가 다음 라운드.
3. FP8 / MXFP4 등 새 dtype — L084 Numerics 가 깐 흐름을 Helion 의 표면이 어떻게 받아낼지. 양자화 커널은 Helion 의 자연스러운 다음 표적.
4. AMD / 다른 backend — Triton 자체가 멀티 백엔드(NVIDIA, AMD, Intel) 로 가는 중이라, Helion 은 그 위에서 자동으로 받을 가능성이 높다.
5. IDE / 디버거 통합 — 자동 lowering 된 Triton 코드와 원래 Helion 코드를 동시에 보는 도구. 사용자가 “내 코드의 어느 줄이 어떤 PTX 를 만드는지” 추적 가능해야 학습 자료로도 쓰인다.

손에 새기기 — 실습 시퀀스

1. 같은 함수 3 가지 구현 비교 — softmax 또는 RMSNorm 을 ① PyTorch eager ② torch.compile ③ Helion 으로 짜고, 각 GPU 시간을 측정. torch.cuda.Event + warmup 5회 + sync 의 L001 패턴 그대로.
2. autotune 효과 보기 — 같은 Helion 커널에 대해 작은 입력 (1024×1024) 과 큰 입력 (8192×8192) 의 캐시된 Triton 코드가 어떻게 다른지 dump 해 비교.
3. FlashAttention forward 짜기 — § 05 의 형태를 그대로 따라 Helion 에서 attention forward 를 짠다. PyTorch reference 와 numerical 비교.
4. Helion → Triton dump — 컴파일러가 만든 Triton 코드를 직접 펼쳐 보고, 어떤 변형이 선택됐는지 추적. HELION_LOGS 같은 환경변수 — 확인 필요.
5. 의사결정 사다리 적용 — 자기 모델 코드에서 한 hot path 를 골라 Helion 으로 교체할 후보를 정한다. NCU 가 무슨 hint 를 주는지가 갈림길.

• 강의 안의 정확한 데모 시퀀스 — 자막이 실패해 발표자가 라이브로 보여준 정확한 코드 시퀀스, 측정 수치, 비교 대상이 다 잡히지 않았다. 직접 영상을 다시 보면 보강 필요.
• Helion repo 의 현재 상태 — 이 노트는 외부 자료 + 도메인 지식 기반. pytorch-labs/helion 의 README, 예시 디렉토리, CHANGELOG 를 직접 한 번 펼쳐 볼 것.
• backward pass 의 위치 — forward 만 짜면 autograd 가 자동 잡아주는지, 아니면 backward 도 직접 짜야 하는지 — 확인 필요.
• FP8 / MXFP4 dtype 지원 — § 09 에서 언급한 양자화 dtype 의 표면이 현재 어디까지 지원되는지.
• Inductor 와의 future 관계 — 둘이 합쳐질 가능성, 아니면 Inductor 가 Helion 을 backend 로 쓰는 가능성 — design 문서에서 확인.
• 실제 production 채택 — vLLM / SGLang / 학습 라이브러리에서 Helion 으로 짠 커널이 머지된 사례. issue tracker 에서 추적 가능.