Skip to content

Commit e9c7b15

Browse files
tsingmicro-public-etsingmicro-public-e
committed
code format
1 parent 917ca9d commit e9c7b15

52 files changed

Lines changed: 2583 additions & 2269 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎third_party/tsingmicro/crt/test/oob/conftest.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,9 @@ def run_kernel_script(script_content: str, tmp_path: Path, caller: str = "") ->
8686
# 1. Run the kernel
8787
proc = subprocess.run(
8888
[sys.executable, str(script_file)],
89-
capture_output=True, text=True, timeout=120,
89+
capture_output=True,
90+
text=True,
91+
timeout=120,
9092
cwd=tmp_path,
9193
env=os.environ.copy(),
9294
)
@@ -98,9 +100,8 @@ def run_kernel_script(script_content: str, tmp_path: Path, caller: str = "") ->
98100
# 2. Extract with test name suffix, then parse
99101
subprocess.run(['sleep', '3'])
100102
extract_sh = os.path.join(workspace, "third_party/tsingmicro/scripts/extract_kcore.sh")
101-
proc = subprocess.run(["bash", extract_sh, "all", "-s", caller],
102-
capture_output=True, text=True, timeout=60,
103-
cwd=workspace)
103+
proc = subprocess.run(["bash", extract_sh, "all", "-s", caller], capture_output=True, text=True, timeout=60,
104+
cwd=workspace)
104105

105106
# if proc.stdout:
106107
# print(proc.stdout, flush=True)

‎third_party/tsingmicro/crt/test/oob/kernels/dropout_oob.py‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,9 @@
77

88
@triton.jit
99
def dropout_kernel_oob(x_ptr, x_keep_ptr, output_ptr, n_elements, p,
10-
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask on load/store
11-
inject_n_elements_scale: tl.constexpr, # 1=normal, 2=double n_elements
12-
BLOCK_SIZE: tl.constexpr,
13-
seed: tl.constexpr):
10+
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask on load/store
11+
inject_n_elements_scale: tl.constexpr, # 1=normal, 2=double n_elements
12+
BLOCK_SIZE: tl.constexpr, seed: tl.constexpr):
1413
pid = tl.program_id(axis=0)
1514
block_start = pid * BLOCK_SIZE
1615
offsets = block_start + tl.arange(0, BLOCK_SIZE)
@@ -28,11 +27,8 @@ def dropout_kernel_oob(x_ptr, x_keep_ptr, output_ptr, n_elements, p,
2827

2928

3029
@triton.jit
31-
def seeded_dropout_kernel_oob(x_ptr, output_ptr, n_elements, p,
32-
inject_mask_mode: tl.constexpr,
33-
inject_n_elements_scale: tl.constexpr,
34-
BLOCK_SIZE: tl.constexpr,
35-
seed: tl.constexpr):
30+
def seeded_dropout_kernel_oob(x_ptr, output_ptr, n_elements, p, inject_mask_mode: tl.constexpr,
31+
inject_n_elements_scale: tl.constexpr, BLOCK_SIZE: tl.constexpr, seed: tl.constexpr):
3632
pid = tl.program_id(axis=0)
3733
block_start = pid * BLOCK_SIZE
3834
offsets = block_start + tl.arange(0, BLOCK_SIZE)

‎third_party/tsingmicro/crt/test/oob/kernels/dsa_vec_add_oob.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ def add_kernel_dsa_oob(
1313
output_ptr,
1414
n_elements,
1515
inject_n_elements_scale: tl.constexpr, # 1=normal, >1=inflate n_elements
16-
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask
17-
inject_n_iter_scale: tl.constexpr, # 1=normal, >1=more iters beyond BLOCK_SIZE
16+
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask
17+
inject_n_iter_scale: tl.constexpr, # 1=normal, >1=more iters beyond BLOCK_SIZE
1818
BLOCK_SIZE: tl.constexpr,
1919
):
2020
pid = tl.program_id(axis=0)

‎third_party/tsingmicro/crt/test/oob/kernels/extern_oob.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,9 @@
77

88

99
@triton.jit
10-
def asin_kernel_oob(x_ptr, y_ptr, n_elements,
11-
inject_pid_shift: tl.constexpr, # 0=normal, N=shift pid by +N
12-
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask
13-
BLOCK_SIZE: tl.constexpr):
10+
def asin_kernel_oob(x_ptr, y_ptr, n_elements, inject_pid_shift: tl.constexpr, # 0=normal, N=shift pid by +N
11+
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask
12+
BLOCK_SIZE: tl.constexpr):
1413
pid = tl.program_id(axis=0) + inject_pid_shift
1514
block_start = pid * BLOCK_SIZE
1615
offsets = block_start + tl.arange(0, BLOCK_SIZE)

‎third_party/tsingmicro/crt/test/oob/kernels/gather_scatter_oob.py‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,21 @@
99

1010
@triton.jit
1111
def gather_scatter_kernel_oob(
12-
Offs, A, B, C,
13-
M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
14-
stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
15-
inject_index_oob: tl.constexpr, # 0=normal, 1=indices >= M (partial OOB)
16-
inject_index_oob_all: tl.constexpr, # 0=normal, 1=all indices = M + 16 (full OOB)
12+
Offs,
13+
A,
14+
B,
15+
C,
16+
M: tl.constexpr,
17+
N: tl.constexpr,
18+
K: tl.constexpr,
19+
stride_am,
20+
stride_ak,
21+
stride_bk,
22+
stride_bn,
23+
stride_cm,
24+
stride_cn,
25+
inject_index_oob: tl.constexpr, # 0=normal, 1=indices >= M (partial OOB)
26+
inject_index_oob_all: tl.constexpr, # 0=normal, 1=all indices = M + 16 (full OOB)
1727
BLOCK_M: tl.constexpr,
1828
):
1929
pid = tl.program_id(0)

‎third_party/tsingmicro/crt/test/oob/kernels/vec_add_oob.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@ def add_kernel_oob(
1111
y_ptr,
1212
output_ptr,
1313
n_elements,
14-
inject_pid_shift: tl.constexpr, # 0=normal, 1=shift pid by +1
15-
inject_block_scale: tl.constexpr, # 1=normal, 2=double block_size
16-
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask, 2=wrong cmp
14+
inject_pid_shift: tl.constexpr, # 0=normal, 1=shift pid by +1
15+
inject_block_scale: tl.constexpr, # 1=normal, 2=double block_size
16+
inject_mask_mode: tl.constexpr, # 0=normal, 1=no mask, 2=wrong cmp
1717
BLOCK_SIZE: tl.constexpr,
1818
):
1919
pid = tl.program_id(axis=0)

‎third_party/tsingmicro/crt/test/oob/test_attention_oob.py‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77
from conftest import run_kernel_script, DmaResult
88

99

10-
11-
def _create_test_script(Z, H, N_CTX, HEAD_DIM, inject_N_CTX_scale,
12-
inject_stride_scale, BLOCK_M=64, BLOCK_N=32):
10+
def _create_test_script(Z, H, N_CTX, HEAD_DIM, inject_N_CTX_scale, inject_stride_scale, BLOCK_M=64, BLOCK_N=32):
1311
"""Generate a self-contained test script for attention_oob.
1412
1513
The kernel uses block pointers (tl.make_block_ptr) for Q/K/V/O.
@@ -75,12 +73,12 @@ def _create_test_script(Z, H, N_CTX, HEAD_DIM, inject_N_CTX_scale,
7573
def main():
7674
"""Generate test scripts for all attention OOB test cases."""
7775
cases = {
78-
"test_normal_no_oob": _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
79-
inject_N_CTX_scale=1, inject_stride_scale=1),
80-
"test_n_ctx_inflated_oob": _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
81-
inject_N_CTX_scale=2, inject_stride_scale=1),
82-
"test_stride_corrupted_oob": _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
83-
inject_N_CTX_scale=1, inject_stride_scale=10),
76+
"test_normal_no_oob":
77+
_create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=1, inject_stride_scale=1),
78+
"test_n_ctx_inflated_oob":
79+
_create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=2, inject_stride_scale=1),
80+
"test_stride_corrupted_oob":
81+
_create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=1, inject_stride_scale=10),
8482
}
8583
out_dir = os.path.dirname(__file__)
8684
for name, script in cases.items():
@@ -95,26 +93,24 @@ def main():
9593

9694

9795
class TestAttentionOOB:
96+
9897
def test_normal_no_oob(self, dma_env, tmp_path):
9998
"""Normal parameters: should not trigger OOB detection."""
100-
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
101-
inject_N_CTX_scale=1, inject_stride_scale=1)
99+
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=1, inject_stride_scale=1)
102100
result = run_kernel_script(script, tmp_path, 'TestAttentionOOB.test_normal_no_oob')
103101
assert result.passed, f"Normal kernel should not OOB: oob={result.oob_count}"
104102

105103
def test_n_ctx_inflated_oob(self, dma_env, tmp_path):
106104
"""Inflate N_CTX by 2x: block pointers think there are 2x more rows,
107105
causing Q load and O store to go beyond actual tensor."""
108-
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
109-
inject_N_CTX_scale=2, inject_stride_scale=1)
106+
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=2, inject_stride_scale=1)
110107
result = run_kernel_script(script, tmp_path, 'TestAttentionOOB.test_n_ctx_inflated_oob')
111108
assert result.detected, \
112109
f"Expected OOB detection with N_CTX_scale=2: {result}"
113110

114111
def test_stride_corrupted_oob(self, dma_env, tmp_path):
115112
"""Inflate strides by 10x: addresses jump far beyond buffer."""
116-
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32,
117-
inject_N_CTX_scale=1, inject_stride_scale=10)
113+
script = _create_test_script(Z=1, H=2, N_CTX=64, HEAD_DIM=32, inject_N_CTX_scale=1, inject_stride_scale=10)
118114
result = run_kernel_script(script, tmp_path, 'TestAttentionOOB.test_stride_corrupted_oob')
119115
assert result.detected, \
120116
f"Expected OOB detection with stride_scale=0: {result}"

‎third_party/tsingmicro/crt/test/oob/test_dropout_oob.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@
77
from conftest import run_kernel_script, DmaResult
88

99

10-
def _create_test_script(n_elements, inject_mask_mode, inject_n_elements_scale,
11-
block_size=256, tensor_size=None):
10+
def _create_test_script(n_elements, inject_mask_mode, inject_n_elements_scale, block_size=256, tensor_size=None):
1211
"""Generate a self-contained test script for dropout_oob.
1312
1413
Args:
@@ -78,6 +77,7 @@ def main():
7877

7978

8079
class TestDropoutOOB:
80+
8181
def test_normal_no_oob(self, dma_env, tmp_path):
8282
"""Normal parameters: should not trigger OOB detection."""
8383
script = _create_test_script(2000, 0, 1)

‎third_party/tsingmicro/crt/test/oob/test_extern_oob.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77
from conftest import run_kernel_script, DmaResult
88

99

10-
11-
def _create_test_script(n_elements, inject_pid_shift, inject_mask_mode,
12-
block_size=256, tensor_size=None):
10+
def _create_test_script(n_elements, inject_pid_shift, inject_mask_mode, block_size=256, tensor_size=None):
1311
"""Generate a self-contained test script for extern_oob.
1412
1513
Args:
@@ -76,6 +74,7 @@ def main():
7674

7775

7876
class TestExternOOB:
77+
7978
def test_normal_no_oob(self, dma_env, tmp_path):
8079
"""Normal parameters: should not trigger OOB detection."""
8180
script = _create_test_script(1024, 0, 0)

‎third_party/tsingmicro/crt/test/oob/test_gather_scatter_oob.py‎

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,7 @@
1313
from conftest import run_kernel_script, DmaResult
1414

1515

16-
17-
def _create_test_script(M, N, K,
18-
inject_offs_scale=1,
19-
inject_index_oob=0,
20-
inject_index_oob_all=0,
21-
BLOCK_M=4):
16+
def _create_test_script(M, N, K, inject_offs_scale=1, inject_index_oob=0, inject_index_oob_all=0, BLOCK_M=4):
2217
"""Generate a self-contained test script for gather_scatter_oob.
2318
2419
Args:
@@ -79,12 +74,9 @@ def main():
7974
"""Generate test scripts for all gather_scatter OOB test cases."""
8075
cases = {
8176
"test_normal_no_oob": _create_test_script(M=64, N=32, K=128),
82-
"test_index_partial_oob": _create_test_script(M=64, N=32, K=128,
83-
inject_index_oob=1),
84-
"test_index_full_oob": _create_test_script(M=64, N=32, K=128,
85-
inject_index_oob_all=1),
86-
"test_indirect_index_oob": _create_test_script(M=64, N=32, K=128,
87-
inject_offs_scale=2),
77+
"test_index_partial_oob": _create_test_script(M=64, N=32, K=128, inject_index_oob=1),
78+
"test_index_full_oob": _create_test_script(M=64, N=32, K=128, inject_index_oob_all=1),
79+
"test_indirect_index_oob": _create_test_script(M=64, N=32, K=128, inject_offs_scale=2),
8880
}
8981
out_dir = os.path.dirname(__file__)
9082
for name, script in cases.items():
@@ -109,26 +101,23 @@ def test_normal_no_oob(self, dma_env, tmp_path):
109101
# 2. Partial index OOB
110102
def test_index_partial_oob(self, dma_env, tmp_path):
111103
"""Some indices shifted by +M, pointing beyond A rows."""
112-
script = _create_test_script(M=64, N=32, K=128,
113-
inject_index_oob=1)
104+
script = _create_test_script(M=64, N=32, K=128, inject_index_oob=1)
114105
result = run_kernel_script(script, tmp_path, 'TestGatherScatterOOB.test_index_partial_oob')
115106
assert result.detected, \
116107
f"Expected OOB detection with index_oob=1: {result}"
117108

118109
# 3. Full index OOB
119110
def test_index_full_oob(self, dma_env, tmp_path):
120111
"""All indices shifted to M+16, pointing well beyond A rows."""
121-
script = _create_test_script(M=64, N=32, K=128,
122-
inject_index_oob_all=1)
112+
script = _create_test_script(M=64, N=32, K=128, inject_index_oob_all=1)
123113
result = run_kernel_script(script, tmp_path, 'TestGatherScatterOOB.test_index_full_oob')
124114
assert result.detected, \
125115
f"Expected OOB detection with index_oob_all=1: {result}"
126-
116+
127117
# 4. indirect index OOB
128118
def test_indirect_index_oob(self, dma_env, tmp_path):
129119
"""All indices shifted to M+16, pointing well beyond A rows."""
130-
script = _create_test_script(M=64, N=32, K=128,
131-
inject_offs_scale=2)
120+
script = _create_test_script(M=64, N=32, K=128, inject_offs_scale=2)
132121
result = run_kernel_script(script, tmp_path, 'TestGatherScatterOOB.test_indirect_index_oob')
133122
assert result.detected, \
134123
f"Expected OOB detection with index_oob_all=1: {result}"

0 commit comments

Comments
 (0)