Skip to content

Instantly share code, notes, and snippets.

View Jokeren's full-sized avatar
💻
Work

Keren Zhou Jokeren

💻
Work
620 followers · 25 following
  • George Mason University
  • Fairfax
  • 22:56 (UTC -12:00)
View GitHub Profile
@Jokeren
Jokeren / benchmark.py
Created October 20, 2026 02:09
Grouped inter-warp scan: joint-routing follow-up benchmark against parent and main (60 stress cases, GB200)
import argparse, hashlib, json, math, os, shutil, statistics, subprocess, traceback
from pathlib import Path
import torch
import triton
from triton.experimental import gluon as g
from triton.experimental.gluon import language as l
from triton.compiler.compiler import CompiledKernel
from triton.testing import do_bench_cudagraph_proton
@g.jit
@Jokeren
Jokeren / README.md
Created October 19, 2026 07:07
Joint scan-prefix routing: GB200 benchmarks against PR 12317 and main

Reproducing the scan comparison

Revisions:

  • Parent: af8850e5ccc2c60b0feb975e349863695c17714b (PR #12317).
  • Main: 2122e24efc1886bfa36337c1f5b07c47d827b1d4.
  • Candidate: eadaa6387055c5510570a96f06d33c23205062f6 (keren/linear-layout-scan-joint-routing), based on the parent.

scan_benchmark.py contains the original 62 scan shapes, layouts, axes, directions, and combiners. targeted_benchmark.py contains five irregular linear layouts, four combiners, and both directions (40 cases). Both benchmark only scan kernels with their input/output loads and stores; the surrounding layout conversions from the original examples are omitted.

@Jokeren
Jokeren / code_compare.py
Created October 16, 2026 18:22
Triton scan PR 12184 with register totals: original 62 examples, slowest first, cf30b2770d versus main 4a18983590
from pathlib import Path
import struct
def machine_code(record):
path = next(v for k, v in record['metadata_group'].items() if k.endswith('.cubin'))
data = Path(path).read_bytes()
assert data[:6] == b'\x7fELF\x02\x01'
shoff = struct.unpack_from('<Q', data, 40)[0]
entsize, count, strings_index = struct.unpack_from('<HHH', data, 58)
sections = [struct.unpack_from('<IIQQQQIIQQ', data, shoff + i * entsize) for i in range(count)]
@Jokeren
Jokeren / code_compare.py
Last active October 16, 2026 17:56
Triton scan PR 12184: original 62 examples, slowest first, aa08440f75 versus main 4a18983590, GB200 GPU 2
from pathlib import Path
import struct
def machine_code(record):
path = next(v for k, v in record['metadata_group'].items() if k.endswith('.cubin'))
data = Path(path).read_bytes()
assert data[:6] == b'\x7fELF\x02\x01'
shoff = struct.unpack_from('<Q', data, 40)[0]
entsize, count, strings_index = struct.unpack_from('<HHH', data, 58)
sections = [struct.unpack_from('<IIQQQQIIQQ', data, shoff + i * entsize) for i in range(count)]
@Jokeren
Jokeren / opcodes.csv
Created October 9, 2026 15:50
Grouped scan carry scheduling: instruction counts and timings on GB200
We can make this file beautiful and searchable if this error is corrected: Unclosed quoted field in line 5.
"ID","Process ID","Process Name","Host Name","Kernel Name","Context","Stream","Block Size","Grid Size","Device","CC","sass__inst_executed_per_opcode"
"","","","","","","","","","","",""
"0","2673322","python3.12","127.0.0.1","kernel","1","7","(128, 1, 1)","(256, 1, 1)","0","10.0","1299456 (ISETP: 326656; LDS: 260096; SEL: 195584; FADD: 161792; FSEL: 96256; STS: 65536; STG: 65536; LDG: 65536; IMAD: 29696; VIADD: 9216; LOP3: 6144; LDC: 5120; S2R: 2048; LEA: 2048; UMOV: 1024; ULEA: 1024; SHF: 1024; SGXT: 1024; S2UR: 1024; LDCU: 1024; EXIT: 1024; BAR: 1024)"
"1","2673322","python3.12","127.0.0.1","kernel","1","7","(128, 1, 1)","(256, 1, 1)","0","10.0","3695616 (SEL: 787456; ISETP: 730112; IMAD: 630784; FSEL: 286720; LDS: 262144; FADD: 163840; VIADD: 157696; UISETP: 131072; UIADD3: 131072; BRA: 131072; STS: 65536; STG: 65536; LDG: 65536; LEA: 46080; UMOV: 17408; SHF: 7168; LOP3: 5120; LDCU: 3072; S2R: 2048; ULEA: 1024; SGXT: 1024; S2UR: 1024; LDC: 1024; EXIT: 1024; BAR: 1024)"
"2","2673322","python3.12","127.0.0.1
@Jokeren
Jokeren / comparison.md
Created October 9, 2026 15:19
Grouped inter-warp scan: full 74-case comparison on GB200

Grouped inter-warp scan experiment

  • Candidate: keren/linear-layout-scan-grouped-interwarp, 8d72f8bcad8e91c36654b41f42d7997f8e09dce5. Artifacts were compiled from the tree subsequently committed at this revision.
  • Parent: keren/linear-layout-scan-register-totals, 80f979012c9088ed029c6ec32d76d76b634a23cc (PR #12193).
  • Main reference: 332c72a78597ca2ccf9f71cf6e20bd4480143345. This is the saved main build used by the previous report; main was not rebuilt or refreshed for this experiment.
  • Physical GPU 2 (third NVIDIA GB200), 256 programs per launch.
  • Timer: do_bench_cudagraph_proton; eight measurements per distinct image, 20 ms per measurement. The launch order alternates between versions.
  • All 74 cases pass exact CPU reference checks. Kernels include input loads and output stores; preceding layout conversions are excluded.
  • Identical ELF text and register/spill/shared resources share measurements. Only cases 72 and 74 change relative to the parent; the other 72 kernels are byte-identical in
@Jokeren
Jokeren / full-comparison.md
Created October 9, 2026 03:52
Triton PR 12193 current policy: full 74-case scan benchmark against main

Full scan benchmark: current carry-ownership policy

Current PR #12193: 80f979012c9088ed029c6ec32d76d76b634a23cc. Previous unconditional-exchange version: c030b7da2b82aec34c31d82465a6bf9c669ea98d. Main reference: 332c72a78597ca2ccf9f71cf6e20bd4480143345.

Implemented policy

Additional register/lane exchanges are enabled for entirely warp-local scans or one-element thread-local segments. Cross-warp scans with multi-element thread-local segments preserve boundary-carry ownership. There is no fixed four-register-basis cap, and convertScanCarries is removed. The proposed shuffle-count threshold (t >= 2 for 32 lanes under the simplified model) has not been implemented as a separate cost model.

Coverage and method

@Jokeren
Jokeren / audit.json
Created October 9, 2026 03:15
Triton PR 12193: conditional register exchange preserving boundary carry owners
[
{
"name": "01_scan",
"shape": [
16,
32
],
"layout": "BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[1, 32], warps_per_cta=[16, 1], order=[1, 0], cga_layout=[[0, 0], [0, 0]])",
"axis": 1,
"kind": "add",
@Jokeren
Jokeren / comparison-public.json
Created October 9, 2026 02:55
Triton PR 12193: uncapped scan register redistribution, 74-case GB200 comparison
{
"metadata": {
"main": {
"commit": "332c72a78597ca2ccf9f71cf6e20bd4480143345",
"gpu": "NVIDIA GB200",
"visible_devices": "2",
"blocks": 256,
"timer": "do_bench_cudagraph_proton",
"samples": 3,
"rep": 20
@Jokeren
Jokeren / comparison-public.json
Created October 9, 2026 02:18
Triton scan totals register promotion: 74-case GB200 comparison and benchmark
{
"metadata": {
"main": {
"commit": "332c72a78597ca2ccf9f71cf6e20bd4480143345",
"gpu": "NVIDIA GB200",
"visible_devices": "2",
"blocks": 256,
"timer": "do_bench_cudagraph_proton",
"samples": 3,
"rep": 20