diff --git a/src/pyjuice/__init__.py b/src/pyjuice/__init__.py index 0a3f70fb..5668fb40 100644 --- a/src/pyjuice/__init__.py +++ b/src/pyjuice/__init__.py @@ -15,6 +15,9 @@ # TensorCircuit layers from pyjuice.layer import InputLayer, ProdLayer, SumLayer, ExternalParamsSumLayer +# Kernel launch-config autotuning (on by default; see `layer/kernels/autotune.py`) +from pyjuice.layer import set_autotune + # Construction methods from pyjuice.nodes import multiply, summate, inputs, set_block_size, structural_properties diff --git a/src/pyjuice/layer/__init__.py b/src/pyjuice/layer/__init__.py index 4d55e373..7d0278d4 100644 --- a/src/pyjuice/layer/__init__.py +++ b/src/pyjuice/layer/__init__.py @@ -5,4 +5,5 @@ from .external_sum_layer import ExternalParamsSumLayer, ExternalNodeInfo, StagedExternalParams, \ EXTERNAL_PARAMS_BUFFER_KWARG, EXTERNAL_PARAMS_GRAD_BUFFER_KWARG, \ EXTERNAL_PARAMS_KWARG, EXTERNAL_PARAMS_GRAD_KWARG -from .layer_group import LayerGroup \ No newline at end of file +from .layer_group import LayerGroup +from .kernels.autotune import set_autotune diff --git a/src/pyjuice/layer/kernels/autotune.py b/src/pyjuice/layer/kernels/autotune.py new file mode 100644 index 00000000..ebf60d2c --- /dev/null +++ b/src/pyjuice/layer/kernels/autotune.py @@ -0,0 +1,169 @@ +"""One-shot launch-config autotuning for the sum- / product-layer Triton kernels. + +Every Triton launch site in `sum_layer.py` / `prod_layer.py` picks its tile sizes from a budget +heuristic plus a handful of hand-measured constants ("cap the node tile below batch 64", "double +the edge tile in the LL regime", ...). Those constants were measured on ONE GPU and one set of +layer shapes, so they are a guess everywhere else. This module lets a site hand over a SHORT list +of candidate configs instead: the first call benchmarks them, the winner is cached, and every +later call is a dict lookup. + +Contract at every call site: + + * ``candidates[0]`` is the heuristic default. It is what gets used whenever tuning is off, is + impossible (CUDA-graph capture), or every candidate fails to launch -- so with tuning disabled + the behaviour is exactly what it was before the autotuner existed. + * All candidates must compute the same values up to floating-point reduction order. Only tile + sizes that partition the *output* (or the batch) are eligible; a tile size that sets a + reduction / max-stabilization group changes the result materially and must stay fixed. Each + call site carries a note saying which of its knobs is which. Even an output-only tile size is + not BIT-identical -- its shape changes how Triton lays the tile out and reduces it, measured at + ~1e-7, the same order as the atomic-add nondeterminism these kernels already have -- which is + why `pick` caches by shape rather than per layer (see there). + * ``bench(cfg)`` must not corrupt live buffers. That is automatic when the kernel's output is a + pure overwrite (re-running it recomputes the same values); a read-accumulate-write output must + be redirected to a scratch buffer (see `scratch_like`). + +Cost: the benchmark runs once per key and pays one Triton compile per candidate config (cached on +disk by Triton across processes), i.e. it lands in the first iteration's warmup. +""" + +import os +import torch + + +# Master switch, settable in code with `pyjuice.set_autotune(...)` or via PYJUICE_AUTOTUNE=0. When +# off, every site uses its heuristic default and nothing is ever benchmarked -- the behaviour from +# before this module existed. Worth turning off for A/B, for debugging, and for test suites, where +# many short-lived models would each pay a warmup they never amortize. +ENABLED = os.environ.get("PYJUICE_AUTOTUNE", "1") != "0" + + +# How much faster than the reference a candidate must measure before it is adopted. These kernels +# run in tens of microseconds, where event-timing noise is several percent even at a median of 7 +# reps: measured on a large HCLT, the {CUDA, Triton} element-flow comparisons cluster in 0.90-1.05 +# and land on either side from run to run, while the comparisons that genuinely favour CUDA sit at +# 1.4-1.6. A 10% margin cleanly separates the two, so a real win is still taken while a tie always +# resolves to the reference -- which matters because the arms of a {CUDA, Triton} comparison are +# numerically equivalent but not bit-identical, so a coin-flip there changes a run's output. +MARGIN = float(os.environ.get("PYJUICE_AUTOTUNE_MARGIN", 1.10)) + + +def _capturing() -> bool: + """True while a CUDA graph is being captured. Benchmarking synchronizes (illegal during + capture) and would bake the warmup launches into the graph, so tuning is skipped there.""" + try: + return torch.cuda.is_current_stream_capturing() + except Exception: + return False + + +def _median_time(run, warmup: int, reps: int): + """Median wall time of `run` in ms, or None if it cannot be launched (e.g. a tile config that + exceeds this GPU's shared memory raises `OutOfResources` at COMPILE time, before any write).""" + ev0, ev1 = torch.cuda.Event(enable_timing = True), torch.cuda.Event(enable_timing = True) + try: + for _ in range(warmup): + run() + torch.cuda.synchronize() + ts = [] + for _ in range(reps): + ev0.record(); run(); ev1.record(); torch.cuda.synchronize() + ts.append(ev0.elapsed_time(ev1)) + except Exception: + return None + ts.sort() + return ts[len(ts) // 2] + + +def best_of(candidates: list, warmup: int = 3, reps: int = 7): + """Benchmark each ``(key, run)`` candidate and return the winning key (None if none can run). + + ``candidates[0]`` is the REFERENCE -- the heuristic tile config, or the plain Triton kernel a + CUDA fast path is competing with -- and it wins unless some other candidate measures at least + `MARGIN` times faster. That tie-break is what makes the choice reproducible: several of these + comparisons sit within a percent of each other, and the arms of a {CUDA, Triton} comparison are + numerically equivalent but NOT bit-identical, so letting noise settle them makes a run's output + depend on how warm the GPU happened to be. `run` may write into scratch; only timing matters. + """ + ref_key, ref_run = candidates[0] + ref_t = _median_time(ref_run, warmup, reps) + + best_key, best_t = None, None + for key, run in candidates[1:]: + t = _median_time(run, warmup, reps) + if t is not None and (best_t is None or t < best_t): + best_key, best_t = key, t + + if ref_t is None: # the reference cannot run on this GPU + return best_key + return best_key if (best_t is not None and best_t * MARGIN < ref_t) else ref_key + + +# Process-wide cache of tuned configs, keyed by SHAPE -- see `pick`. +_CACHE = dict() + + +def _full_key(key): + return (torch.cuda.current_device(), key) + + +def set_autotune(enabled: bool = True, clear_cache: bool = False): + """Enable or disable launch-config autotuning process-wide; returns the previous setting. + + Choices already measured stay cached (and keep being used) unless `clear_cache` is set. + """ + global ENABLED + was, ENABLED = ENABLED, bool(enabled) + if clear_cache: + _CACHE.clear() + return was + + +def cached(key): + """The config already chosen for `key`, or None if this key still has to go through `pick`. + Lets a call site skip setting up for a benchmark (allocating a scratch output) on the steady + state path, where the answer is already known.""" + return _CACHE.get(_full_key(key)) + + +def pick(key, candidates: list, bench, warmup: int = 3, reps: int = 7): + """Return the best of `candidates` (config values), benchmarking them at most ONCE per `key`. + + `candidates[0]` is the heuristic default, kept unless `best_of`'s margin is cleared. Never + raises: a candidate that fails to launch is simply skipped. + + `key` must describe the SHAPE of the launch -- the kernel, the tile/block/edge/batch counts and + every constexpr flag -- and must NOT identify a particular layer object. Two knock-on reasons: + + * a config picked for one layer is equally good for any other launch of the same kernel at the + same shape, so keying on shape both cuts the tuning cost and gets a cache hit far more often; + * more importantly, these candidates are NOT bit-identical. They agree to ~1e-7 (changing the + tile shape changes how Triton lays out and reduces it), so two structurally identical models + that tuned independently would disagree in the last ulp -- and with the winner decided by + measurement, they sometimes would. Keying on shape makes them share one answer instead. + + Across processes the choice can still differ, exactly as the existing atomic-add + nondeterminism in these kernels already does, and at the same ~1e-7 magnitude. + """ + key = _full_key(key) + cfg = _CACHE.get(key) + if cfg is not None: + return cfg + + if not ENABLED or len(candidates) < 2 or _capturing(): + # Not cached on purpose: capture is transient, so a later ordinary call still tunes. + return candidates[0] + + best = best_of([(c, (lambda c = c: bench(c))) for c in candidates], warmup, reps) + cfg = candidates[0] if best is None else best + _CACHE[key] = cfg + return cfg + + +def scratch_like(tensor: torch.Tensor): + """A throwaway buffer to benchmark a read-accumulate-write kernel into, or None if it cannot + be allocated (in which case the caller must skip tuning rather than touch the live output).""" + try: + return torch.empty_like(tensor) + except torch.cuda.OutOfMemoryError: + return None diff --git a/src/pyjuice/layer/prod_layer.py b/src/pyjuice/layer/prod_layer.py index 4f638df8..fb801670 100644 --- a/src/pyjuice/layer/prod_layer.py +++ b/src/pyjuice/layer/prod_layer.py @@ -18,6 +18,7 @@ from pyjuice.nodes import ProdNodes from pyjuice.utils.parameter_list import FastParamList from .kernels import prod as kernels +from .kernels import autotune from .layer import Layer from .backend.node_partition import partition_nodes_by_n_edges from .backend.index_set import batched_index_set, batched_index_cum @@ -315,26 +316,57 @@ def _forward_backward(self, node_vals: torch.Tensor, element_vals: torch.Tensor, if batch_size < 64: BLOCK_M = min(BLOCK_M, _SMALL_BATCH_PROD_TILE_M) - grid = (triton.cdiv(n_nblocks * self.block_size, BLOCK_M), triton.cdiv(batch_size, BLOCK_B)) + def _launch_2d(cfg, out): + bm, bb = cfg + grid = (triton.cdiv(n_nblocks * self.block_size, bm), triton.cdiv(batch_size, bb)) + kernels._forward_backward_kernel_2d[grid]( + node_vals_ptr = out, + element_vals_ptr = element_vals, + local_ids_ptr = local_ids, + nids_ptr = nids, + cids_ptr = cids, + tot_n_nodes = tot_n_nodes, + tot_n_eles = tot_n_eles, + n_nblocks = n_nblocks, + num_edges = num_edges, + batch_size = batch_size, + BLOCK_M = bm, + BLOCK_B = bb, + block_size = block_size, + accum = accum, + partial_eval = partial_eval, + prop_logsumexp = prop_logsumexp + ) - kernels._forward_backward_kernel_2d[grid]( - node_vals_ptr = node_vals, - element_vals_ptr = element_vals, - local_ids_ptr = local_ids, - nids_ptr = nids, - cids_ptr = cids, - tot_n_nodes = tot_n_nodes, - tot_n_eles = tot_n_eles, - n_nblocks = n_nblocks, - num_edges = num_edges, - batch_size = batch_size, - BLOCK_M = BLOCK_M, - BLOCK_B = BLOCK_B, - block_size = block_size, - accum = accum, - partial_eval = partial_eval, - prop_logsumexp = prop_logsumexp - ) + # Both knobs are pure OUTPUT tiling -- each program owns a distinct (node, batch) slice + # and reduces over all `num_edges` on its own -- so the candidates differ only in + # reduction layout, and only the heuristic's two guesses are in question: how far to cap + # the node tile, and whether a fatter batch tile pays for the lower program count. + # `BLOCK_M` must divide `block_size` (the kernel derives the node block from `pid_m`), + # so it stays a power of two <= `block_size`. See `kernels/autotune.py`. + cfgs = [(BLOCK_M, BLOCK_B)] + budget_BLOCK_M = min(max(2048 // (BLOCK_B * num_edges), 1), self.block_size) + for bm in (8, 32, budget_BLOCK_M): + bm = min(bm, self.block_size) + if (bm, BLOCK_B) not in cfgs: + cfgs.append((bm, BLOCK_B)) + wide_BLOCK_B = min(BLOCK_B * 2, triton.next_power_of_2(batch_size)) + if (BLOCK_M, wide_BLOCK_B) not in cfgs: + cfgs.append((BLOCK_M, wide_BLOCK_B)) + + # `accum` makes the output read-accumulate-write, so the timing runs must go to a + # scratch buffer; without it the kernel just overwrites `node_vals` with the same values + # it is about to write anyway, so it can be timed in place. + key = (kernels._forward_backward_kernel_2d, n_nblocks, num_edges, block_size, + batch_size, accum, partial_eval, prop_logsumexp, cfgs[0]) + cfg = autotune.cached(key) + if cfg is None: + bench_out = node_vals if not accum else autotune.scratch_like(node_vals) + cfg = cfgs[0] if bench_out is None else \ + autotune.pick(key, cfgs, lambda c: _launch_2d(c, bench_out)) + del bench_out + + _launch_2d(cfg, node_vals) else: diff --git a/src/pyjuice/layer/sum_layer.py b/src/pyjuice/layer/sum_layer.py index e3a12293..d31a190a 100644 --- a/src/pyjuice/layer/sum_layer.py +++ b/src/pyjuice/layer/sum_layer.py @@ -25,6 +25,7 @@ class _TritonOutOfResources(Exception): from .kernels import sum_backward_element_sparse as bk_ele_sparse from .kernels import sum_backward_param_block_sparse as bk_par_bsparse from .kernels import sum_backward_param_sparse as bk_par_sparse +from .kernels import autotune from .layer import Layer from .backend.node_partition import partition_nodes_by_n_edges from .backend.index_set import batched_index_set, index_cum @@ -759,7 +760,71 @@ def _forward_block_sparse(self, node_mars: torch.Tensor, element_mars: torch.Ten else: use_bf16 = False - grid = (triton.cdiv(batch_size, BLOCK_B), triton.cdiv(layer_n_nodes, TILE_SIZE_M)) + # Which of the three block-sparse forward kernels runs is decided ONCE, from the heuristic + # config -- they differ numerically (different dot / reduction structure), so the autotuned + # candidates below must not be able to flip it. `use_bf16` is likewise already fixed above. + if TILE_SIZE_M >= 16 and TILE_SIZE_K >= 16 and BLOCK_B >= 16: + fw_kernel, tile_floor = fw_bsparse._fw_triton_block_sparse_tlmm_kernel, 16 + elif TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8: + fw_kernel, tile_floor = fw_bsparse._fw_triton_block_sparse_csmm1_kernel, 8 + else: + fw_kernel, tile_floor = fw_bsparse._fw_triton_block_sparse_csmm2_kernel, 1 + + def _launch_fw(cfg): + tm, bb = cfg + g = (triton.cdiv(batch_size, bb), triton.cdiv(layer_n_nodes, tm)) + for pid_m_start in range(0, g[1], 32768): + curr_grid = (g[0], min(pid_m_start + 32768, g[1]) - pid_m_start) + fw_kernel[curr_grid]( + node_mars, + element_mars, + params, + nids, + cids_start, + cids_increment, + pids_start, + pids_increment, + local_ids, + batch_size, + partial_eval = partial_eval, + BLOCK_B = bb, + TILE_SIZE_K = TILE_SIZE_K, + K_NUM_TILES = K_NUM_TILES, + TILE_SIZE_M = tm, + BLOCK_SIZE_M = BLOCK_SIZE_M, + use_bf16 = use_bf16, + propagation_alg_id = propagation_alg_id, + pflow_tempered_enabled = pflow_tempered_enabled, + pid_m_offset = pid_m_start, + **propagation_alg_kwargs, + **pflow_tempered_kwargs, + num_stages = 1 + ) + + # Autotune the two OUTPUT-TILING knobs around the heuristic: `TILE_SIZE_M` tiles the output + # nodes and `BLOCK_B` the batch, while the `emars_max` stabilizer is taken over `TILE_SIZE_K` + # edges (left untouched), so the candidates differ only in reduction layout (~1e-7) -- what + # really changes is the per-m-tile recomputation and the program count, which is exactly what + # the hand-picked tile caps above are guessing at. Candidates stay at or above `tile_floor` so + # the kernel choice cannot change. Benchmarked in place: the kernel overwrites `node_mars` + # with the values it is about to write for real anyway. The measurement itself runs only + # AFTER the CUDA fast paths have declined this layer (`autotune.pick` at the bottom) -- + # benchmarking here would waste warmup on layers that end up on CUDA and, worse, perturb the + # neighbouring {CUDA, Triton} measurement, whose two arms are numerically equivalent but not + # bit-identical, so nudging that tie shows up as a changed result. + fw_cfgs = [(TILE_SIZE_M, BLOCK_B)] + for tm in (TILE_SIZE_M // 2, TILE_SIZE_M * 2): + if tile_floor <= tm <= self.block_size and (tm, BLOCK_B) not in fw_cfgs: + fw_cfgs.append((tm, BLOCK_B)) + for bb in (BLOCK_B // 2, BLOCK_B * 2): + if tile_floor <= bb <= BATCH_SIZE_NP2 and (TILE_SIZE_M, bb) not in fw_cfgs: + fw_cfgs.append((TILE_SIZE_M, bb)) + + # The heuristic default is part of the key so that the `OutOfResources` retry below -- which + # re-enters with a smaller default -- looks up a fresh entry instead of the config that just + # failed (which would loop forever). + fw_key = (fw_kernel, self.block_size, TILE_SIZE_K, K_NUM_TILES, batch_size, num_nblocks, + partial_eval, use_bf16, propagation_alg_id, pflow_tempered_enabled, fw_cfgs[0]) # Optional CUDA (CuTe/TMA) fast path for the `tlmm` regime. It is numerically equivalent to # the Triton tlmm kernel and only valid here: LL propagation (`propagation_alg_id == 0`), the @@ -782,26 +847,13 @@ def _forward_block_sparse(self, node_mars: torch.Tensor, element_mars: torch.Ten if choice is None: # Autotune once: fastest of {valid CUDA tile configs} vs Triton. Every candidate # computes the same result into `node_mars`, so it stays correct afterwards. - cands = [(("cuda", c), - (lambda c=c: cuda_kernels.tlmm_forward_sum( - node_mars, element_mars, params, nids, ebase, pbase, - batch_size, self.block_size, K_NUM_TILES, c))) - for c in valid_cfgs] - - def _triton_tlmm_cand(): - for s in range(0, grid[1], 32768): - cg = (grid[0], min(s + 32768, grid[1]) - s) - fw_bsparse._fw_triton_block_sparse_tlmm_kernel[cg]( - node_mars, element_mars, params, nids, cids_start, cids_increment, - pids_start, pids_increment, local_ids, batch_size, - partial_eval = partial_eval, BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, BLOCK_SIZE_M = BLOCK_SIZE_M, - use_bf16 = use_bf16, propagation_alg_id = propagation_alg_id, - pflow_tempered_enabled = pflow_tempered_enabled, pid_m_offset = s, - **propagation_alg_kwargs, **pflow_tempered_kwargs, num_stages = 1) - cands.append((("triton", -1), _triton_tlmm_cand)) - choice = cuda_kernels.autotune(cands) or ("triton", -1) + cands = [(("triton", -1), (lambda: _launch_fw(fw_cfgs[0])))] + cands += [(("cuda", c), + (lambda c=c: cuda_kernels.tlmm_forward_sum( + node_mars, element_mars, params, nids, ebase, pbase, + batch_size, self.block_size, K_NUM_TILES, c))) + for c in valid_cfgs] + choice = autotune.best_of(cands) or ("triton", -1) self._cached_fw_cuda_choice[choice_key] = choice if choice[0] == "cuda": @@ -842,28 +894,12 @@ def _triton_tlmm_cand(): # Autotune the CUDA SPLIT configs against the Triton small-batch launch (csmm1/2, # whichever this tile shape selects -- mirrors the fall-through below). Every # candidate overwrites node_mars with the same result, so it stays correct. - cands = [(("cuda", c), (lambda c=c: cuda_kernels.smallbatch_forward_sum( - node_mars, element_mars, params, nids, sb_ebase, sb_pbase, - batch_size, self.block_size, num_edges, c))) - for c in range(n_sb_cfg)] - - def _triton_sb_cand(): - sb_kern = (fw_bsparse._fw_triton_block_sparse_csmm1_kernel - if (TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8) - else fw_bsparse._fw_triton_block_sparse_csmm2_kernel) - for s in range(0, grid[1], 32768): - cg = (grid[0], min(s + 32768, grid[1]) - s) - sb_kern[cg]( - node_mars, element_mars, params, nids, cids_start, cids_increment, - pids_start, pids_increment, local_ids, batch_size, - partial_eval = partial_eval, BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, BLOCK_SIZE_M = BLOCK_SIZE_M, - use_bf16 = use_bf16, propagation_alg_id = propagation_alg_id, - pflow_tempered_enabled = pflow_tempered_enabled, pid_m_offset = s, - **propagation_alg_kwargs, **pflow_tempered_kwargs, num_stages = 1) - cands.append((("triton", -1), _triton_sb_cand)) - choice = cuda_kernels.autotune(cands) or ("triton", -1) + cands = [(("triton", -1), (lambda: _launch_fw(fw_cfgs[0])))] + cands += [(("cuda", c), (lambda c=c: cuda_kernels.smallbatch_forward_sum( + node_mars, element_mars, params, nids, sb_ebase, sb_pbase, + batch_size, self.block_size, num_edges, c))) + for c in range(n_sb_cfg)] + choice = autotune.best_of(cands) or ("triton", -1) self._cached_fw_cuda_choice[choice_key] = choice if choice[0] == "cuda": @@ -876,92 +912,7 @@ def _triton_sb_cand(): # OOM-safe tuned launch: if the larger tuned tiles exceed this GPU's shared-memory/ # register budget, fall back to the default configuration (recompiled untuned). try: - for pid_m_start in range(0, grid[1], 32768): - pid_m_end = min(pid_m_start + 32768, grid[1]) - block_m_size = pid_m_end - pid_m_start - - curr_grid = (grid[0], block_m_size) - - if TILE_SIZE_M >= 16 and TILE_SIZE_K >= 16 and BLOCK_B >= 16: - fw_bsparse._fw_triton_block_sparse_tlmm_kernel[curr_grid]( - node_mars, - element_mars, - params, - nids, - cids_start, - cids_increment, - pids_start, - pids_increment, - local_ids, - batch_size, - partial_eval = partial_eval, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - use_bf16 = use_bf16, - propagation_alg_id = propagation_alg_id, - pflow_tempered_enabled = pflow_tempered_enabled, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs, - **pflow_tempered_kwargs, - num_stages = 1 - ) - - elif TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8: - fw_bsparse._fw_triton_block_sparse_csmm1_kernel[curr_grid]( - node_mars, - element_mars, - params, - nids, - cids_start, - cids_increment, - pids_start, - pids_increment, - local_ids, - batch_size, - partial_eval = partial_eval, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - use_bf16 = use_bf16, - propagation_alg_id = propagation_alg_id, - pflow_tempered_enabled = pflow_tempered_enabled, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs, - **pflow_tempered_kwargs, - num_stages = 1 - ) - - else: - fw_bsparse._fw_triton_block_sparse_csmm2_kernel[curr_grid]( - node_mars, - element_mars, - params, - nids, - cids_start, - cids_increment, - pids_start, - pids_increment, - local_ids, - batch_size, - partial_eval = partial_eval, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - use_bf16 = use_bf16, - propagation_alg_id = propagation_alg_id, - pflow_tempered_enabled = pflow_tempered_enabled, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs, - **pflow_tempered_kwargs, - num_stages = 1 - ) + _launch_fw(autotune.pick(fw_key, fw_cfgs, _launch_fw)) except _TritonOutOfResources: # `OutOfResources` is raised at compile time before any write, so retry is safe. if not (FORWARD_SUM_TUNED and not getattr(self, "_fw_tuning_oom", False) @@ -1547,6 +1498,89 @@ def _cumbase(start, incr): grid = (triton.cdiv(batch_size, BLOCK_B), triton.cdiv(layer_n_nodes, TILE_SIZE_M)) + # As in the forward, which of the four element-flow kernels runs (tempered or not, dot or + # csmm2) is decided ONCE from the heuristic config -- they differ numerically -- so the + # autotuned candidates below cannot flip it. `TL_DOT` is likewise already fixed above. + ele_tempered = not (abs(eflow_temperature - 1.0) < 1e-6) + if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8: + ele_kernel = (bk_ele_bsparse._bk_triton_block_sparse_tempered_ele_kernel if ele_tempered + else bk_ele_bsparse._bk_triton_block_sparse_ele_kernel) + ele_floor = 8 + else: + ele_kernel = (bk_ele_bsparse._bk_triton_block_sparse_tempered_ele_csmm2_kernel if ele_tempered + else bk_ele_bsparse._bk_triton_block_sparse_ele_csmm2_kernel) + ele_floor = 1 + # The two kernel families take different extra arguments. + ele_extra = dict(node_mars_tempered = kwargs["node_mars_tempered"], + eflow_temperature = eflow_temperature) if ele_tempered else \ + dict(node_mars = node_mars, allow_modify_flows = allow_modify_flows, + logspace_flows = logspace_flows, allow_neg_flows = allow_neg_flows, + propagation_alg_id = propagation_alg_id, **propagation_alg_kwargs) + + def _launch_ele(cfg, out): + tm, bb = cfg + g = (triton.cdiv(batch_size, bb), triton.cdiv(layer_n_nodes, tm)) + for pid_m_start in range(0, g[1], 32768): + curr_grid = (g[0], min(pid_m_start + 32768, g[1]) - pid_m_start) + ele_kernel[curr_grid]( + node_flows = node_flows, + element_flows = out, + element_mars = element_mars, + mparams = params, + chids = chids, + parids_start = parids_start, + parids_increment = parids_increment, + parpids_start = parpids_start, + parpids_increment = parpids_increment, + local_ids = local_ids, + batch_size = batch_size, + partial_eval = partial_eval, + ptr_inc_step = ptr_inc_step, + BLOCK_B = bb, + TILE_SIZE_K = TILE_SIZE_K, + K_NUM_TILES = K_NUM_TILES, + TILE_SIZE_M = tm, + BLOCK_SIZE_M = BLOCK_SIZE_M, + BLOCK_SIZE_K = BLOCK_SIZE_K, + TL_DOT = TL_DOT, + accumulate_ch_flows = accumulate_ch_flows, + pid_m_offset = pid_m_start, + num_stages = 1, + **ele_extra + ) + + # `TILE_SIZE_M` tiles the child-node outputs and `BLOCK_B` the batch; the parent reduction + # (and its stabilizer) runs over `TILE_SIZE_K`, which is left untouched -- so the candidates + # only trade per-tile work against program count. That is precisely the question + # `BACKWARD_ELE_FLOW_TUNED` (double the node tile) and `_SMALL_BATCH_ELE_TILE_M` answer with + # a constant above; here it is measured instead -- but only once the CUDA fast paths below + # have declined this layer (see `_tuned_ele_cfg`), for the reason spelled out in the + # forward: measuring here would perturb the neighbouring {CUDA, Triton} tie. + ele_cfgs = [(TILE_SIZE_M, BLOCK_B)] + for tm in (TILE_SIZE_M // 2, TILE_SIZE_M * 2): + if ele_floor <= tm <= cs_block_size and (tm, BLOCK_B) not in ele_cfgs: + ele_cfgs.append((tm, BLOCK_B)) + for bb in (BLOCK_B // 2, BLOCK_B * 2): + if ele_floor <= bb <= BATCH_SIZE_NP2 and (TILE_SIZE_M, bb) not in ele_cfgs: + ele_cfgs.append((TILE_SIZE_M, bb)) + + ele_key = (ele_kernel, self.block_size, cs_block_size, TILE_SIZE_K, K_NUM_TILES, + ptr_inc_step, batch_size, num_nblocks, partial_eval, TL_DOT, accumulate_ch_flows, + allow_modify_flows, logspace_flows, allow_neg_flows, propagation_alg_id, + ele_cfgs[0]) + + def _tuned_ele_cfg(): + cfg = autotune.cached(ele_key) + if cfg is not None: + return cfg + # `accumulate_ch_flows` makes `element_flows` read-accumulate-write, so the timing runs + # must go to a scratch buffer; otherwise the kernel overwrites it with the values it is + # about to write for real anyway and can be timed in place. + out = element_flows if not accumulate_ch_flows else autotune.scratch_like(element_flows) + if out is None: + return ele_cfgs[0] # no scratch -> leave this launch untuned + return autotune.pick(ele_key, ele_cfgs, lambda c: _launch_ele(c, out)) + # An external parameterization owns this computation. ONE interception for every regime: it is # handed the operands of all of them -- the per-tile tables the CuTe kernel wants, the general # `parids` walk a Triton kernel wants, and the shapes -- and picks among whichever of its own @@ -1599,21 +1633,7 @@ def _cuda_ele(tgt): batch_size, self.block_size, cs_block_size, K_NUM_TILES) def _triton_ele(tgt): - for s in range(0, grid[1], 32768): - cg = (grid[0], min(s + 32768, grid[1]) - s) - bk_ele_bsparse._bk_triton_block_sparse_ele_kernel[cg]( - node_flows = node_flows, element_flows = tgt, node_mars = node_mars, - element_mars = element_mars, mparams = params, chids = chids, - parids_start = parids_start, parids_increment = parids_increment, - parpids_start = parpids_start, parpids_increment = parpids_increment, - local_ids = local_ids, batch_size = batch_size, partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, BLOCK_B = BLOCK_B, TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, BLOCK_SIZE_K = BLOCK_SIZE_K, TL_DOT = TL_DOT, - num_stages = 1, propagation_alg_id = propagation_alg_id, - accumulate_ch_flows = accumulate_ch_flows, allow_neg_flows = allow_neg_flows, - pid_m_offset = s, **propagation_alg_kwargs) + _launch_ele(ele_cfgs[0], tgt) choice_key = (signature, batch_size) choice = self._cached_bk_ele_choice.get(choice_key) @@ -1624,9 +1644,9 @@ def _triton_ele(tgt): or self._bk_ele_scratch.dtype != element_flows.dtype): self._bk_ele_scratch = torch.empty_like(element_flows) scr = self._bk_ele_scratch - choice = cuda_kernels.autotune( - [(("cuda", 0), (lambda: _cuda_ele(scr))), - (("triton", -1), (lambda: _triton_ele(scr)))]) or ("triton", -1) + choice = autotune.best_of( + [(("triton", -1), (lambda: _triton_ele(scr))), + (("cuda", 0), (lambda: _cuda_ele(scr)))]) or ("triton", -1) self._cached_bk_ele_choice[choice_key] = choice if choice[0] == "cuda": _cuda_ele(element_flows) @@ -1671,22 +1691,10 @@ def _cuda_ele_sb(tgt, c): tgt, element_mars, node_flows, node_mars, params, chids, sb_ebase, sb_pbase, batch_size, self.block_size, cs_block_size, num_edges, c) + # Compare against the Triton launch that would actually run on fall-through (the + # tuned config and the kernel `ele_kernel` selects), not a hard-coded csmm2 one. def _triton_ele_sb(tgt): - for s in range(0, grid[1], 32768): - cg = (grid[0], min(s + 32768, grid[1]) - s) - bk_ele_bsparse._bk_triton_block_sparse_ele_csmm2_kernel[cg]( - node_flows = node_flows, element_flows = tgt, node_mars = node_mars, - element_mars = element_mars, mparams = params, chids = chids, - parids_start = parids_start, parids_increment = parids_increment, - parpids_start = parpids_start, parpids_increment = parpids_increment, - local_ids = local_ids, batch_size = batch_size, partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, BLOCK_B = BLOCK_B, TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, BLOCK_SIZE_K = BLOCK_SIZE_K, TL_DOT = TL_DOT, - num_stages = 1, propagation_alg_id = propagation_alg_id, - accumulate_ch_flows = accumulate_ch_flows, allow_neg_flows = allow_neg_flows, - pid_m_offset = s, **propagation_alg_kwargs) + _launch_ele(ele_cfgs[0], tgt) choice_key = (signature, batch_size, "sb") choice = self._cached_bk_ele_choice.get(choice_key) @@ -1697,146 +1705,16 @@ def _triton_ele_sb(tgt): or self._bk_ele_scratch.dtype != element_flows.dtype): self._bk_ele_scratch = torch.empty_like(element_flows) scr = self._bk_ele_scratch - cands = [(("cuda", c), (lambda c = c: _cuda_ele_sb(scr, c))) for c in range(n_sb_cfg)] - cands.append((("triton", -1), (lambda: _triton_ele_sb(scr)))) - choice = cuda_kernels.autotune(cands) or ("triton", -1) + cands = [(("triton", -1), (lambda: _triton_ele_sb(scr)))] + cands += [(("cuda", c), (lambda c = c: _cuda_ele_sb(scr, c))) for c in range(n_sb_cfg)] + choice = autotune.best_of(cands) or ("triton", -1) self._cached_bk_ele_choice[choice_key] = choice if choice[0] == "cuda": _cuda_ele_sb(element_flows, choice[1]) return None # choice == ("triton", -1): fall through to the Triton launch below - for pid_m_start in range(0, grid[1], 32768): - pid_m_end = min(pid_m_start + 32768, grid[1]) - block_m_size = pid_m_end - pid_m_start - - curr_grid = (grid[0], block_m_size) - - if abs(eflow_temperature - 1.0) < 1e-6: - - if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8: - bk_ele_bsparse._bk_triton_block_sparse_ele_kernel[curr_grid]( - node_flows = node_flows, - element_flows = element_flows, - node_mars = node_mars, - element_mars = element_mars, - mparams = params, - chids = chids, - parids_start = parids_start, - parids_increment = parids_increment, - parpids_start = parpids_start, - parpids_increment = parpids_increment, - local_ids = local_ids, - batch_size = batch_size, - partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, - allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_K = BLOCK_SIZE_K, - TL_DOT = TL_DOT, - num_stages = 1, - propagation_alg_id = propagation_alg_id, - accumulate_ch_flows = accumulate_ch_flows, - allow_neg_flows = allow_neg_flows, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs - ) - else: - bk_ele_bsparse._bk_triton_block_sparse_ele_csmm2_kernel[curr_grid]( - node_flows = node_flows, - element_flows = element_flows, - node_mars = node_mars, - element_mars = element_mars, - mparams = params, - chids = chids, - parids_start = parids_start, - parids_increment = parids_increment, - parpids_start = parpids_start, - parpids_increment = parpids_increment, - local_ids = local_ids, - batch_size = batch_size, - partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, - allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_K = BLOCK_SIZE_K, - TL_DOT = TL_DOT, - num_stages = 1, - propagation_alg_id = propagation_alg_id, - accumulate_ch_flows = accumulate_ch_flows, - allow_neg_flows = allow_neg_flows, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs - ) - - else: - - if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and BLOCK_B >= 8: - bk_ele_bsparse._bk_triton_block_sparse_tempered_ele_kernel[curr_grid]( - node_flows = node_flows, - element_flows = element_flows, - node_mars_tempered = kwargs["node_mars_tempered"], - element_mars = element_mars, - mparams = params, - chids = chids, - parids_start = parids_start, - parids_increment = parids_increment, - parpids_start = parpids_start, - parpids_increment = parpids_increment, - local_ids = local_ids, - batch_size = batch_size, - partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_K = BLOCK_SIZE_K, - TL_DOT = TL_DOT, - accumulate_ch_flows = accumulate_ch_flows, - pid_m_offset = pid_m_start, - eflow_temperature = eflow_temperature, - num_stages = 1, - ) - else: - bk_ele_bsparse._bk_triton_block_sparse_tempered_ele_csmm2_kernel[curr_grid]( - node_flows = node_flows, - element_flows = element_flows, - node_mars_tempered = kwargs["node_mars_tempered"], - element_mars = element_mars, - mparams = params, - chids = chids, - parids_start = parids_start, - parids_increment = parids_increment, - parpids_start = parpids_start, - parpids_increment = parpids_increment, - local_ids = local_ids, - batch_size = batch_size, - partial_eval = partial_eval, - ptr_inc_step = ptr_inc_step, - BLOCK_B = BLOCK_B, - TILE_SIZE_K = TILE_SIZE_K, - K_NUM_TILES = K_NUM_TILES, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_K = BLOCK_SIZE_K, - TL_DOT = TL_DOT, - accumulate_ch_flows = accumulate_ch_flows, - pid_m_offset = pid_m_start, - eflow_temperature = eflow_temperature, - num_stages = 1, - ) + _launch_ele(_tuned_ele_cfg(), element_flows) return None @@ -1966,21 +1844,124 @@ def _backward_block_sparse_par_flows(self, node_flows: torch.Tensor, params: tor # index them as `nblock * num_edges + edge`, so num_edges must equal the tensors' actual width. # Caching the slice -- rather than re-slicing every call -- avoids per-call copy kernels that # otherwise regress the backward at small batch. The cids-read runs on cache miss (warmup only). - if _BLOCK_SPARSE_EDGE_TRIM: - tkey = (id(cids), TILE_SIZE_K) + raw_num_edges, raw_cids, raw_pids, raw_pfids = num_edges, cids, pids, pfids + + def _par_edges(tk): + """`(num_edges, cids, pids, pfids)` trimmed for edge tile `tk`.""" + if not _BLOCK_SPARSE_EDGE_TRIM: + return raw_num_edges, raw_cids, raw_pids, raw_pfids + tkey = (id(raw_cids), tk) trimmed = self._cached_bk_par_trim.get(tkey) if trimmed is None: - real_max = int((cids != 0).any(dim = 0).sum()) - eff = triton.cdiv(real_max, TILE_SIZE_K) * TILE_SIZE_K - if 0 < eff < num_edges: - trimmed = (eff, cids[:, :eff].contiguous(), pids[:, :eff].contiguous(), pfids[:, :eff].contiguous()) + real_max = int((raw_cids != 0).any(dim = 0).sum()) + eff = triton.cdiv(real_max, tk) * tk + if 0 < eff < raw_num_edges: + trimmed = (eff, raw_cids[:, :eff].contiguous(), raw_pids[:, :eff].contiguous(), + raw_pfids[:, :eff].contiguous()) else: - trimmed = (num_edges, cids, pids, pfids) + trimmed = (raw_num_edges, raw_cids, raw_pids, raw_pfids) self._cached_bk_par_trim[tkey] = trimmed - num_edges, cids, pids, pfids = trimmed + return trimmed + + num_edges, cids, pids, pfids = _par_edges(TILE_SIZE_K) grid = (triton.cdiv(num_edges, TILE_SIZE_K), triton.cdiv(layer_n_nodes, TILE_SIZE_M)) + # As in the forward / element-flow backward, which of the four parameter-flow kernels runs + # (tempered or not, dot or csmm2) is decided ONCE from the heuristic config so the autotuned + # candidates cannot flip it. Within the non-tempered dot regime the read-add-store variant + # replaces the atomic one only when this partition's flow slots are provably disjoint. + par_tempered = not (abs(pflow_temperature - 1.0) < 1e-6) + if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and TILE_SIZE_B >= 8: + if par_tempered: + par_kernel = bk_par_bsparse._bk_triton_block_sparse_tempered_par_kernel + elif self._par_flow_collision_free(pfids): + par_kernel = bk_par_bsparse._bk_triton_block_sparse_par_kernel_rmw + else: + par_kernel = bk_par_bsparse._bk_triton_block_sparse_par_kernel + par_floor = 8 + else: + par_kernel = (bk_par_bsparse._bk_triton_block_sparse_tempered_par_csmm2_kernel if par_tempered + else bk_par_bsparse._bk_triton_block_sparse_par_csmm2_kernel) + par_floor = 1 + # The two kernel families take different extra arguments; neither tempered kernel has an + # `allow_neg_flows` parameter. (The launch this replaces passed one to the tempered csmm2 + # kernel, which Triton rejects outright -- that combination, tempered pflows with tiles + # below 8, could not have run.) + par_extra = dict(pflow_temperature = pflow_temperature) if par_tempered else \ + dict(allow_modify_flows = allow_modify_flows, logspace_flows = logspace_flows, + propagation_alg_id = propagation_alg_id, allow_neg_flows = allow_neg_flows, + **propagation_alg_kwargs) + + def _launch_par(cfg, out): + tk, warps = cfg + ne, cs, ps, fs = _par_edges(tk) + g = (triton.cdiv(ne, tk), triton.cdiv(layer_n_nodes, TILE_SIZE_M)) + for pid_m_start in range(0, g[1], 32768): + curr_grid = (g[0], min(pid_m_start + 32768, g[1]) - pid_m_start) + par_kernel[curr_grid]( + node_flows = node_flows, + element_mars = element_mars, + mparams = params, + param_flows = out, + nids = nids, + cids = cs, + pids = ps, + pfids = fs, + batch_size = batch_size, + num_edges = ne, + TILE_SIZE_B = TILE_SIZE_B, + B_NUM_TILES = B_NUM_TILES, + TILE_SIZE_K = tk, + TILE_SIZE_M = TILE_SIZE_M, + BLOCK_SIZE_M = self.block_size, + TL_DOT = TL_DOT, + negate_pflows = negate_pflows, + pid_m_offset = pid_m_start, + num_stages = 1, + **({"node_mars_tempered": kwargs["node_mars_tempered"]} if par_tempered + else {"node_mars": node_mars}), + **par_extra, + **({} if warps is None else {"num_warps": warps}) + ) + + # `TILE_SIZE_K` tiles the OUTPUT columns and `num_warps` is a launch-only knob, so the + # candidates agree to reduction order; `TILE_SIZE_M` (the node group the `log_n_fdm_max` + # stabilizer is taken over) and `TILE_SIZE_B` (the batch-reduction grouping) are the two that + # would change the arithmetic, and both are held at whatever the heuristic above chose. The + # candidates are exactly the alternatives `BACKWARD_PAR_FLOW_TUNED` / `_SMALL_BATCH_PAR_TILE_K` + # pick between with a hard-coded, GPU-specific rule: a wider edge tile (fewer redundant + # node_mars/node_flows reads) or a narrower one (more programs, better occupancy). + # `None` means "leave `num_warps` to Triton" -- the default the untuned launches use. + default_warps = par_kernel_extra.get("num_warps") + par_cfgs = [(TILE_SIZE_K, default_warps)] + for tk in (TILE_SIZE_K // 2, TILE_SIZE_K * 2): + if par_floor <= tk <= raw_num_edges and (TL_DOT == 0 or tk >= 16): + par_cfgs.append((tk, default_warps)) + par_cfgs.append((TILE_SIZE_K, 8 if default_warps is None else None)) + par_cfgs = list(dict.fromkeys(par_cfgs)) + + # Like the forward, the heuristic default is part of the key so the `OutOfResources` retry -- + # which re-enters with the untuned default -- cannot look up the config that just failed. + par_key = (par_kernel, self.block_size, raw_num_edges, num_nblocks, batch_size, + TILE_SIZE_M, TILE_SIZE_B, B_NUM_TILES, TL_DOT, allow_modify_flows, + logspace_flows, negate_pflows, allow_neg_flows, propagation_alg_id, par_cfgs[0]) + + def _tuned_par_cfg(): + cfg = autotune.cached(par_key) + if cfg is not None: + return cfg + # `param_flows` is read-accumulate-write, so the timing runs must go to a scratch clone. + # It is the full parameter array (can be GBs), so the scratch is local and freed right + # after; if it cannot be allocated, this launch is simply left untuned. + scr = autotune.scratch_like(param_flows) + if scr is None: + return par_cfgs[0] + try: + return autotune.pick(par_key, par_cfgs, lambda c: _launch_par(c, scr)) + finally: + del scr + # Optional CUDA fast path (CuTe/fp16/TMA), autotuned vs Triton INTO A SCRATCH buffer so the # autotune timing runs never corrupt the live param_flows (the kernel is read-accumulate-write). # Only intercepts when CUDA wins; Triton / unsupported shapes fall through to the dispatch below @@ -2032,18 +2013,7 @@ def _cuda_par(tgt): nbase, cbase, pbase, fbase, batch_size, self.block_size, num_edges, 0) def _triton_par(tgt): - for s in range(0, grid[1], 32768): - cg = (grid[0], min(s + 32768, grid[1]) - s) - bk_par_bsparse._bk_triton_block_sparse_par_kernel_rmw[cg]( - node_flows = node_flows, node_mars = node_mars, element_mars = element_mars, - mparams = params, param_flows = tgt, nids = nids, cids = cids, pids = pids, - pfids = pfids, batch_size = batch_size, num_edges = num_edges, - allow_modify_flows = allow_modify_flows, logspace_flows = logspace_flows, - TILE_SIZE_B = TILE_SIZE_B, B_NUM_TILES = B_NUM_TILES, TILE_SIZE_K = TILE_SIZE_K, - TILE_SIZE_M = TILE_SIZE_M, BLOCK_SIZE_M = self.block_size, TL_DOT = TL_DOT, - propagation_alg_id = propagation_alg_id, negate_pflows = negate_pflows, - allow_neg_flows = allow_neg_flows, pid_m_offset = s, - **propagation_alg_kwargs, **par_kernel_extra, num_stages = 1) + _launch_par(par_cfgs[0], tgt) choice_key = (par_sig, batch_size) choice = self._cached_bk_par_choice.get(choice_key) @@ -2054,8 +2024,8 @@ def _triton_par(tgt): # scratch can't be allocated (memory-constrained GPU), fall back to Triton. try: scr = torch.empty_like(param_flows) - choice = cuda_kernels.autotune( - [("cuda", (lambda: _cuda_par(scr))), ("triton", (lambda: _triton_par(scr)))]) or "triton" + choice = autotune.best_of( + [("triton", (lambda: _triton_par(scr))), ("cuda", (lambda: _cuda_par(scr)))]) or "triton" del scr except torch.cuda.OutOfMemoryError: choice = "triton" @@ -2089,150 +2059,23 @@ def _triton_par(tgt): f"no Triton fallback for this parameterization." ) - for pid_m_start in range(0, grid[1], 32768): - pid_m_end = min(pid_m_start + 32768, grid[1]) - block_m_size = pid_m_end - pid_m_start - - curr_grid = (grid[0], block_m_size) - - if abs(pflow_temperature - 1.0) < 1e-6: - - if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and TILE_SIZE_B >= 8: - # Use the non-atomic read-add-store variant when this partition's param-flow - # slots are provably collision-free (untied); otherwise the atomic kernel. - # The check is computed once per partition and cached (see the helper). - if self._par_flow_collision_free(pfids): - par_kernel = bk_par_bsparse._bk_triton_block_sparse_par_kernel_rmw - else: - par_kernel = bk_par_bsparse._bk_triton_block_sparse_par_kernel - try: - par_kernel[curr_grid]( - node_flows = node_flows, - node_mars = node_mars, - element_mars = element_mars, - mparams = params, - param_flows = param_flows, - nids = nids, - cids = cids, - pids = pids, - pfids = pfids, - batch_size = batch_size, - num_edges = num_edges, - allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, - TILE_SIZE_B = TILE_SIZE_B, - B_NUM_TILES = B_NUM_TILES, - TILE_SIZE_K = TILE_SIZE_K, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = self.block_size, - TL_DOT = TL_DOT, - propagation_alg_id = propagation_alg_id, - negate_pflows = negate_pflows, - allow_neg_flows = allow_neg_flows, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs, - **par_kernel_extra, - num_stages = 1 - ) - except _TritonOutOfResources: - # The tuned launch config exceeds this GPU's shared memory. Disable the - # tuning for this layer (cached) and retry with the default heuristic. - # `OutOfResources` is raised at compile time, before any `param_flows` - # write, so re-running from scratch is safe (no partial accumulation). - if "num_warps" not in par_kernel_extra: - raise - self._par_tuning_oom = True - warnings.warn("pyjuice: tuned parameter-flow backward launch exceeds GPU " - "shared memory; falling back to the default configuration.") - return self._backward_block_sparse_par_flows( - node_flows, params, node_mars, element_mars, param_flows, - nids, cids, pids, pfids, allow_modify_flows = allow_modify_flows, - propagation_alg = propagation_alg, logspace_flows = logspace_flows, - negate_pflows = negate_pflows, allow_neg_flows = allow_neg_flows, - force_use_fp32 = force_use_fp32, - pflow_temperature = pflow_temperature, **kwargs) - - else: - bk_par_bsparse._bk_triton_block_sparse_par_csmm2_kernel[curr_grid]( - node_flows = node_flows, - node_mars = node_mars, - element_mars = element_mars, - mparams = params, - param_flows = param_flows, - nids = nids, - cids = cids, - pids = pids, - pfids = pfids, - batch_size = batch_size, - num_edges = num_edges, - allow_modify_flows = allow_modify_flows, - logspace_flows = logspace_flows, - TILE_SIZE_B = TILE_SIZE_B, - B_NUM_TILES = B_NUM_TILES, - TILE_SIZE_K = TILE_SIZE_K, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = self.block_size, - TL_DOT = TL_DOT, - propagation_alg_id = propagation_alg_id, - negate_pflows = negate_pflows, - allow_neg_flows = allow_neg_flows, - pid_m_offset = pid_m_start, - **propagation_alg_kwargs, - num_stages = 1 - ) - - else: - - if TILE_SIZE_M >= 8 and TILE_SIZE_K >= 8 and TILE_SIZE_B >= 8: - bk_par_bsparse._bk_triton_block_sparse_tempered_par_kernel[curr_grid]( - node_flows = node_flows, - node_mars_tempered = kwargs["node_mars_tempered"], - element_mars = element_mars, - mparams = params, - param_flows = param_flows, - nids = nids, - cids = cids, - pids = pids, - pfids = pfids, - batch_size = batch_size, - num_edges = num_edges, - TILE_SIZE_B = TILE_SIZE_B, - B_NUM_TILES = B_NUM_TILES, - TILE_SIZE_K = TILE_SIZE_K, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = self.block_size, - TL_DOT = TL_DOT, - negate_pflows = negate_pflows, - pid_m_offset = pid_m_start, - pflow_temperature = pflow_temperature, - num_stages = 1 - ) - - else: - bk_par_bsparse._bk_triton_block_sparse_tempered_par_csmm2_kernel[curr_grid]( - node_flows = node_flows, - node_mars_tempered = kwargs["node_mars_tempered"], - element_mars = element_mars, - mparams = params, - param_flows = param_flows, - nids = nids, - cids = cids, - pids = pids, - pfids = pfids, - batch_size = batch_size, - num_edges = num_edges, - TILE_SIZE_B = TILE_SIZE_B, - B_NUM_TILES = B_NUM_TILES, - TILE_SIZE_K = TILE_SIZE_K, - TILE_SIZE_M = TILE_SIZE_M, - BLOCK_SIZE_M = self.block_size, - TL_DOT = TL_DOT, - negate_pflows = negate_pflows, - allow_neg_flows = allow_neg_flows, - pid_m_offset = pid_m_start, - pflow_temperature = pflow_temperature, - num_stages = 1 - ) + # OOM-safe tuned launch: `OutOfResources` is raised at COMPILE time, before any write to + # `param_flows`, so retrying from scratch cannot double-accumulate. + try: + _launch_par(_tuned_par_cfg(), param_flows) + except _TritonOutOfResources: + if par_cfgs[0][1] is None: + raise + self._par_tuning_oom = True + warnings.warn("pyjuice: tuned parameter-flow backward launch exceeds GPU " + "shared memory; falling back to the default configuration.") + return self._backward_block_sparse_par_flows( + node_flows, params, node_mars, element_mars, param_flows, + nids, raw_cids, raw_pids, raw_pfids, allow_modify_flows = allow_modify_flows, + propagation_alg = propagation_alg, logspace_flows = logspace_flows, + negate_pflows = negate_pflows, allow_neg_flows = allow_neg_flows, + force_use_fp32 = force_use_fp32, partition_id = partition_id, + pflow_temperature = pflow_temperature, **kwargs) return None @@ -2631,9 +2474,9 @@ def _triton_par_sb(tgt): or self._bk_par_scratch.dtype != param_flows.dtype): self._bk_par_scratch = torch.empty_like(param_flows) scr = self._bk_par_scratch - cands = [(("cuda", cfg), (lambda cfg = cfg: _cuda_par_sb(scr, cfg))) for cfg in range(n_cfg)] - cands.append((("triton", -1), (lambda: _triton_par_sb(scr)))) - choice = cuda_kernels.autotune(cands) or ("triton", -1) + cands = [(("triton", -1), (lambda: _triton_par_sb(scr)))] + cands += [(("cuda", cfg), (lambda cfg = cfg: _cuda_par_sb(scr, cfg))) for cfg in range(n_cfg)] + choice = autotune.best_of(cands) or ("triton", -1) self._cached_bk_par_sparse_choice[choice_key] = choice if choice[0] == "cuda": _cuda_par_sb(param_flows, choice[1]) diff --git a/src/pyjuice/nodes/external_params/block_scale.py b/src/pyjuice/nodes/external_params/block_scale.py index a5f1e057..ef7f7f2d 100644 --- a/src/pyjuice/nodes/external_params/block_scale.py +++ b/src/pyjuice/nodes/external_params/block_scale.py @@ -353,16 +353,18 @@ def _build_plan(self, layer, ns_tensors, node_mars, element_mars, params, extern """ Resolve every per-layer launch argument once, and check each kernel's assumptions. - TWO kernels serve this type, and which applies is a property of the shape: + THREE kernels serve this type, and which apply is a property of the shape: * the CuTe/TMA fork, for batches it can tile (`batch % 64 == 0`) on sm_90+ with CUTLASS. It carries the normalizer as a gate-factored contraction against a precomputed `sigma`; * a plain-CUDA small-batch kernel, one warp per 32 nodes, which accumulates the normalizer - inline and needs neither `batch % 64` nor `num_edges % 64` nor CUTLASS. + inline and needs neither `batch % 64` nor `num_edges % 64` nor CUTLASS; + * the portable Triton fork, which needs no CUDA toolchain at all and applies everywhere, + including the shapes the other two decline (a gate finer than the node block, in + particular, is served by it alone). - Both are collected here and the choice is MEASURED, so where they overlap the faster one wins - rather than the one that happened to be checked first. Where neither applies this raises: there - is no Triton fallback for this parameterization. + All three are collected here and the choice is MEASURED, so where they overlap the faster one + wins rather than the one that happened to be checked first. """ from .kernels.c import get_cute_module, get_sb_module import pyjuice.layer.kernels.c as ck @@ -1063,22 +1065,24 @@ def pre_backward_layer(self, layer, ns_tensors, node_flows, element_flows, node_ state = self._bw_state(layer, node_mars, kwargs) - from .kernels.c import get_module, get_ele_bw_module, get_par_bw_module, get_sb_bw_module - - plain, ele_mod, par_mod = get_module(), get_ele_bw_module(), get_par_bw_module() + from .kernels.c import get_ele_bw_module, get_par_bw_module, get_sb_bw_module + from .kernels.shift_logz import shift_logz + + # NO CUDA extension is required here. The CuTe and small-batch forks are accelerators, and + # every kernel on this path also has a Triton fork: a gate table is always built, so + # `"triton"` is always among the element candidates; `_par_triton_hook` is installed + # unconditionally; and the normalizer shift -- the last piece that used to pin the whole + # gated backward to nvcc -- is now `shift_logz`, which falls back to Triton on its own. + # Every use of `ele_mod` / `sb_mod` / `par_mod` below checks for itself, so a missing one + # costs speed, not correctness. + ele_mod, par_mod = get_ele_bw_module(), get_par_bw_module() sb_mod = get_sb_bw_module() - if plain is None or (ele_mod is None and sb_mod is None): - raise NotImplementedError( - "the block-scale backward needs its CUDA extensions -- the plain one holding the " - "normalizer shift, plus at least one of the CuTe/TMA forks (large batch) and the " - "small-batch forks -- and they are unavailable here. There is no Triton fallback." - ) batch_size, block_size = state["batch_size"], state["block_size"] ext_base, gate_cbs, node_cbs = state["ext_base"], state["gate_cbs"], state["node_cbs"] for nids, log_z, rows in state["shift_args"]: - plain.lowrank_shift_logz(node_mars, nids, log_z, block_size, 1.0) + shift_logz(node_mars, nids, log_z, block_size, 1.0) cache = getattr(layer, "_bs_bw_gate_cache", None) if cache is None: @@ -1461,7 +1465,7 @@ def post_backward_layer(self, layer, ns_tensors, ns_grad_tensors, node_flows, el layer._ext_bw_par_sb_hook = None layer._ext_bw_par_triton_hook = None - from .kernels.c import get_module + from .kernels.shift_logz import shift_logz state = layer._bs_bw_state @@ -1516,7 +1520,7 @@ def _go(gt, bb, tgt, nids = nids, log_z = log_z, rows = rows, sigma = sigma, layer._bs_grad_ext = None for nids, log_z, rows in state["shift_args"]: - get_module().lowrank_shift_logz(node_mars, nids, log_z, state["block_size"], -1.0) + shift_logz(node_mars, nids, log_z, state["block_size"], -1.0) return None diff --git a/src/pyjuice/nodes/external_params/kernels/shift_logz.py b/src/pyjuice/nodes/external_params/kernels/shift_logz.py new file mode 100644 index 00000000..97a30e02 --- /dev/null +++ b/src/pyjuice/nodes/external_params/kernels/shift_logz.py @@ -0,0 +1,77 @@ +""" +The +-log-Z normalizer shift, in Triton. + +A port of `lowrank_shift_logz` from `c/lowrank_backward.cu`, with the same contract: + + node_mars[nids[r] + m, b] += sign * log_z[r, m, b] + +`BlockScaleSumParams` stores `log N - log Z` in `node_mars` during the forward, and its backward needs +`log N`; it shifts by `+log Z` on the way in and by `-log Z` on the way out. That made a CUDA toolchain +a hard requirement for the gated backward -- every OTHER kernel on that path already has a Triton fork, +so this one elementwise add was the whole of what an nvcc-less machine was missing. + +Bit-identical to the CUDA kernel: `sign` is exactly +-1, so `sign * log_z` is a negation or a no-op and +the remaining add rounds once either way. +""" + +import triton +import triton.language as tl + +from pyjuice.utils.kernel_launcher import triton_jit + + +@triton_jit +def _shift_logz_kernel(node_mars, nids, log_z, batch_size, num_rows, + block_size: tl.constexpr, sign, BLOCK_M: tl.constexpr, BLOCK_B: tl.constexpr): + # `offs_m` walks the [rows x block_size] node rows, which is exactly `log_z`'s leading dimension. + offs_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + offs_b = tl.program_id(1) * BLOCK_B + tl.arange(0, BLOCK_B) + mask_m = offs_m < num_rows * block_size + mask = mask_m[:,None] & (offs_b < batch_size)[None,:] + + # `block_size` is a constexpr power of two, so these are a shift and a mask. + row = offs_m // block_size + m = offs_m % block_size + + # int64 for the same reason the staging transpose needs it: `num_rows * block_size * batch_size` + # is a whole activation buffer and passes 2^31 on a large model at a large batch. + ob = offs_b.to(tl.int64) + lz = tl.load(log_z + offs_m.to(tl.int64)[:,None] * batch_size + ob[None,:], mask = mask, other = 0.0) + + nid = tl.load(nids + row, mask = mask_m, other = 0) + ptr = node_mars + (nid + m)[:,None] * batch_size + ob[None,:] + # Each (node, sample) belongs to exactly one program, so the read-modify-write cannot race. + tl.store(ptr, tl.load(ptr, mask = mask, other = 0.0) + sign * lz, mask = mask) + + +def shift_logz_triton(node_mars, nids, log_z, block_size: int, sign: float) -> None: + """`node_mars[nids[r] + m, b] += sign * log_z[r, m, b]`, matching `lowrank_shift_logz` from the + CUDA extension argument for argument so the two are drop-in for each other.""" + batch_size = node_mars.size(1) + num_rows = nids.size(0) + + assert log_z.is_contiguous(), "shift_logz: `log_z` must be contiguous" + assert log_z.numel() == num_rows * block_size * batch_size, \ + "shift_logz: `log_z` must hold num_rows * block_size * batch_size entries" + + if num_rows == 0 or block_size == 0 or batch_size == 0: + return None + + BLOCK_B = min(triton.next_power_of_2(batch_size), 128) + BLOCK_M = max(1024 // BLOCK_B, 1) + + _shift_logz_kernel[(triton.cdiv(num_rows * block_size, BLOCK_M), triton.cdiv(batch_size, BLOCK_B))]( + node_mars = node_mars, nids = nids, log_z = log_z, batch_size = batch_size, + num_rows = num_rows, block_size = block_size, sign = sign, + BLOCK_M = BLOCK_M, BLOCK_B = BLOCK_B) + + +def shift_logz(node_mars, nids, log_z, block_size: int, sign: float) -> None: + """The same shift, taking the CUDA extension when it is built and the Triton port otherwise.""" + from .c import get_module + + mod = get_module() + if mod is not None: + mod.lowrank_shift_logz(node_mars, nids, log_z, block_size, sign) + else: + shift_logz_triton(node_mars, nids, log_z, block_size, sign) diff --git a/tests/conftest.py b/tests/conftest.py index ce5a0f3f..1c7f17ee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -67,6 +67,13 @@ def pytest_configure(config): import torch torch.set_num_threads(8) + # Kernel launch-config autotuning off for the suite. Each test builds a short-lived model, so it + # would pay the one-off benchmark (several Triton compiles per launch signature) and never + # amortize it; and the tuned configs are only equal up to floating-point reduction order, which + # tests comparing two independently-built models against each other assert away. Tests that + # exercise the tuner itself turn it back on locally. + os.environ.setdefault("PYJUICE_AUTOTUNE", "0") + def pytest_collection_modifyitems(session, config, items): # LPT scheduling: reorder the collected tests slowest-first using the cached durations, so the diff --git a/tests/layer/par_flow_backward_test.py b/tests/layer/par_flow_backward_test.py index d1b591e4..9636a5f4 100644 --- a/tests/layer/par_flow_backward_test.py +++ b/tests/layer/par_flow_backward_test.py @@ -11,6 +11,7 @@ import pyjuice.nodes.distributions as dists import pyjuice.layer.kernels.sum_backward_param_block_sparse as parmod from pyjuice.layer.sum_layer import SumLayer +from pyjuice.layer.kernels import autotune def _build_small_hclt(num_latents=64, num_cats=8, num_vars=16, device="cuda:0"): @@ -65,9 +66,13 @@ def test_par_flow_rmw_matches_atomic(): assert torch.all((pf_rmw - pf_atomic).abs() < 1e-3) -def test_par_flow_oom_fallback(): - # If the tuned launch raises OutOfResources (simulating a smaller GPU), the backward - # must transparently fall back to the default configuration and stay correct. +@pytest.mark.parametrize("tuning", [True, False]) +def test_par_flow_oom_fallback(tuning): + # If the tuned launch raises OutOfResources (simulating a smaller GPU), the backward must + # transparently avoid that configuration and stay correct. Two mechanisms cover this, and both + # are exercised here: with the autotuner ON, the offending config simply loses the benchmark + # (`OutOfResources` is raised at compile time, so the candidate is skipped); with it OFF, the + # launch itself raises and `_par_tuning_oom` triggers a retry on the untuned config. from triton.runtime.errors import OutOfResources pc, device = _build_small_hclt() @@ -90,23 +95,29 @@ def run(*a, **k): if l.is_sum() and hasattr(l, "_par_tuning_oom"): del l._par_tuning_oom + was_enabled, cache = autotune.ENABLED, dict(autotune._CACHE) try: + # A fresh cache, so the proxy is actually benchmarked rather than served a prior choice. + autotune.ENABLED, autotune._CACHE = tuning, dict() parmod._bk_triton_block_sparse_par_kernel_rmw = OOMProxy() with warnings.catch_warnings(): warnings.simplefilter("ignore") - got = _backward(pc, x) # must fall back, not crash - got2 = _backward(pc, x) # cached fallback path + got = _backward(pc, x) # must avoid the failing config, not crash + got2 = _backward(pc, x) # cached choice / cached fallback path finally: parmod._bk_triton_block_sparse_par_kernel_rmw = real_rmw + autotune.ENABLED, autotune._CACHE = was_enabled, cache assert torch.all((got - ref).abs() < 1e-3) assert torch.all((got2 - ref).abs() < 1e-3) - n_oom = sum(1 for g in pc.inner_layer_groups for l in g - if l.is_sum() and getattr(l, "_par_tuning_oom", False)) - assert n_oom >= 1 + if not tuning: + n_oom = sum(1 for g in pc.inner_layer_groups for l in g + if l.is_sum() and getattr(l, "_par_tuning_oom", False)) + assert n_oom >= 1 if __name__ == "__main__": test_par_flow_collision_free_gate() test_par_flow_rmw_matches_atomic() - test_par_flow_oom_fallback() + test_par_flow_oom_fallback(True) + test_par_flow_oom_fallback(False) diff --git a/tests/model/external_params/external_blockscale_sweep_test.py b/tests/model/external_params/external_blockscale_sweep_test.py index 532f3b48..48e580d1 100644 --- a/tests/model/external_params/external_blockscale_sweep_test.py +++ b/tests/model/external_params/external_blockscale_sweep_test.py @@ -43,7 +43,8 @@ def _cute_available(): needs_cute = pytest.mark.skipif( not _cute_available(), - reason = "needs the CuTe/TMA extension (nvcc + CUTLASS + sm_90+); no fallback exists") + reason = "compares the CuTe/TMA fork (nvcc + CUTLASS + sm_90+) against the oracle; the " + "Triton fork that covers a machine without it is pinned separately") def _build(num_latents, block_size, gate_cbs, seed): diff --git a/tests/model/external_params/external_blockscale_test.py b/tests/model/external_params/external_blockscale_test.py index 4e7c7c6a..ed4eafeb 100644 --- a/tests/model/external_params/external_blockscale_test.py +++ b/tests/model/external_params/external_blockscale_test.py @@ -19,9 +19,12 @@ * every tile config computes the same contraction, so they must all agree -- which is what makes the launcher's autotuning safe. -There is no Triton fallback for this parameterization: it is a fork of the CuTe/TMA sum kernel, so the -tests skip wholesale where that kernel does not apply, and the last test pins the boundary of what is -supported so it is visible when it moves. +A portable Triton fork now covers every shape, so the parameterization itself needs no CUDA toolchain +apart from the normalizer shift (`lowrank_shift_logz`, the one kernel here with no Triton port). These +tests still skip wholesale without the CuTe/TMA extension, because what they check is that all the +forks agree with the oracle -- and the CuTe one is the fork most of them exercise. The Triton fork is +pinned separately by `test_shapes_only_the_triton_fork_reaches` and +`test_backward_needs_only_the_normalizer_shift_extension`. """ import pytest @@ -50,7 +53,8 @@ def _cute_available(): needs_cute = pytest.mark.skipif( not _cute_available(), - reason = "needs the CuTe/TMA extension (nvcc + CUTLASS + sm_90+); no fallback exists") + reason = "compares the CuTe/TMA fork (nvcc + CUTLASS + sm_90+) against the oracle; the " + "Triton fork that covers a machine without it is pinned separately") # --------------------------------------------------------------------------------- circuits @@ -736,6 +740,82 @@ def test_backward_matches_reference(num_latents, block_size, gate_cbs, batch, sc assert d_pf < 3e-3, f"param flows off by {d_pf} (relative)" +@cuda_only +@needs_cute +def test_the_whole_path_runs_without_any_cuda_extension(): + """Forward AND backward with EVERY CUDA extension gone -- the state of a machine with no nvcc. + + Every kernel this parameterization needs has a Triton fork: the forward, the element flows, the + parameter flows, and (since it was ported) the +-log-Z normalizer shift. So the CUDA extensions + are accelerators, and a toolchain-less machine must get the same answers, just slower. Two + things used to break that -- a guard demanding at least one CUDA flow fork, and the shift being + CUDA-only -- and this pins both. + """ + import pyjuice.nodes.external_params.kernels.c as ck + + num_latents, block_size, gate_cbs, batch = 256, 128, 8, 128 + accessors = ["get_module", "get_ele_bw_module", "get_sb_bw_module", "get_par_bw_module", + "get_cute_module", "get_sb_module"] + + pc, root, ns, data, phi, lls = _run(num_latents, block_size, gate_cbs, batch, scale = 1.5) + pc.backward(data, flows_memory = 0.0) + ns.update_param_flows(pc.param_flows) + with_cuda_lls, with_cuda_pf = lls.detach().double(), ns.get_param_flows().double() + + saved = {n: getattr(ck, n) for n in accessors} + try: + for n in accessors: + setattr(ck, n, lambda: None) + # A fresh circuit: the forks are chosen (and cached) per layer at the first pass. + pc2, _, ns2, data2, phi2, lls2 = _run(num_latents, block_size, gate_cbs, batch, scale = 1.5) + pc2.backward(data2, flows_memory = 0.0) + ns2.update_param_flows(pc2.param_flows) + triton_lls, triton_pf = lls2.detach().double(), ns2.get_param_flows().double() + finally: + for n, f in saved.items(): + setattr(ck, n, f) + + ref_ef, ref_pf = _flow_reference(pc, ns, phi, gate_cbs, batch) + d_pf = float(((triton_pf.to(ref_pf.device) - ref_pf).abs() / ref_pf.clamp(min = 1e-30)).max()) + assert d_pf < 3e-3, f"Triton-only param flows off by {d_pf} (relative)" + + # And they agree with what the CUDA forks produced, to those kernels' fp16/bf16 floor. + d_lls = float((triton_lls - with_cuda_lls).abs().max()) + d = float(((triton_pf - with_cuda_pf).abs() / with_cuda_pf.abs().clamp(min = 1e-30)).max()) + assert d_lls < 2e-3, f"Triton-only and CUDA-served lls differ by {d_lls}" + assert d < 3e-3, f"Triton-only and CUDA-served param flows differ by {d} (relative)" + + +@cuda_only +@pytest.mark.parametrize("rows,block_size,batch", [(4, 128, 128), (7, 64, 1), (3, 256, 512)]) +def test_the_logz_shift_port_matches_the_cuda_kernel(rows, block_size, batch): + """`shift_logz_triton` against `lowrank_shift_logz`, BIT for bit. + + `sign` is exactly +-1, so the multiply is a negation or a no-op and the remaining add rounds once + either way -- there is no reason for these to differ, and if they ever do, the gated backward's + `log N` is wrong by exactly that much on machines without the extension.""" + from pyjuice.nodes.external_params.kernels.c import get_module + from pyjuice.nodes.external_params.kernels.shift_logz import shift_logz_triton + + mod = get_module() + if mod is None: + pytest.skip("needs the CUDA extension to compare against") + + dev = torch.device("cuda:0") + torch.manual_seed(0) + num_nodes = 8192 + nids = (torch.randperm(num_nodes // block_size, device = dev)[:rows] * block_size).to(torch.int64) + base = torch.randn(num_nodes, batch, device = dev) + log_z = torch.randn(rows * block_size * batch, device = dev) + + for sign in (1.0, -1.0): + cuda_out, triton_out = base.clone(), base.clone() + mod.lowrank_shift_logz(cuda_out, nids, log_z, block_size, sign) + shift_logz_triton(triton_out, nids, log_z, block_size, sign) + assert torch.equal(cuda_out, triton_out), \ + f"sign={sign}: max|d| = {(cuda_out - triton_out).abs().max().item()}" + + @cuda_only @needs_cute @pytest.mark.parametrize("block_size,batch", [(128, 64), (128, 128)])