Skip to content

Standardize softmax.py to avoid numpy dependency #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hannahli-nv merged 4 commits into from
Dec 23, 2025
Merged

Standardize softmax.py to avoid numpy dependency #16

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
58e348a
Update softmax.py
lessw2020 Dec 21, 2025
ef173a2
fomat
lessw2020 Dec 21, 2025
17ed210
use ct.float32
lessw2020 Dec 22, 2025
fb27207
use ct.int32 instead of torch.int32
lessw2020 Dec 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions src/tilegym/ops/cutile/softmax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
# SPDX-License-Identifier: MIT


import math

import cuda.tile as ct
import numpy as np
import torch

from tilegym.backend import register_impl
Expand All @@ -26,13 +27,13 @@ def softmax_kernel(
# Static persistent scheduling: each block processes multiple rows
pid = ct.bid(0)
num_programs = ct.num_blocks(0)
offsets = ct.arange(TILE_SIZE, dtype=torch.int32)
offsets = ct.arange(TILE_SIZE, dtype=ct.int32)

for row_idx in range(pid, n_rows, num_programs):
# Load the row tile using index-based access
row = ct.gather(input, (row_idx, offsets), check_bounds=True, padding_value=-np.inf)
row = ct.gather(input, (row_idx, offsets), check_bounds=True, padding_value=-math.inf)
# Convert to float32 for computation
row = ct.astype(row, torch.float32)
row = ct.astype(row, ct.float32)

# Subtract maximum for numerical stability
row_max = ct.max(row, 0, keepdims=True)
Expand Down Expand Up @@ -72,7 +73,7 @@ def softmax_kernel_tma(
row = ct.load(input, index=(row_idx, 0), shape=(1, TILE_SIZE), padding_mode=ct.PaddingMode.NEG_INF)

# Convert to float32 for computation
row = ct.astype(row, np.float32)
row = ct.astype(row, ct.float32)

# Subtract maximum for numerical stability
row_max = ct.max(row, 1, keepdims=True)
Expand Down