diff --git a/flashinfer/cake_causal_conv1d_train_backend/kcp.py b/flashinfer/cake_causal_conv1d_train_backend/kcp.py index d79ffb75e9c..c231862b552 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/kcp.py +++ b/flashinfer/cake_causal_conv1d_train_backend/kcp.py @@ -51,6 +51,24 @@ involved; each kernel stays within one BF16 ulp of its own oracle. ``dweight`` is the rank's fixed-order FP32 partial; it is reproducible across launches and within ``5e-2 + 1e-4 |dw|`` of an FP64 oracle, not bitwise the no-CP partial. + +Shipped host path (this revision; +``HOST_TRIM_DEFAULT = 1``): the context-derived host work of the eager halo +call -- the validated ``cu`` / ``hist_index`` / row-offset tables and halo +splits, the plan scalars, the grids, the record names and the loaded modules +with their positional argument plans -- is computed once per context and +reused, and the kernels launch through a raw-stream FFI context of the current +stream (no ``torch.cuda.Stream`` object per launch). ``_HOST_CACHE`` is keyed +on the context object's identity (the entry pins the context) and, inside it, +on ``(kind, S, C, device, cu_seqlens_local_cpu)`` -- host integers of the +context and the shapes of the call, never a tensor's contents; it holds at most +``HOST_CACHE_MAX = 64`` contexts (the oldest dropped) and +:func:`host_cache_clear` empties it; ``_FFI_DEVICES`` caches the FFI device per +device index. The operand checks, the gather, the exchange, its wait and the +kernel launches run as before with the same arguments on the same stream, so +no bit changes; a replayed graph (``graph="replay"``) contains no host work. +``_set_test_knobs({"host_trim": 0})`` (with ``CAKE_KCP_TEST_KNOBS=1``) selects +the previous host path for the A/B record. """ from __future__ import annotations @@ -85,7 +103,7 @@ GRAPH_MODES = ( "off", "replay", -) # graph replay: keyed CUDA-graph capture / replay of the halo calls; public keyword graph= since this revision +) # graph replay: keyed CUDA-graph capture / replay of the halo calls; public keyword graph= since an earlier revision GRAPH_DEFAULT = ( "off" # shipped: the eager path; the public keyword graph="replay" opts in per call ) @@ -221,17 +239,122 @@ def _context(ctx: Any, *, device: torch.device, T_r: int) -> tuple[torch.Tensor, # cached key; at most GRAPH_CACHE_MAX keys, the oldest dropped first); ``graph_cache_clear()`` releases them and must run # before ``torch.distributed.destroy_process_group()`` (a graph holding the NCCL collective blocks ``ncclCommDestroy``). -_TEST_KNOBS: dict[str, Any] = {"graph": "off"} +# Trimmed host path; shipped default: +# 1 = the trimmed host path of the eager halo call (same launches, bitwise; the module docstring's "Shipped host path"), +# 0 = the previous host path, selectable through _set_test_knobs for the A/B record. +HOST_TRIMS = (0, 1) +HOST_TRIM_DEFAULT = ( + 1 # shipped; 0 = the previous path, via _set_test_knobs({"host_trim": 0}) only +) +_TEST_KNOBS: dict[str, Any] = {"graph": "off", "host_trim": HOST_TRIM_DEFAULT} GRAPH_STATS: dict[str, int] = {"eager": 0, "capture": 0, "replay": 0, "fallback": 0} GRAPH_CACHE_MAX = 64 LAST_CALL: dict[str, dict[str, Any]] = {} +# -- the trimmed host path (host_trim=1): the per-context host state of the eager halo calls. ``_HOST_CACHE[id(ctx)]`` holds the +# context (pinned, so its id cannot be reused) and one entry per ``(kind, S, C, device, cu_seqlens_local_cpu)``: the +# validated context tables and tensors (``cu``, ``hist_row_off``, ...), the plan scalars, the record names, the loaded +# modules with their positional argument plans and grids. Everything in an entry derives from host integers of the +# context and the shapes of the call; no tensor content is read. At most HOST_CACHE_MAX contexts (the oldest dropped). +HOST_CACHE_MAX = 64 +_HOST_CACHE: dict[int, tuple[Any, dict[tuple, dict[str, Any]]]] = {} +_FFI_DEVICES: dict[int, Any] = {} + + +def _host_cache_entry(ctx: Any, key: tuple) -> dict[str, Any]: + """The (possibly empty, to be filled) host-cache entry of ``ctx`` under ``key``.""" + hit = _HOST_CACHE.get(id(ctx)) + if hit is None or hit[0] is not ctx: + if len(_HOST_CACHE) >= HOST_CACHE_MAX: + del _HOST_CACHE[next(iter(_HOST_CACHE))] + hit = _HOST_CACHE[id(ctx)] = (ctx, {}) + entry = hit[1].get(key) + if entry is None: + entry = hit[1][key] = {} + return entry + + +def host_cache_clear() -> int: + """Drop the trimmed host path's per-context cache; returns the number of contexts dropped.""" + n = len(_HOST_CACHE) + _HOST_CACHE.clear() + return n + + +def _ffi_device(device_index: int) -> Any: + dev = _FFI_DEVICES.get(int(device_index)) + if dev is None: + dev = _FFI_DEVICES[int(device_index)] = tvm_ffi.device( + f"cuda:{int(device_index)}" + ) + return dev + + +def _raw_stream(device_index: int) -> int: + get = getattr(torch._C, "_cuda_getCurrentRawStream", None) + if get is not None: + return int(get(int(device_index))) + return int(torch.cuda.current_stream(int(device_index)).cuda_stream) + + +def _launch_plan( + name: str, grid: tuple[int, int, int], device_index: int +) -> tuple[Any, str, tuple[tuple[str | None, int], ...]]: + """``(module, ffi entry, plan)`` of record ``name`` for the trimmed launch: ``plan`` lists ``(None, grid value)`` + per grid argument and ``(logical name, 0)`` per kernel parameter, in ``run`` order (the same arguments + :func:`run_module` / :func:`_launch_with_grid` build per call).""" + record = _loader.MODULES[name] + logical = record["logical"] + grid_values = { + "grid_x": int(grid[0]), + "grid_y": int(grid[1]), + "grid_z": int(grid[2]), + } + plan: list[tuple[str | None, int]] = [] + for kind, key in record["arg_plan"]: + if kind == "grid": + plan.append((None, grid_values[key])) + else: + if key not in logical: + raise ValueError( + f"module {name}: kernel parameter {key!r} has no logical name" + ) + plan.append((logical[key], 0)) + module = load(name, _loader.device_arch(device_index)) + return module, record["ffi_entry"], tuple(plan) + + +def _launch_trim( + info: tuple[Any, str, tuple[tuple[str | None, int], ...]], + values: dict[str, Any], + device_index: int, +) -> None: + """Launch a cached :func:`_launch_plan` with the per-call ``values`` on the device's current torch stream (raw-stream + FFI context: the same stream ``tvm_ffi.use_torch_stream`` would set, without building a torch.cuda.Stream).""" + module, entry, plan = info + args = [] + for logical_name, grid_value in plan: + if logical_name is None: + args.append(grid_value) + else: + try: + args.append(values[logical_name]) + except KeyError: + raise ValueError( + f"the launch needs logical argument {logical_name!r}" + ) from None + with tvm_ffi.use_raw_stream(_ffi_device(device_index), _raw_stream(device_index)): + getattr(module, entry)(*args) + def _set_test_knobs(knobs: Optional[dict]) -> dict: - """Set the module's private test knobs (``graph``: ``"off"`` shipped | ``"replay"``); ``None`` restores the - defaults. Refused without ``CAKE_KCP_TEST_KNOBS=1`` in the environment. Returns the active knobs.""" + """Set the module's private test knobs (``graph``: ``"off"`` shipped | ``"replay"``; ``host_trim``: ``1`` shipped = + the trimmed host path of the eager call | ``0`` the previous host path); a knob + the dict does not name keeps its shipped default; ``None`` restores the defaults. Refused without + ``CAKE_KCP_TEST_KNOBS=1`` in the environment. Returns the active knobs.""" if knobs is None: _TEST_KNOBS["graph"] = "off" + _TEST_KNOBS["host_trim"] = HOST_TRIM_DEFAULT return dict(_TEST_KNOBS) if os.environ.get(TEST_KNOBS_ENV) != "1": raise ValueError( @@ -239,17 +362,23 @@ def _set_test_knobs(knobs: Optional[dict]) -> dict: ) if not isinstance(knobs, dict): raise TypeError(f"_set_test_knobs expects a dict, got {type(knobs).__name__}") - unknown = sorted(set(knobs) - {"graph"}) + unknown = sorted(set(knobs) - {"graph", "host_trim"}) if unknown: raise ValueError( - f"_set_test_knobs: unknown knob(s) {unknown}; known: ['graph']" + f"_set_test_knobs: unknown knob(s) {unknown}; known: ['graph', 'host_trim']" ) graph = str(knobs.get("graph", "off")) if graph not in GRAPH_MODES: raise ValueError( f"_set_test_knobs: graph must be one of {GRAPH_MODES}, got {graph!r}" ) + host_trim = knobs.get("host_trim", HOST_TRIM_DEFAULT) + if isinstance(host_trim, bool) or host_trim not in HOST_TRIMS: + raise ValueError( + f"_set_test_knobs: host_trim must be one of {HOST_TRIMS}, got {host_trim!r}" + ) _TEST_KNOBS["graph"] = graph + _TEST_KNOBS["host_trim"] = int(host_trim) return dict(_TEST_KNOBS) @@ -418,6 +547,8 @@ def _kcp_forward( x = _require_tensor(x_local, "x_local") device = x.device + if _TEST_KNOBS["host_trim"] == 1: + return _kcp_forward_trim(x, weight, ctx, out=out, device=device) arch = _device_arch(device) S, C, x_flat = _tokens(x, "x_local", device=device) w = _weight(weight, C, device=device) @@ -471,6 +602,93 @@ def _kcp_forward( return y, hist +def _kcp_forward_trim( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + out: Optional[torch.Tensor], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Trimmed host path: :func:`_kcp_forward` with the context-derived host work cached per ``(ctx, S, C, device, + cu_seqlens_local_cpu)``: the validated ``cu`` / ``hist_index`` / ``hist_row_off`` tables, the halo splits, the + record, its plan scalars and grid, the loaded module with its argument plan. The operand checks, the gather, the + exchange, the wait and the kernel launch are the untrimmed call's, in the same order on the same stream.""" + + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) + e = _host_cache_entry( + ctx, ("fwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) + ) + if not e: + arch = _device_arch(device) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + e.update( + cu=cu, + n_frag=n_frag, + hist_index=_require_tensor( + ctx.halo_tables["hist_index"], "ctx.halo_tables['hist_index']" + ), + recv_rows=int(halo.n_hist_rows), + send_splits=tuple(halo.fwd_send_rows), + recv_splits=tuple(halo.fwd_recv_rows), + ) + if S != 0: + e["hist_row_off"] = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + (name,) = route_modules(arch, KCP_FWD) + device_index = _device_index(device) + record = _loader.MODULES[name] + scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"]) + grid = launch_grid(record, device_index=device_index) + if scalars["n_ch_tiles"] > grid[0]: + raise ValueError( + f"C={C} spans more channel tiles than the forward's grid on this device holds" + ) + e["scalars"] = _scalars(scalars) + e["device_index"] = device_index + e["launch"] = _launch_plan(name, grid, device_index) + hist_index = e["hist_index"] + send = ( + x_flat.index_select(0, hist_index) + if hist_index.numel() + else x_flat.new_empty((0, C)) + ) + hist, work = _exchange( + send, + recv_rows=e["recv_rows"], + send_splits=e["send_splits"], + recv_splits=e["recv_splits"], + ctx=ctx, + ) + if S == 0: + if work is not None: + work.wait() + return y, hist + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": e["cu"], + "y": y.view(S, C), + "hist": hist + if hist.shape[0] + else _pool(0, C, dtype=torch.bfloat16, device=device), + "hist_row_off": e["hist_row_off"], + **e["scalars"], + } + if work is not None: + work.wait() + _launch_trim(e["launch"], values, e["device_index"]) + return y, hist + + def cake_causal_conv1d_kcp_backward( x_local: torch.Tensor, weight: torch.Tensor, @@ -522,6 +740,10 @@ def _kcp_backward( x = _require_tensor(x_local, "x_local") device = x.device + if _TEST_KNOBS["host_trim"] == 1: + return _kcp_backward_trim( + x, weight, ctx, dy, hist, dx=dx, dweight=dweight, device=device + ) arch = _device_arch(device) S, C, x_flat = _tokens(x, "x_local", device=device) w = _weight(weight, C, device=device) @@ -634,6 +856,176 @@ def _kcp_backward( return dx_out, dw_out +def _kcp_backward_trim( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + dy: torch.Tensor, + hist: torch.Tensor, + *, + dx: Optional[torch.Tensor], + dweight: Optional[torch.Tensor], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Trimmed host path: :func:`_kcp_backward` with the context-derived host work cached per ``(ctx, S, C, device, + cu_seqlens_local_cpu)`` (tables, the head grid, the backward records, plan scalars, partial rows, the loaded modules + with their argument plans); the operand checks, allocations, exchange, wait and launches are the untrimmed call's.""" + + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) + hist_t = _require_tensor(hist, "hist") + if ( + hist_t.ndim != 2 + or int(hist_t.shape[1]) != C + or hist_t.dtype != torch.bfloat16 + or not hist_t.is_contiguous() + ): + raise ValueError( + f"hist must be the forward's contiguous BF16 [n_hist_rows, {C}] pool, got {hist_t.dtype} {tuple(hist_t.shape)}" + ) + if hist_t.shape[0] == 0: + hist_t = _pool(0, C, dtype=torch.bfloat16, device=device) + dx_out = _output( + dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device + ) + dw_out = _output( + dweight, "dweight", shape=(C, WIDTH), dtype=torch.float32, device=device + ) + e = _host_cache_entry( + ctx, ("bwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) + ) + if not e: + arch = _device_arch(device) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + device_index = _device_index(device) + n_head = int(halo.n_head_rows) + e.update( + cu=cu, + n_frag=n_frag, + device_index=device_index, + n_head=n_head, + tail_rows=int(halo.n_tail_rows), + send_splits=tuple(halo.bwd_send_rows), + recv_splits=tuple(halo.bwd_recv_rows), + head=None, + ) + if S > 0 and n_head > 0: + hist_row_off = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + head_row_off = _table( + ctx, + "head_row_off", + shape=(n_frag, HALO, HEAD_SLOTS), + dtype=torch.int32, + device=device, + ) + (head_name,) = route_modules(arch, KCP_BWD_HEAD) + geometry = _loader.MODULES[head_name]["grid_rule"]["geometry"] + channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) + grid = (max(1, n_frag * _ceil_div(C, channels_per_cta)), 1, 1) + e["head"] = ( + _launch_plan(head_name, grid, device_index), + hist_row_off, + head_row_off, + ) + if S != 0: + e["hist_row_off"] = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + e["tail_row_off"] = _table( + ctx, + "tail_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + names = selected_modules(arch, KCP_BWD, C=C, S=S) + record = _loader.MODULES[names[0]] + scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"]) + grid_x = launch_grid(record, device_index=device_index)[0] + if scalars["n_ch_tiles"] > grid_x: + raise ValueError( + f"C={C} spans more channel tiles than the backward's grid on this device holds" + ) + e["rows"] = partial_rows(grid_x, scalars["n_ch_tiles"]) + e["scalars"] = _scalars(scalars) + e["launches"] = tuple( + _launch_plan( + name, + launch_grid(_loader.MODULES[name], device_index=device_index), + device_index, + ) + for name in names + ) + n_head = e["n_head"] + device_index = e["device_index"] + dpre_head = _pool(n_head, C, dtype=torch.float32, device=device) + if e["head"] is not None: + info, hist_row_off, head_row_off = e["head"] + _launch_trim( + info, + { + "x": x_flat, + "w": w, + "cu": e["cu"], + "dy": dy_flat, + "hist": hist_t, + "hist_row_off": hist_row_off, + "dpre_head": dpre_head, + "head_row_off": head_row_off, + "S": S, + "C": C, + "num_docs": e["n_frag"], + }, + device_index, + ) + dpre_tail, work = _exchange( + dpre_head[:n_head], + recv_rows=e["tail_rows"], + send_splits=e["send_splits"], + recv_splits=e["recv_splits"], + ctx=ctx, + ) + if S == 0: + if work is not None: + work.wait() + dw_out.zero_() + return dx_out, dw_out + partial = torch.empty((e["rows"], C, WIDTH), dtype=torch.float32, device=device) + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": e["cu"], + "dy": dy_flat, + "dx": dx_out.view(S, C), + "dw": dw_out, + "partial": partial, + "hist": hist_t, + "hist_row_off": e["hist_row_off"], + "dpre_tail": dpre_tail + if dpre_tail.shape[0] + else _pool(0, C, dtype=torch.float32, device=device), + "tail_row_off": e["tail_row_off"], + **e["scalars"], + } + if work is not None: + work.wait() + for info in e["launches"]: + _launch_trim(info, values, device_index) + return dx_out, dw_out + + class CakeCausalConv1dKcpFunction(torch.autograd.Function): """Autograd wrapper: the forward saves ``(x_local, weight)`` and the history pool ``hist`` (no second exchange in the backward); the backward runs the head program, the backward halo exchange and the main backward program.""" diff --git a/flashinfer/experimental/cake_kda_train/README.md b/flashinfer/experimental/cake_kda_train/README.md index ad19d1dc90f..2b6adab1045 100644 --- a/flashinfer/experimental/cake_kda_train/README.md +++ b/flashinfer/experimental/cake_kda_train/README.md @@ -307,7 +307,8 @@ the programmatic-stream-serialization attribute). Rules: "serial" `H > 8`, The packed batch of one training step may be split over the ranks of a `torch.distributed` process group. `kcp_partition` cuts the documents into per-rank fragments (`"contiguous"`, `"zigzag"` or `"zigzag_per_document"` -layouts, or explicit cut points); `build_kcp_context` gives every rank the +layouts, an explicit `fragments` table, or `cut_points` with the +`"zigzag_per_document"` layout); `build_kcp_context` gives every rank the device tables of its fragments, the collectives that carry the per-document state recurrences and the convolution's history rows across the rank boundaries, and the segment plan of its fragment walks. The per-rank entry @@ -330,9 +331,11 @@ dq, dk, dv, dbeta, dg = chunk_kda_kcp_bwd(q_local, k_local, v_local, g_local, be dout_local, saved=saved) ``` -CUDA-graph replay: every KCP entry point -- `chunk_kda_kcp`, -`chunk_kda_kcp_fwd`, `chunk_kda_kcp_bwd` and `cake_causal_conv1d_kcp` -- takes -`graph="replay"` (default `"off"`, the eager path). The first call of a +CUDA-graph replay: the three KDA entry points -- `chunk_kda_kcp`, +`chunk_kda_kcp_fwd` and `chunk_kda_kcp_bwd` -- take `graph="replay"` +(default `"off"`, the eager path); the convolution halo takes `graph=` only +through the backend module's function, the `cake_causal_conv1d_kcp` facade +has no such parameter. The first call of a validity key runs eagerly, the second is captured into a CUDA graph and every later call of the key replays it: the same launches and collectives, bitwise the eager result, with no host work between them. The key is the call's kind, @@ -347,17 +350,25 @@ and `LAST_CALL[kind]["graph_event"]` say what happened; a capture failure disables that key for the process), `exchange="ring"`, `debug_sequence` and calls without a caller-owned workspace are never captured. Memory: every cached key pins a private memory pool -- 0.7-9.0 GiB per rank per key on the -the A/B rows (the per-layer forward and backward of CP4-CP16 at the B1 / -B3 batch sizes), at most `GRAPH_CACHE_MAX = 64` keys (LRU). Call +A/B rows (the per-layer forward and backward of CP4-CP16 at the B1 / +B3 batch sizes), at most `GRAPH_CACHE_MAX = 64` keys (the oldest key is +evicted). Call `graph_cache_clear()` (both modules) before `torch.distributed.destroy_process_group()` and whenever a layer's operands -are gone for good. A captured call runs its one-CTA-per-SM tape programs -without the device cache-config bracket of the shipped walk geometry (the -bracket cannot be issued inside a capture): on the CP8 / CP16 rows replay is -still 1.9-9.1 % faster than the eager default, on the CP4 row 1.3-2.3 % slower -(the eager default keeps the full geometry gain; `graph="off"` is the shipped -default). The A/B record was taken with pointer-stable harness inputs; its -gain is an upper bound of what a framework sees. +are gone for good. A replayed call keeps the tape programs at one CTA per SM: +the device cache-config bracket of the eager path cannot be issued inside a +capture, so right after each tape launch is captured the host sets the +preferred shared-memory carveout as a launch attribute of the kernel node it +just added (`cudaStreamGetCaptureInfo` + `cudaGraphKernelNodeSetAttribute`, +host-only `ctypes` calls into the runtime that is already loaded) and the +replayed graph carries the cap itself. Without those entry points, or when the +runtime cannot name the captured node, the capture continues without the cap +after one warning (`LAST_CALL[kind]["cache_config"] == "unavailable"`); the +results are bitwise identical either way. The previous replay form stays +reachable for A/B records through the private test knob `graph_cap` +(`CAKE_KCP_TEST_KNOBS=1`); `graph="off"` is the shipped default. The A/B +record was taken with pointer-stable harness inputs; its gain is an upper +bound of what a framework sees. Limits of a partition: at most 65536 fragments per rank (the layout table's document search), at most 64 segments per fragment and no walk chain longer @@ -367,6 +378,15 @@ than 256 chunk steps (the segment plan's bounds); the byte queries one step's traffic before any allocation. Every rank of the group calls the same entry points in the same order; the collectives are issued inside them. +Saved set of a rank (`KcpSaved`, `boundary_mode="saved"`): the fresh `S_in` +(the initial states of its fragments), with `save_transfer` the owner's c1 +receive tensor and `M_seg`, and the call's scalars; the int32 tables are +rebuilt from the context. The token-row `Aqk` / `Akk` of the single-rank +saved set are not part of it: the backward recomputes the internal rows in +both modes (saved set and checkpoint replay run the same local programs), so +`kcp_saved_bytes` counts the tables, `S_in` and the `save_transfer` tensors +(and the convolution's history rows with `with_conv`). + Walk-launch geometry: the walk kernels' CTA lists are issued longest-segment-first and the two tape programs (`fwdh_tape`, `dhu_tape`) run at one CTA per SM. Because the JIT binding exposes no per-function carveout @@ -378,6 +398,19 @@ itself (occupancy only -- results never depend on the geometry). The earlier geometry stays reachable for A/B records through the private test knob `walk_geometry` (`CAKE_KCP_TEST_KNOBS=1`). +Host path: the eager KCP calls reuse the call-independent host work -- the +validated options, the workspace's last tables and pool views, the local +stages' launch templates and the table set's own device tables -- keyed on +object identity plus host integers and bounded (`ID_CACHE_MAX = 256` entries +per table-set cache; the conv halo module keeps at most `HOST_CACHE_MAX = 64` +contexts). The launch templates live on the workspace's local arena and +reference the table set only weakly, so a training loop that builds a new +context and workspace per step releases everything with them. The kernels, +their arguments and their order are unchanged, so the results are bitwise +those of the previous host path, which stays reachable for A/B records +through the private test knob `host_trim` (`CAKE_KCP_TEST_KNOBS=1`; the conv +module's `_set_test_knobs({"host_trim": 0})`). + Options of the KDA entry points: `save_transfer` keeps the owner's c1 receive tensor (and `M_seg`) for the backward instead of re-sending it in the backward exchange -- the default with the all-to-all exchange (an earlier revision). @@ -387,22 +420,35 @@ automatically; an explicit `save_transfer=True` together with the ring is a `ValueError`. `KcpSaved.save_transfer` always reports the value that was used. The segment plan splits every fragment into `P` segments: above the 256-chunk -chain cap the equal-segment wave rule, below it on one-fragment ranks the +chain cap the equal-segment wave rule on ranks with at most three fragments, +below it on one-fragment ranks the largest `P` that runs one CTA per SM (`floor(SM / H)`), on every other rank -one CTA-per-SM occupancy (`ceil(SM / (n_frag H))`); -`ctx.describe(H)["p_rule"]` records the rule (`"wave_fill1"`). +one CTA-per-SM occupancy (`ceil(SM / (n_frag H))`) -- except that a rank +with two or three fragments below the cap whose `P = 1` walks run one wave +(`n_frag H <= SM`) takes `P = 1` when the longest-processing-time makespan of +its actual fragment chunk counts at `P = 1` is not above the makespan at the +occupancy `P` (such a rank runs no segment composition and no backward +re-walk; at H = 64 on 152 SMs equal two-fragment ranks tie, unequal ones +keep `P > 1`); +`ctx.describe(H)["p_rule"]` records the rule (`"wave_fill1_p1tie"`). The +previous rules stay selectable through the private test knob `p_rule`, as +does the fused form of the carry merges (`merge_fuse="carry"`: the two carry +families of a phase in one launch; off by default). The KCP kernels live beside the single-rank programs: -* `cake_kcp_jit.py` -- the `MODULES` registry of the seven KCP programs (the +* `cake_kcp_jit.py` -- the `MODULES` registry of the nine KCP programs (the fragment walks of the forward and backward state recurrences in their - `offset` and `tape` forms, the TF32 hi/lo transfer walk and the affine - merge at `cols_per_cta` 8 and 16; selection by `always` or by `columns` on - the merge's `ncols`, grids by `kcp_segments_heads` or `kcp_chains`), the - logical vocabulary of the KCP arguments, and `run_program(name, args, - device_index=)`, the lowering of the host backend's walk and merge requests - to the registry's launches. Its `TEST_MODULES` are test-only programs of - `build.test.json` (below). + `offset` and `tape` forms, the TF32 hi/lo transfer walk, the affine + merge at `cols_per_cta` 8 and 16, and the two chunk stages of the + context-parallel route, `kcp_fwd_o` (the output from the K1 tapes) and + `kcp_dav` (`dA_qk` and the `dv1` K2 reads); selection by `always` or by + `columns` on the merge's `ncols`, grids by `kcp_segments_heads`, + `kcp_chains`, `kcp_sm_bounded` or `kcp_pairs`), the logical vocabulary of + the KCP arguments, `run_program(name, args, device_index=)`, the lowering + of the host backend's walk and merge requests to the registry's launches, + and `run_module` for the two chunk stages. Its `TEST_MODULES` are + test-only programs of `build.test.json` (below). * `cake_kcp_plan.py` -- the partition, the segment plan and the context's tables and collectives, with the byte queries. * `cake_kcp.py` -- the per-rank launch sequences, the saved set and workspace @@ -415,7 +461,7 @@ The KCP kernels live beside the single-rank programs: entry only under the environment variable `CAKE_KCP_TEST_KNOBS=1`, which also admits the test knobs (`precision_mode`, `transfer_walk`) of the package's own tests; without the file those tests skip. -* `csrc/cake_kda_train__{kernel,binding}.cu` of the eleven KCP +* `csrc/cake_kda_train__{kernel,binding}.cu` of the thirteen KCP programs, rendered like the single-rank ones (digests in `build.json` and `build.test.json`). diff --git a/flashinfer/experimental/cake_kda_train/build.json b/flashinfer/experimental/cake_kda_train/build.json index bec77ab336b..75d76ab20c1 100644 --- a/flashinfer/experimental/cake_kda_train/build.json +++ b/flashinfer/experimental/cake_kda_train/build.json @@ -55,9 +55,9 @@ } }, "generator": { - "cake_revision": "25e558c241041e5918e9b5144d62e25ad0eedc18", + "cake_revision": "2147cb368e75839a4023c7545de2fc5466c5c288", "cuda": "13.4.92", - "rendered_at": "2026-10-21T13:33:51Z" + "rendered_at": "2026-10-21T03:20:39Z" }, "libraries": [ { @@ -312,6 +312,69 @@ "csrc/cake_kda_train_0fb763e9a81e9066d776_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "2da0b8ee88cebac0e5b357c56196e6bb5a0a8c2dc1ea1371dcdebe9857e429c1", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_train_1214bfc28e50985bfa0c_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_1214bfc28e50985bfa0c_sm_100a.so", + "sha256": { + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu": "9f2aac4088fe8d71cd1daf6934b36f3840b96c728db110cf740d2a45fefdfb5b", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu": "e9b3689fadf5197ac1286a092d18a75866f8ca951bf155d0ca4514338dc0ee4b" + }, + "sources": [ + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "4836a553a611e44837ce06ab2f70f5addf432a98fdfecb6e8209b886c833a18c", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_train_1214bfc28e50985bfa0c_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_1214bfc28e50985bfa0c_sm_103a.so", + "sha256": { + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu": "9f2aac4088fe8d71cd1daf6934b36f3840b96c728db110cf740d2a45fefdfb5b", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu": "e9b3689fadf5197ac1286a092d18a75866f8ca951bf155d0ca4514338dc0ee4b" + }, + "sources": [ + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "dab112a5b9f0c8c4b9282faa78f3005b2ba9cf5f19fed97eac18970ad8e0be5e", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.0", + "name": "cake_kda_train_1214bfc28e50985bfa0c_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_1214bfc28e50985bfa0c_sm_107a.so", + "sha256": { + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu": "9f2aac4088fe8d71cd1daf6934b36f3840b96c728db110cf740d2a45fefdfb5b", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu": "e9b3689fadf5197ac1286a092d18a75866f8ca951bf155d0ca4514338dc0ee4b" + }, + "sources": [ + "csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu", + "csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "1ae0ef57299ee2aa7d7cb8be93fee412d3fc4e7247f542dc53e981e47557881b", @@ -4029,6 +4092,69 @@ "csrc/cake_kda_train_f1777a88a21b832b29ee_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "e6b155bbbf0200ac1d8b16a94da63352ebb65b9f58890618ac3691ac4814cb5e", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_train_f227dd21f1f3cb464b29_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_f227dd21f1f3cb464b29_sm_100a.so", + "sha256": { + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu": "7fef7e466d7e81f98bf1c3a19450f94b11b9327c706914597a5c9bbeb1aadd1c", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu": "49b7f927255d4dff348744c660d7cf6689fcf5fc6dcfdfc86247689ed4177164" + }, + "sources": [ + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "6608e2615106b089efdb681f004954cad9610dc41d747a012f3f776c3e760cdd", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_train_f227dd21f1f3cb464b29_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_f227dd21f1f3cb464b29_sm_103a.so", + "sha256": { + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu": "7fef7e466d7e81f98bf1c3a19450f94b11b9327c706914597a5c9bbeb1aadd1c", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu": "49b7f927255d4dff348744c660d7cf6689fcf5fc6dcfdfc86247689ed4177164" + }, + "sources": [ + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "0f2fba2d9da2abbd8728c0b972c380d79884abd08fc3d16f8d8edeffe2572a2c", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.0", + "name": "cake_kda_train_f227dd21f1f3cb464b29_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_train_f227dd21f1f3cb464b29_sm_107a.so", + "sha256": { + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu": "7fef7e466d7e81f98bf1c3a19450f94b11b9327c706914597a5c9bbeb1aadd1c", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu": "49b7f927255d4dff348744c660d7cf6689fcf5fc6dcfdfc86247689ed4177164" + }, + "sources": [ + "csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu", + "csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "79c04925757402c0784766aeb83b9e2f962b2ac353424c151b83ad2820243552", diff --git a/flashinfer/experimental/cake_kda_train/cake_kcp.py b/flashinfer/experimental/cake_kda_train/cake_kcp.py index 79d4a3efb80..f8fd6212149 100644 --- a/flashinfer/experimental/cake_kda_train/cake_kcp.py +++ b/flashinfer/experimental/cake_kda_train/cake_kcp.py @@ -32,8 +32,8 @@ owner walks each chain backwards (``X <- M_f^T X + d_f``) and returns every crossing-out fragment's ``dS_out`` (``b2``); the local adjoint recurrence runs from it. One compute stream; the collectives are asynchronous on the caller's -process group and waited at their first consumer. Shipped schedule (an earlier -adoption of levers L-A + L-B): ``wait_first`` -- no merge +process group and waited at their first consumer. Shipped schedule (adopted in an +earlier revision): ``wait_first`` -- no merge or tape walk is issued while a collective is in flight, and the tape walks of a phase are one fused launch from the carry plane; the earlier ``overlap`` form (the document-start tape walk under ``c1``, the document-last adjoint tape walk @@ -41,7 +41,7 @@ the test knobs ``schedule`` / ``fused_walks``. Nothing launched between an issue and its wait reads that collective's buffers, so the bits are independent of the schedule and a checkpoint replay re-runs the same program order on every rank. -CUDA-graph replay (adopted in this revision; keyword +CUDA-graph replay (adopted in an earlier revision; keyword ``graph="replay"`` of :func:`chunk_kda_kcp`, :func:`chunk_kda_kcp_fwd` and :func:`chunk_kda_kcp_bwd`, shipped default ``GRAPH_DEFAULT = "off"``): the first call of a validity key runs eagerly, the second is captured into a @@ -60,7 +60,7 @@ ``exchange="ring"``, ``debug_sequence`` and calls without a caller-owned workspace are never captured. -Shipped walk-launch geometry (adopted in this revision; +Shipped walk-launch geometry (adopted in an earlier revision; ``WALK_GEOMETRY_DEFAULT = "lpt_tape_one_cta_per_sm"``): the offset walk lists of K1 / K2 / K3 launch in LPT order (longest segment first, ties by segment number; the tape lists already did) and the tape programs ``fwdh_tape`` / ``dhu_tape`` @@ -73,18 +73,84 @@ The bracket is host-side runtime state, so it is never issued inside a CUDA stream capture (``cudaDeviceSetCacheConfig`` would invalidate the capture, -see below): under ``graph="replay"`` the tape one-CTA-per-SM launch cap -(the walk-launch geometry) is not applied because ``cudaDeviceSetCacheConfig`` is incompatible -with stream capture; at CP4 C B3 replay therefore gains about 2 pp less than the -eager default, at CP8 / CP16 replay's own gain dominates. The -test knob ``graph_cap`` selects whether the preference is instead set once before -the capture of a key and reset right after it ends (``"capture_pref"``: PreferL1 -for every kernel of the captured call, 6.8-18 % slower on every composition row -of the gate -- rejected, a gated knob only) or left alone (``"none"``, the form of -record); every call records its form as ``LAST_CALL[kind]["cache_config"]`` -(``"bracket"`` / ``"capture_pref"`` / ``"none"`` / ``"unavailable"``). When no CUDA runtime soname loads, the -shipped geometry keeps the LPT order, launches the tape programs without the -cap after one ``RuntimeWarning`` and records ``"unavailable"``. +see below). Under ``graph="replay"`` the captured call carries the tape +one-CTA-per-SM cap itself (the shipped +``GRAPH_CAP_DEFAULT = "node_attr"``): right after each tape launch is +captured, the host reads the kernel node it just added +(``cudaStreamGetCaptureInfo``; CUDA 12 runtimes export that signature as +``cudaStreamGetCaptureInfo_v2``, CUDA 12.3+ also as ``_v3``) and sets its +preferred shared-memory carveout launch attribute +(``cudaLaunchAttributePreferredSharedMemoryCarveout`` = +``cudaSharedmemCarveoutMaxL1``, ``cudaGraphKernelNodeSetAttribute`` with a +``cudaGraphKernelNodeGetAttribute`` read-back), so every replay runs the tape +programs at one CTA per SM like the eager bracket does; the eager path is +untouched. The test knob ``graph_cap`` keeps the other forms selectable for +the A/B record: ``"none"`` (the previous form of record: no +device preference, the tape programs replay uncapped), ``"capture_pref"`` +(PreferL1 set once before the capture of a key and reset right after it ends: +6.8-18 % slower on every composition row of the gate -- rejected) and +``"node_attr_post"`` (the same attribute set on the kept node handles after +the capture ends and before the graph is instantiated, +``torch.cuda.CUDAGraph(keep_graph=True)``). Every call records its form as +``LAST_CALL[kind]["cache_config"]`` (``"bracket"`` / ``"node_attr"`` / +``"node_attr_post"`` / ``"capture_pref"`` / ``"none"`` / ``"unavailable"``); +the node forms also record ``cache_config_nodes`` = the number of capped tape +launches of the captured call. Degrade contract: when the +runtime lacks one of the capture-graph entry points, or the capture-info call +does not report exactly one kernel-node dependency after a tape launch, the +capture continues UNCAPPED (the ``"none"`` behaviour) after one +``RuntimeWarning`` per process and the call records ``cache_config`` = +``"unavailable"`` (``cache_config_degraded`` = the reason) -- never an eager +fallback; an error returned by the set-attribute call keeps the capture-failure +contract (that key runs eagerly, counted in ``GRAPH_STATS["fallback"]``). When +no CUDA runtime soname loads, the shipped geometry keeps the LPT order, +launches the tape programs without the cap after one ``RuntimeWarning`` and +records ``"unavailable"``. Host-only ``ctypes`` entry points +of this module: +``cudaDeviceSetCacheConfig``, ``cudaDeviceGetCacheConfig``, +``cudaGetLastError``, ``cudaRuntimeGetVersion``, ``cudaStreamGetCaptureInfo`` +(``_v2`` / ``_v3``), ``cudaGraphNodeGetType``, +``cudaGraphKernelNodeSetAttribute``, ``cudaGraphKernelNodeGetAttribute``. + +Shipped host path (this revision; +``HOST_TRIM_DEFAULT = 1``): the host work of an eager call that does not depend +on the call is done once and reused. The entry point validates ``**options`` +once and the autograd node carries the :class:`KcpOptions`; +``KcpWorkspace.ensure`` returns the last tables for the same (context, H, +device, direction, knobs) and keeps the ``pool[:n]`` views of an unchanged pool +tensor; the local stages launch from a per-(arena, table set, direction) value +template with a cached argument plan, arch and raw-stream FFI context; the +K1-K4 launches skip the positional ABI that the ``run_program`` provider form +discards and pass the table set's own device tables. Every cache is keyed on +object identity plus host integers and never reads a tensor's contents: +``_Tables.dev_list`` (the walk / chunk-start list object, device, dtype), +``_Tables.indexed_meta`` (the index and tile list objects, H, device), +``_Tables.device_tiles`` by identity (the tile table object, device), +``Chains.device_tables`` (the chain set, device), ``_Tables.walk_all_p0`` / +``walk_all_rest`` (the P = 1 tape-replay sublists, one object per table set), +the workspace's last ``ensure`` key and its ``(pool name, n)`` views while the +pool tensor is the same object (the views of a replaced pool are dropped), the +local arena's ``begin_call`` templates ((table set, direction) -> the arena +views and plan scalars) and the per-device-index arch / FFI device of this +module and of :mod:`.cake_kcp_jit`. Owners and bounds: a table +set's identity-keyed caches hold at most ``ID_CACHE_MAX = 256`` entries each +(emptied beyond; a caller passing fresh objects cannot grow them) and die with +the table set; the chain tables die with their chain set; the ``ensure`` key and +the pool views die with the workspace; the templates live on the arena they +describe (``TEMPLATE_CACHE_ATTR``) and reference the table set only weakly (a +dead table set's entry goes on the next miss), so they die with the arena -- the +backend itself holds no workspace, arena or table set (an earlier form +kept them on the backend for the process lifetime); the conv halo module bounds its +per-context cache at +``HOST_CACHE_MAX = 64`` contexts. The same kernels launch with the same +arguments in the same order on the same stream, so no bit changes (bitwise on +every A/B and confirmation row of the adoption); under ``graph="replay"`` the +trimmed segments are not part of the replayed call. +``_kcp_test_knobs={"host_trim": 0}`` selects the previous host path for the +A/B record; the D1 explicit-record merge form (``cols_rule="by_chains"`` beyond +``COLS_RULE_CHAINS`` state chains) needs the positional ABI and a diagnostics +launcher on the arena keeps ``KdaWorkspace.run``: both keep the previous path +under either value of the knob. The metadata (partition, context, head ownership, chain tables, split sizes, segment plan, compact-buffer addressing, byte queries) is @@ -111,15 +177,17 @@ import ctypes import dataclasses +import functools import os import warnings +import weakref from collections.abc import Sequence from typing import Any import torch import torch.distributed as dist -from . import cake_backend, cake_kcp_plan +from . import cake_backend, cake_jit, cake_kcp_plan D = cake_backend.D # head dimension (128) BT = cake_backend.BT # tokens per chunk (64) @@ -136,20 +204,20 @@ "bf16x3", "bf16x3_hw64", ) # K3 program family (non-default only via _kcp_test_knobs) # exchange="auto" picks the ring when every chain has at most this many hops (SPEC-2) -KCP_SAVED_VERSION = 1 # saved set {Aqk, Akk, tables, S_in [, c1 receive tensor, M_seg]} +KCP_SAVED_VERSION = 2 # saved set {tables, S_in [, c1 receive tensor, M_seg]}; version 1 also held the token-row Aqk / Akk TEST_KNOBS_ENV = "CAKE_KCP_TEST_KNOBS" -# Host / plan levers (test knob class: reachable only through _kcp_test_knobs): +# Host / plan knobs (test knob class: reachable only through _kcp_test_knobs): SCHEDULES = ( "overlap", "wait_first", -) # L-A: "wait_first" (shipped) = merges and tape walks after each wait; "overlap" = the earlier R7 overlaps of the design note §3.5 / §6.5 (test knob) +) # schedule: "wait_first" (shipped) = merges and tape walks after each wait; "overlap" = the earlier overlapped issue order of the design note §3.5 / §6.5 (test knob) COLS_RULES = ( "fixed", "by_chains", ) # D1: K4 cols_per_cta of the 128-column state launches: the backend's, or 8 up to COLS_RULE_CHAINS chains else 16 COLS_RULE_CHAINS = 24 # D1 threshold (DECISIONS.md D1: <= 24 concurrent chains -> c8, otherwise c16; map launches stay c16) # Walk-launch geometry: the launch geometry of the K1 / K2 / K3 walk kernels, a test -# knob until this revision made "lpt_tape_one_cta_per_sm" the shipped default. "shipped" = +# knob until an earlier revision made "lpt_tape_one_cta_per_sm" the shipped default. "shipped" = # the earlier lists and launches (the offset / staging lists in segment order, the runtime's own occupancy); "lpt_chains" = # every walk list in LPT order (longest segment first, ties by segment number; the tape lists already are); "one_cta_per_sm" = # the K1 / K2 kernels (115712 B of dynamic shared memory: two CTAs fit an SM only at the full 228 KiB carveout) launched at @@ -187,31 +255,62 @@ GRAPH_MODES = ( "off", "replay", -) # graph replay: "replay" = keyed CUDA-graph capture / replay of the per-rank call; public keyword graph= since this revision +) # graph replay: "replay" = keyed CUDA-graph capture / replay of the per-rank call; public keyword graph= since an earlier revision GRAPH_DEFAULT = "off" # shipped: the eager path; "replay" is opt-in per call (section "CUDA-graph replay of a call" below) GRAPH_CAPS = ( "none", "capture_pref", -) # the cache-config form of a captured call -- "none" (no device preference under replay) | "capture_pref" (PreferL1 set before the capture of a key, PreferNone right after its end) -GRAPH_CAP_DEFAULT = "none" # of record ("capture_pref" lost 6.8-18 % on every row of the composition gate and stays a gated knob, recorded as rejected); _kcp_test_knobs only -P2_KNOBS = { + "node_attr", + "node_attr_post", +) # the cache-config form of a captured call -- "none" (no device preference under replay) | "capture_pref" (PreferL1 set before the capture of a key, PreferNone right after its end) | the kernel-node forms ("node_attr" = the replay form of record): "node_attr" (the carveout launch attribute set on each captured tape kernel node right after its launch, inside the capture) | "node_attr_post" (the same attribute set on the kept node handles after the capture ends, before the graph is instantiated) +GRAPH_CAP_DEFAULT = "node_attr" # the graph="replay" cache-config form of record -- the tape one-CTA-per-SM cap as a launch attribute of the captured kernel nodes; "none" (the previous form of record), "capture_pref" (rejected: 6.8-18 % slower on every composition row) and "node_attr_post" stay selectable through _kcp_test_knobs for the A/B record; the eager path is untouched +GRAPH_NODE_CAPS = ( + "node_attr", + "node_attr_post", +) # the kernel-node forms (the cap as a kernel-node launch attribute); "node_attr" shipped +GRAPH_NODE_CARVEOUT = 0 # cudaSharedmemCarveoutMaxL1 (driver_types.h): prefer the maximum L1 = the smallest shared-memory configuration that fits the kernel -> one 115712-B tape CTA per SM, the per-node form of WALK_CACHE_CONFIG_CAP +_LAUNCH_ATTR_SHMEM_CARVEOUT = 14 # cudaLaunchAttributePreferredSharedMemoryCarveout (driver_types.h, CUDA >= 12.4: valid for graph nodes and launches) +_GRAPH_NODE_TYPE_KERNEL = 0 # cudaGraphNodeTypeKernel +_STREAM_CAPTURE_ACTIVE = 1 # cudaStreamCaptureStatusActive +# Trimmed host path: host-side trimming of the eager call's Python / +# dispatch time (the start-phase host idle of the profile). Shipped default: +# 1 = the trimmed host path; 0 = the previous host path, selectable through _kcp_test_knobs for the A/B record. Host side only: +# the same kernels with the same arguments in the same order on the same stream; caches hold host ints keyed on the context / +# tables / options objects (never a tensor's contents); bitwise by construction (module docstring, "Shipped host path"). +HOST_TRIMS = (0, 1) +HOST_TRIM_DEFAULT = ( + 1 # shipped: the trimmed host path; 0 = the previous path, via _kcp_test_knobs only +) +# Fused carry merges (test knob): "carry" launches the two forward carry +# families carry_start + carry_in of a rank as ONE K4 launch (chain family "carry_both") and the two adjoint families +# dcarry_end + dcarry_cross as one ("dcarry_both") -- the same m / c / y pools, the start (end) chains keep x0_index = -1 +# (zero state) beside the S_in (dS_out) tiles of the crossing chains; one launch saved per phase on ranks holding both a +# document-start (document-last) and a crossing-in (crossing-out) fragment at P > 1. Chain order inside a launch is the +# CTA issue order only (every CTA walks its own chain / column block): bitwise. Applies under schedule wait_first (the +# overlap schedule issues carry_start / dcarry_end under the collective) and in the saved-mode backward's re-walk merges. +MERGE_FUSES = ("none", "carry") +MERGE_FUSE_DEFAULT = "none" # shipped: the separate launches; "carry" via _kcp_test_knobs (for the A/B record) +P2_KNOBS = { # knob -> admitted values "schedule": SCHEDULES, "fused_walks": (0, 1), "cols_rule": COLS_RULES, "walk_geometry": WALK_GEOMETRIES, "graph": GRAPH_MODES, "graph_cap": GRAPH_CAPS, -} # knob -> admitted values -# Shipped defaults of the schedule / fused-walk pair (adoption of levers L-A + L-B together, the keep + "host_trim": HOST_TRIMS, + "merge_fuse": MERGE_FUSES, +} +# Shipped defaults of the schedule / fused-walk pair (adopted together, the keep # rule): the "wait_first" schedule with the fused tape walks (one launch per phase). "overlap" / fused_walks=0 stay -# selectable through the test knobs for the A/B record; cols_rule stays "fixed" (D1 in band everywhere: not adopted, +# selectable through the test knobs for the A/B record; cols_rule stays "fixed" (the other column rule was in band everywhere: not adopted, # knob kept). The schedule and the walk grouping never change a published bit (the merges read c1_recv / the staging # planes, the tapes read no collective buffer): G3 holds under every combination. SCHEDULE_DEFAULT = "wait_first" FUSED_WALKS_DEFAULT = 1 -# Program names of the dispatch table (the K1-K4 programs of the design note §14) and the local chunk stages -# (``cake_backend`` stages at the per-rank shape, launched through ``KdaWorkspace.run``). +# Program names of the dispatch table (the K1-K4 programs of the design note §14) and the local chunk stages: the +# existing backend's stages at the per-rank shape (launched with the records its plan selects on +# :data:`LOCAL_PATH_FWD` / :data:`LOCAL_PATH_BWD`) and the two KCP-owned records of :data:`KCP_LOCAL_RECORDS`. KCP_PROGRAMS = ( "kcp_fwdh_offset", # K1 ``_offset``: walk from zero, final state -> compact tiles (B_f / b_seg) "kcp_fwdh_tape", # K1 ``_tape``: walk from an indexed initial state, h / v_new tapes @@ -220,8 +319,30 @@ "kcp_transfer_walk", # K3: FP32-class transfer M_f (^T) (or the FP32-class offset with init_mode = 1) "kcp_affine_merge", # K4: exact FP32 affine chain (state / map / transposed modes) ) -LOCAL_STAGES_FWD = ("fwd_prep", "fwd_intra", "wy", "fwd_o") -LOCAL_STAGES_BWD = ("bwd_prep", "wy", "dav", "dqkg", "glm_intra", "gate_epilogue") +LOCAL_STAGES_FWD = ("fwd_intra", "wy_fwd", "fwd_o") +LOCAL_STAGES_BWD = ( + "bwd_prep", + "fwd_intra", + "wy", + "dav", + "dqkg", + "glm_intra", + "gate_epilogue", +) +# The paths of the existing backend's plan the local stages take their records from (the forward without a saved +# set, the checkpoint replay for both backward modes: :mod:`.cake_kcp_plan`) and the arena layouts they run on. +LOCAL_PATH_FWD = cake_kcp_plan.LOCAL_PATH_FWD +LOCAL_PATH_BWD = cake_kcp_plan.LOCAL_PATH_BWD +LOCAL_LAYOUT_FWD = "fwd" +LOCAL_LAYOUT_BWD = "replay" +# The KCP-owned local records (:mod:`.cake_kcp_jit`), by the stage name the launch sequences use: the output program +# over the K1 ``h`` / ``v_new`` tapes and the ``dA_qk`` / ``dv1`` program whose ``dv1`` K2 reads (the forms the +# existing backend's output and dA / dv programs had before its fused recurrences took them over). +KCP_LOCAL_RECORDS = {"fwd_o": "kcp_fwd_o", "dav": "kcp_dav"} +# The WY records the KCP launches directly where the plan fuses the WY representation into a state recurrence the +# KCP does not run (its recurrences are the K1 / K2 walks): the forward's program and the replay's. +WY_DIRECT_RECORDS = {"wy_fwd": "wy_fwd", "wy": "wy_bwd_replay"} +SCRATCH_TABLES_ATTR = "_kcp_scratch_tables" # the arena attribute holding the table set the intra record writes COLLECTIVES = ("c1", "c2", "b1", "b2") @@ -389,7 +510,7 @@ def transpose_tiles( "w", "u", ) # the prep rows SPEC-3 compares (the design note §6.4) -SPEC3_ENV = "CAKE_KCP_SPEC3" # "0": skip the debug assertion; "warn": record a failure without raising; default: assert +SPEC3_ENV = "CAKE_KCP_SPEC3" # "0": skip the probe; default: record; "warn": warn on a difference; "assert": raise _PROGRAM_OF = { "kcp_fwdh_offset": "fwdh_offset", "kcp_fwdh_tape": "fwdh_tape", @@ -401,9 +522,14 @@ def transpose_tiles( class KcpPreconditionError(RuntimeError): - """SPEC-3 failed: ``bwd_prep`` + ``wy`` did not reproduce the forward's ``gk`` / ``kg`` / ``qg`` / ``w`` / ``u`` rows - bitwise (the design note §6.4); saved mode and replay mode would differ. Options (a) / (b) / (c) of §6.4 are a parent - decision (``NEED_PARENT``), never a silent relaxation.""" + """The SPEC-3 probe found the backward's ``gk`` / ``kg`` / ``qg`` / ``w`` / ``u`` rows different from the forward's + on a real token row and ``CAKE_KCP_SPEC3=assert`` asked for an error (the design note §6.4). The requirement of + record is that the two backward modes (saved set, checkpoint replay) produce identical gradients: both run the + same prep, intra and WY programs on the same operands, which the GPU tests check bitwise. The forward's and the + backward's ``gk`` buffers differ by construction on the current program set: the forward's intra program + (``write_prep = 1``) writes the chunk-final gate row alone (row 63 of every chunk: the chunk's total gate in log2 + units; 0 on pad chunks) and the backward's prep program writes every token row's prefix, whose row 63 is the same + total; the K1 / K2 / K3 walks read row 63 only. ``kg`` / ``qg`` / ``w`` / ``u`` agree bitwise.""" @dataclasses.dataclass(frozen=True) @@ -418,6 +544,10 @@ class _Block: stride: int +TEMPLATE_CACHE_ATTR = "_kcp_local_templates" # the arena attribute holding the trimmed path's begin_call templates +# {(id(table set), backward): (weakref.ref(table set), template)}; the views belong to that arena, so the cache dies with it + + class KernelBackend(Backend): """The shipped backend: the KCP programs K1-K4 through a program provider (``launch(program, args, *, grid, device_index, precision_mode)`` with the positional ABI of the design note §14) and the local @@ -425,18 +555,34 @@ class KernelBackend(Backend): registry (:mod:`.cake_kcp_jit`, imported at the first launch so this module imports without it); a bring-up provider launching the source builds directly can be passed (``path="direct"``; never a gated row). - Local stages: a forward call runs ``fwd_prep`` / ``fwd_intra`` / ``wy`` / ``fwd_o`` on the arena's ``fwd`` layout; - a backward call runs on the ``bwd`` layout — saved mode ``bwd_prep`` (``copy_saved = 1`` from the saved ``Aqk`` / - ``Akk``) + ``wy``, replay mode ``bwd_prep`` (``copy_saved = 0``) + ``fwd_intra`` (``write_int``) + ``wy`` as the - existing replay path (``cake_backend.chunk_kda_bwd``) — then ``dav`` / ``dqkg`` / ``glm_intra`` / ``gate_epilogue``. - The K1 / K2 / K3 walks read the layout's ``w`` / ``kg`` / ``u`` / ``gk`` (``qg`` / ``do_int`` / ``dv1``) rows and write - its ``h`` / ``v_new`` (``dh`` / ``dv2``) tapes. - - SPEC-3 (debug assertion, once per process on the first saved-mode backward): before that backward's ``bwd_prep`` - the forward's ``fwd_prep`` + ``fwd_intra`` + ``wy`` run on the ``fwd`` layout and their five rows are kept; after - ``bwd_prep`` + ``wy`` the ``bwd`` layout's rows must equal them bitwise on every real token row (pad rows are - reported separately). ``CAKE_KCP_SPEC3=0`` skips it, ``=warn`` records without raising; the result is in - ``LAST_CALL["bwd"]["spec3"]`` and :attr:`spec3`. + Local stages (:data:`LOCAL_STAGES_FWD` / :data:`LOCAL_STAGES_BWD`): a forward call runs the existing backend's + intra program (``write_prep = 1``: the GLM prep's rows and the internal ``Aqk`` / ``Akk`` rows; the token-row + copies are never written, ``write_ext = 0``), its forward WY program and the KCP's output program ``kcp_fwd_o`` + (the output from the K1 ``h`` / ``v_new`` tapes) on the arena's ``fwd`` layout, with the records the plan selects + on the forward path without a saved set (:data:`LOCAL_PATH_FWD`). A backward call runs on the ``replay`` layout + with the records of the replay path (:data:`LOCAL_PATH_BWD`) in both modes: the replay prep (the GLM prep's rows + and the internal ``dout`` copy ``do_int``), the intra program (``write_prep = 0``: the internal ``Aqk`` / ``Akk`` + rows), the replay's WY program, then after the K1 tape replay the KCP's ``kcp_dav`` (``dA_qk`` and the intra-chunk + ``dv1`` K2 reads) and after K2 the fused gradient, GLM intra and the gate epilogue. Where the plan fuses the WY + representation into a state recurrence the KCP does not run, the WY program is launched directly + (:data:`WY_DIRECT_RECORDS`) on its grid rule. The intra record that computes the chunk layout from ``cu_seqlens`` + writes its tables into a scratch set on the arena (:data:`SCRATCH_TABLES_ATTR`); every other program reads the + context's CPU-built tables. The saved set carries no token-row ``Aqk`` / ``Akk`` (:class:`KcpSaved`): the + backward recomputes the internal rows in both modes, so the two modes run the same programs on the same operands. + The K1 / K2 / K3 walks read the layout's ``w`` / ``kg`` / ``u`` / ``gk`` (``qg`` / ``do_int`` / ``dv1``) rows and + write its ``h`` / ``v_new`` (``dh`` / ``dv2``) tapes. + + SPEC-3 (debug probe, once per process on the first saved-mode backward): before that backward's prep the + forward's intra program and WY program run on the ``fwd`` layout and their five rows (``gk`` / ``kg`` / ``qg`` / + ``w`` / ``u``) are kept; after the backward's WY program the ``replay`` layout's rows are compared with them on + every real token row (pad rows are reported separately). The requirement of record is that the two backward + modes produce identical gradients (the GPU tests check saved mode against the replay bitwise). The probe's + ``equal`` is ``False`` by construction: the forward's intra program writes ``gk`` at row 63 of every chunk alone + (the chunk-final gate row, the only ``gk`` row the K1 / K2 / K3 walks read), the backward's prep program writes + every row's prefix with the same row 63, so ``rows["gk"]["real_rows_differ"]`` is the real rows minus the full + chunks while ``kg`` / ``qg`` / ``w`` / ``u`` report 0; hence the probe records by default. ``CAKE_KCP_SPEC3=0`` + skips it, ``=warn`` warns on a difference, ``=assert`` raises + :class:`KcpPreconditionError`; the result is in ``LAST_CALL["bwd"]["spec3"]`` and :attr:`spec3`. """ name = "kernels" @@ -494,11 +640,47 @@ def begin_call(self, call: KcpCall) -> None: raise ValueError( f"local plan geometry (NC_pad, R) = {(p.NC_pad, p.R)} differs from the KCP context's {(ctx.NC_pad, ctx.R)}" ) - views = ws.views("bwd" if call.backward else "fwd") + views = ws.views(LOCAL_LAYOUT_BWD if call.backward else LOCAL_LAYOUT_FWD) + if call.host_trim: # the call-independent part of the value dict is built once per (arena, tables, direction) + # the template is cached on the ARENA whose views it holds (it dies with the arena), keyed by the + # table set's identity with a WEAK reference to it (a hit needs that very table set alive); the backend + # keeps no reference to an arena or a table set; dead entries are purged on the next miss. + cache = ws.__dict__.get(TEMPLATE_CACHE_ATTR) + if cache is None: + cache = {} + setattr(ws, TEMPLATE_CACHE_ATTR, cache) + key = (id(call.tb), call.backward) + hit = cache.get(key) + if hit is None or hit[0]() is not call.tb: + for stale in [k for k, h in cache.items() if h[0]() is None]: + del cache[stale] + hit = ( + weakref.ref(call.tb), + self._local_template(views, ctx, p, call.backward), + ) + cache[key] = hit + values: dict[str, Any] = dict(hit[1]) + values.update(call.ops) + values["scale"] = call.scale + call.local_values = values + call.scratch["views"] = views + return + values = self._local_template(views, ctx, p, call.backward) + values.update(call.ops) + values["scale"] = call.scale + call.local_values = values + call.scratch["views"] = views + + @staticmethod + def _local_template(views: Any, ctx: Any, p: Any, backward: bool) -> dict[str, Any]: + """The values of ``begin_call`` that do not depend on the call: the arena views of the layout, the context's + four chunk tables (the existing backend's vocabulary; the KCP's segment table is a K1-K4 argument, not a + local one), the plan scalars of the existing backend's logical vocabulary and the ``h2d`` / ``dh2d`` views. + The operands and ``scale`` are added per call (the trimmed host path caches this part per (arena, table set, + direction)).""" values: dict[str, Any] = dict(vars(views)) for name in cake_backend.TABLE_NAMES[:4]: values[name] = ctx.tables[name] - values["seg_chunk_start"] = call.tb.hp.device_tables["seg_chunk_start"] values.update( num_seqs=p.N, num_chunks_pad=p.NC_pad, @@ -506,82 +688,95 @@ def begin_call(self, call: KcpCall) -> None: group=1, num_segments=p.NSEG, segments=p.P, + seg_arg=p.seg_arg, + proportional=p.proportional, + num_pairs=p.R // cake_jit.PAIR_ROWS, ) - values.update(call.ops) - values.update(h2d=views.h.view(-1, D), scale=call.scale) - if call.backward: + values["h2d"] = views.h.view(-1, D) + if backward: values["dh2d"] = views.dh.view(-1, D) - call.local_values = values - call.scratch["views"] = views + return values def stage(self, name: str, call: KcpCall, **extra: Any) -> None: + """One local stage of :data:`LOCAL_STAGES_FWD` / :data:`LOCAL_STAGES_BWD`: the plan's record on the call's path + (or the direct WY record), the KCP-owned record of ``fwd_o`` / ``dav``.""" values, views = call.local_values, call.scratch["views"] + local = call.local + p = local.plan + path = LOCAL_PATH_BWD if call.backward else LOCAL_PATH_FWD if call.out is not None: values["out"] = call.out if call.grads is not None: values.update(call.grads) - if ( - name == "fwd_prep" and call.backward - ): # replay: the backward layout's prep without the saved copies - name, extra = "bwd_prep", {"saved": None} if name == "bwd_prep": - saved = extra.pop("saved", None) - if saved is None: - extra.update(aqk_src=views.akk_int, akk_src=views.akk_int, copy_saved=0) - else: - if saved.Aqk is None or saved.Akk is None: - raise ValueError( - "saved: the saved set has no Aqk / Akk rows (forward without save_for_backward?)" - ) - _check( - saved.Aqk, - "saved.Aqk", - (call.T_r, call.H, BT), - torch.bfloat16, - call.device, - ) - _check( - saved.Akk, - "saved.Akk", - (call.T_r, call.H, BT), - torch.bfloat16, - call.device, - ) - extra.update(aqk_src=saved.Aqk, akk_src=saved.Akk, copy_saved=1) - if self._spec3_pending: - self._spec3_forward_rows( - call - ) # before the backward's prep: the layouts share the arena + # both backward modes run the replay prep (the GLM prep's rows, the internal dout copy K2 and dav read); + # the saved set carries no token rows: the intra program recomputes the internal rows in both modes + if extra.pop("saved", None) is not None and self._spec3_pending: + self._spec3_forward_rows( + call + ) # before the backward's prep: the layouts share the arena elif name == "fwd_intra": - if int(extra.get("write_ext", 0)): - aqk = call.scratch["Aqk"] = _alloc( - (call.T_r, call.H, BT), torch.bfloat16, call.device - ) - akk = call.scratch["Akk"] = _alloc( - (call.T_r, call.H, BT), torch.bfloat16, call.device + # the token-row Aqk / Akk are never written on this route (the saved set carries none); any BF16 tensor + # serves as the unused pointer (as cake_backend does) + extra.update(write_ext=0, aqk_ext=views.akk_int, akk_ext=views.akk_int) + # the internal rows every KCP consumer reads; the GLM prep's rows on the forward, the backward's prep + # program wrote them already + extra.setdefault("write_int", 1) + extra.setdefault("write_prep", 0 if call.backward else 1) + if cake_kcp_plan.intra_writes_tables(p, path): + extra.update( + cu_seqlens=call.ctx.cu_seqlens_local, **_scratch_tables(local) ) - extra.update(aqk_ext=aqk, akk_ext=akk) - else: # not written: any BF16 tensor serves as the unused pointer (as cake_backend does) - extra.update(aqk_ext=views.akk_int, akk_ext=views.akk_int) - call.local.run(name, {**values, **extra}) + elif name in KCP_LOCAL_RECORDS: + self.programs().run_record( + KCP_LOCAL_RECORDS[name], + {**values, **extra}, + grid=kcp_local_grid(name, p), + device_index=local.device_index, + ) + return + record, grid = _local_record(p, name, path) + self._launch(call, record, grid, {**values, **extra}) if name == "wy" and call.scratch.get("spec3_ref") is not None: self._spec3_compare(call) + @staticmethod + def _launch( + call: KcpCall, + record: str, + grid: tuple[int, int, int], + values: dict[str, Any], + ) -> None: + """A record of the existing backend's registry on the local arena: the trimmed launch, or the arena's + launcher (a diagnostics launcher keeps the base path).""" + local = call.local + if call.host_trim and local.launcher is cake_jit.run_module: + _launch_record_trim(local, record, grid, values) + else: + local.launcher(record, values, grid=grid, device_index=local.device_index) + # -- SPEC-3 def _spec3_forward_rows(self, call: KcpCall) -> None: + """The probe's reference: the forward's intra program (``write_prep = 1``) and WY program on the ``fwd`` + layout, their five rows kept (the backward's layout shares the arena, so this runs before its prep).""" self._spec3_pending = False - ws = call.local - a = ws.views("fwd") + local = call.local + p = local.plan + a = local.views(LOCAL_LAYOUT_FWD) vals = {**call.local_values, **vars(a)} vals.update( h2d=a.h.view(-1, D), aqk_ext=a.akk_int, akk_ext=a.akk_int, - out=call.grads["dv"], + write_ext=0, + write_int=1, + write_prep=1, ) - ws.run("fwd_prep", vals) - ws.run("fwd_intra", {**vals, "write_ext": 0, "write_int": 1}) - ws.run("wy", vals) + if cake_kcp_plan.intra_writes_tables(p, LOCAL_PATH_FWD): + vals.update(cu_seqlens=call.ctx.cu_seqlens_local, **_scratch_tables(local)) + for stage in ("fwd_intra", "wy_fwd"): + record, grid = _local_record(p, stage, LOCAL_PATH_FWD) + self._launch(call, record, grid, vals) call.scratch["spec3_ref"] = {n: getattr(a, n).clone() for n in _SPEC3_ROWS} def _spec3_compare(self, call: KcpCall) -> None: @@ -608,10 +803,16 @@ def _spec3_compare(self, call: KcpCall) -> None: if n_real: result["equal"] = False self.spec3 = call.scratch["spec3"] = result - if not result["equal"] and os.environ.get(SPEC3_ENV, "1") != "warn": - raise KcpPreconditionError( - f"SPEC-3: bwd_prep + wy differ from fwd_prep + wy on real rows: {result['rows']}" + if not result["equal"]: + mode = os.environ.get(SPEC3_ENV, "1") + message = ( + "SPEC-3 probe: the backward's prep + WY rows differ from the forward's intra + WY rows on real " + f"token rows: {result['rows']}" ) + if mode == "assert": + raise KcpPreconditionError(message) + if mode == "warn": + warnings.warn(message, RuntimeWarning, stacklevel=2) # -- KCP programs (positional ABI of the design note §14) def abi( @@ -643,6 +844,8 @@ def abi( def run(self, name: str, args: dict[str, Any]) -> None: call: KcpCall = args["_call"] + if call.host_trim and self._run_trim(name, args, call): + return prog, abi, grid, variant = self.abi(name, args) self.programs().launch( prog, @@ -654,6 +857,50 @@ def run(self, name: str, args: dict[str, Any]) -> None: **variant, ) + def _run_trim(self, name: str, args: dict[str, Any], call: KcpCall) -> bool: + """Trimmed host path: the launch without :meth:`abi` when the provider launches from the abstract ``args`` anyway (the + registry's ``run_program`` form discards the positional ABI, whose walk / index / chain device tables cost 100-300 us of + host time per launch). Keeps the D1 bookkeeping of the merge (alignment check, ``cols_per_cta``, ``merge_log``). + Returns False when the base path is needed (another provider, an explicit ``_c`` merge record).""" + programs = self.programs() + jit = getattr(programs, "jit", None) + if ( + not isinstance(programs, _FiPrograms) + or jit is None + or hasattr(jit, "launch_kcp") + or not hasattr(jit, "run_program") + ): + return False + prog = _PROGRAM_OF[name] + variant: dict[str, Any] = {} + if prog == "affine_merge": + ch: Chains = args["chains"] + cols = merge_cols_per_cta(ch.ncols, ch.n_chains, call.cols_rule) + if cols is not None: + return False # the explicit record form needs the positional ABI (the base path appends the merge_log entry) + if id(ch) not in self._checked_chains: + ch.check_alignment() + self._checked_chains.add(id(ch)) + args["cols_per_cta"] = None + call.merge_log.append( + (str(args.get("family")), int(ch.ncols), int(ch.n_chains), None) + ) + elif prog == "transfer_walk": + variant = { + "precision_mode": call.precision_mode, + "transfer_walk": call.transfer_walk, + } + programs.launch( + prog, + [], + grid=0, + device_index=call.device.index, + name=name, + args=args, + **variant, + ) + return True + def _walk_common( self, call: KcpCall, args: dict[str, Any] ) -> tuple[torch.Tensor, torch.Tensor, int, Any]: @@ -953,7 +1200,21 @@ def merge_cols_per_cta(ncols: int, n_chains: int, rule: str) -> int | None: "none": 0, "unavailable": 0, "capture_skips": 0, -} # calls recorded per cache-config form; capture_skips = capped launches issued bare inside a capture + "node_attr": 0, + "node_attr_post": 0, + "node_attr_launches": 0, + "node_attr_degraded": 0, +} # calls recorded per cache-config form; capture_skips = capped launches issued bare inside a capture; node_attr_launches = captured tape launches whose kernel node got the carveout attribute (the kernel-node cap) +_CUDART_CAPTURE_INFO_NAMES = ( + "cudaStreamGetCaptureInfo_v3", + "cudaStreamGetCaptureInfo_v2", + "cudaStreamGetCaptureInfo", +) # export names of the stream-capture-info call: _v3 (CUDA 12.3+, 7 arguments incl. edgeData_out), _v2 (CUDA 11.3-12.x, 6 arguments), plain (CUDA 13 (13000+) = the 7-argument form; older runtimes = the pre-11.3 3-argument form, unusable) +_CUDART_GRAPH_ENTRY_POINTS = ( + "cudaGraphNodeGetType", + "cudaGraphKernelNodeSetAttribute", + "cudaGraphKernelNodeGetAttribute", +) # with the capture-info call: the capture-graph entry points of the kernel-node cap forms def _cudart() -> Any: @@ -996,11 +1257,374 @@ def _cudart_available() -> Any | None: return _CUDART_CACHE["rt"] +def _capture_info_name(rt: Any) -> str | None: + """The export name of the stream-capture-info call on runtime ``rt`` (:data:`_CUDART_CAPTURE_INFO_NAMES`, first match), or + ``None``.""" + return next( + (name for name in _CUDART_CAPTURE_INFO_NAMES if hasattr(rt, name)), None + ) + + +def _runtime_version(rt: Any) -> int: + """``cudaRuntimeGetVersion`` of ``rt`` (0 when the call is missing or fails).""" + fn = getattr(rt, "cudaRuntimeGetVersion", None) + if fn is None: + return 0 + ver = ctypes.c_int(0) + try: + rc = int(fn(ctypes.byref(ver))) + except (TypeError, ctypes.ArgumentError): + return 0 + return int(ver.value) if rc == 0 else 0 + + +def _capture_info_arity(rt: Any, name: str) -> int: + """Arguments of the stream-capture-info export ``name`` on ``rt``: 7 for ``_v3`` (``..., dependencies_out, edgeData_out, + numDependencies_out``), 6 for ``_v2``, and for the plain name 7 on a CUDA 13 runtime (13000+; it exports the v3 form) or 0 + (unusable pre-11.3 form) on older runtimes.""" + if name == "cudaStreamGetCaptureInfo_v3": + return 7 + if name == "cudaStreamGetCaptureInfo_v2": + return 6 + return 7 if _runtime_version(rt) >= 13000 else 0 + + +def graph_entry_points() -> dict[str, Any]: + """The ctypes entry points of the ``node_attr`` forms as resolved in this process (the evidence record): the + capture-info export name (or ``None``), the fixed entry points, and whether the typed runtime is available.""" + rt = _cudart_graph_available() + return { + "capture_info": _CUDART_CACHE.get("capture_info_name"), + "capture_info_arity": _CUDART_CACHE.get("capture_info_arity"), + "runtime_version": _CUDART_CACHE.get("runtime_version"), + "entry_points": list(_CUDART_GRAPH_ENTRY_POINTS), + "available": rt is not None, + "cache_config_entry_points": [ + "cudaDeviceSetCacheConfig", + "cudaDeviceGetCacheConfig", + "cudaGetLastError", + "cudaRuntimeGetVersion", + ], + } + + +def _cudart_graph() -> Any: + """The CUDA runtime of :func:`_cudart_available` with the capture-graph entry points of the ``node_attr`` forms typed + the stream-capture-info call (``cudaStreamGetCaptureInfo`` on CUDA 13, ``..._v2`` on CUDA 12: + the capture status / graph of a stream and the nodes the next captured operation would depend on), ``cudaGraphNodeGetType``, + ``cudaGraphKernelNodeSetAttribute`` / ``...GetAttribute`` (a kernel node's launch attribute). Raises ``RuntimeError`` when + the runtime or one of the entry points is missing.""" + rt = _cudart_available() + if rt is None: + raise RuntimeError( + "graph_cap node_attr: no CUDA runtime is loadable in this process" + ) + if hasattr(rt, "cudaRuntimeGetVersion") and isinstance(rt, ctypes.CDLL): + rt.cudaRuntimeGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)] + rt.cudaRuntimeGetVersion.restype = ctypes.c_int + info_name = _capture_info_name(rt) + missing = ( + [" | ".join(_CUDART_CAPTURE_INFO_NAMES)] if info_name is None else [] + ) + [name for name in _CUDART_GRAPH_ENTRY_POINTS if not hasattr(rt, name)] + if missing: + raise RuntimeError( + f"graph_cap node_attr: the CUDA runtime exports no {missing}" + ) + arity = _capture_info_arity(rt, info_name) + if arity == 0: + raise RuntimeError( + f"graph_cap node_attr: {info_name} of this runtime (version {_runtime_version(rt)}) has the pre-11.3 form without dependency nodes" + ) + info = getattr(rt, info_name) + deps_t = ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p)) + info.argtypes = ( + [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_ulonglong), + ctypes.POINTER(ctypes.c_void_p), + deps_t, + ] + + ([deps_t] if arity == 7 else []) + + [ctypes.POINTER(ctypes.c_size_t)] + ) + rt.cudaGraphNodeGetType.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int)] + rt.cudaGraphKernelNodeSetAttribute.argtypes = [ + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ] + rt.cudaGraphKernelNodeGetAttribute.argtypes = [ + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ] + for fn in ( + info, + rt.cudaGraphNodeGetType, + rt.cudaGraphKernelNodeSetAttribute, + rt.cudaGraphKernelNodeGetAttribute, + ): + fn.restype = ctypes.c_int + _CUDART_CACHE["capture_info_name"] = info_name + _CUDART_CACHE["capture_info_arity"] = arity + _CUDART_CACHE["runtime_version"] = _runtime_version(rt) + return rt + + +def _cudart_graph_available() -> Any | None: + """:func:`_cudart_graph`'s runtime, or ``None`` (one ``RuntimeWarning``) when the runtime or an entry point is missing: the + ``node_attr`` forms then capture without the cap and the call records ``cache_config="unavailable"``.""" + if "graph" not in _CUDART_CACHE: + try: + _CUDART_CACHE["graph"] = _cudart_graph() + except RuntimeError as exc: + _CUDART_CACHE["graph"] = None + warnings.warn( + f"{exc}: the graph_cap node_attr forms capture the tape programs without the cap (cache_config='unavailable')", + RuntimeWarning, + stacklevel=2, + ) + return _CUDART_CACHE["graph"] + + +def capture_kernel_node(device_index: int) -> dict[str, Any]: + """The kernel node just captured on the current stream of ``device_index`` (the kernel-node cap): the stream-capture-info call must + report an active capture whose next operation depends on exactly one node, and that node must be a kernel node (the launch + issued right before this call). Returns the evidence record with the node handle; raises ``RuntimeError`` otherwise (:meth:`_FiPrograms.cap_captured_node` + catches it: the launch stays uncapped and the capture goes on).""" + rt = _cudart_graph_available() + if rt is None: + raise RuntimeError( + "capture_kernel_node: the capture-graph entry points are unavailable" + ) + info_name = _capture_info_name(rt) + if info_name is None: + raise RuntimeError( + f"capture_kernel_node: the runtime exports none of {_CUDART_CAPTURE_INFO_NAMES}" + ) + arity = _CUDART_CACHE.get("capture_info_arity") or _capture_info_arity( + rt, info_name + ) + if arity == 0: + raise RuntimeError( + f"capture_kernel_node: {info_name} has the pre-11.3 form without dependency nodes" + ) + info = getattr(rt, info_name) + stream = int(torch.cuda.current_stream(device_index).cuda_stream) + status, capture_id, graph = ( + ctypes.c_int(-1), + ctypes.c_ulonglong(0), + ctypes.c_void_p(), + ) + deps, edges, n_deps = ( + ctypes.POINTER(ctypes.c_void_p)(), + ctypes.POINTER(ctypes.c_void_p)(), + ctypes.c_size_t(0), + ) + args = [ + ctypes.c_void_p(stream), + ctypes.byref(status), + ctypes.byref(capture_id), + ctypes.byref(graph), + ctypes.byref(deps), + ] + if arity == 7: + args.append( + ctypes.byref(edges) + ) # const cudaGraphEdgeData **edgeData_out (unused: plain edges) + args.append(ctypes.byref(n_deps)) + rc = int(info(*args)) + rec = { + "device_index": int(device_index), + "stream": stream, + "capture_info": info_name, + "arity": arity, + "rc_info": rc, + "status": int(status.value), + "capture_id": int(capture_id.value), + "graph": graph.value, + "n_deps": int(n_deps.value), + } + if rc != 0: + rt.cudaGetLastError() + raise RuntimeError(f"{info_name} failed with CUDA error {rc}") + if status.value != _STREAM_CAPTURE_ACTIVE or n_deps.value != 1: + raise RuntimeError( + f"capture_kernel_node: expected one dependency node of an active capture, got status {status.value} and {n_deps.value} nodes" + ) + node = deps[0] + node_type = ctypes.c_int(-1) + rc = int(rt.cudaGraphNodeGetType(ctypes.c_void_p(node), ctypes.byref(node_type))) + if rc != 0: + rt.cudaGetLastError() + raise RuntimeError(f"cudaGraphNodeGetType failed with CUDA error {rc}") + if node_type.value != _GRAPH_NODE_TYPE_KERNEL: + raise RuntimeError( + f"capture_kernel_node: the just-captured node has type {node_type.value}, not a kernel node" + ) + rec.update(node=int(node), node_type=int(node_type.value)) + return rec + + +def set_kernel_node_carveout( + node: int, carveout: int = GRAPH_NODE_CARVEOUT +) -> dict[str, Any]: + """Set the preferred shared-memory carveout launch attribute of kernel node ``node`` (the kernel-node cap) and read it back; returns + the evidence record (``rc`` of the set, the read-back value). ``cudaLaunchAttributeValue`` is a 64-byte union whose first + unsigned int is ``sharedMemCarveout``.""" + rt = _cudart_graph_available() + if rt is None: + raise RuntimeError( + "set_kernel_node_carveout: the capture-graph entry points are unavailable" + ) + value = (ctypes.c_uint * 16)() + value[0] = int(carveout) & 0xFFFFFFFF + rc = int( + rt.cudaGraphKernelNodeSetAttribute( + ctypes.c_void_p(int(node)), _LAUNCH_ATTR_SHMEM_CARVEOUT, ctypes.byref(value) + ) + ) + if rc != 0: + rt.cudaGetLastError() + raise RuntimeError( + f"cudaGraphKernelNodeSetAttribute(PreferredSharedMemoryCarveout={carveout}) failed with CUDA error {rc}" + ) + got = (ctypes.c_uint * 16)() + got[0] = 0xFFFFFFFE + rc_get = int( + rt.cudaGraphKernelNodeGetAttribute( + ctypes.c_void_p(int(node)), _LAUNCH_ATTR_SHMEM_CARVEOUT, ctypes.byref(got) + ) + ) + if rc_get != 0: + rt.cudaGetLastError() + return { + "node": int(node), + "set": int(carveout), + "rc": rc, + "rc_get": rc_get, + "readback": ctypes.c_int(got[0]).value if rc_get == 0 else None, + } + + def _capturing() -> bool: """``True`` inside a CUDA stream capture of the current thread (the bracket is never issued there).""" return torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() +# -- the trimmed host path's FFI stream context. tvm_ffi.use_torch_stream() builds a torch.cuda.Stream and +# parses its device string on every launch (~20 us); use_raw_stream with the current raw stream handle of the device sets the +# same environment stream (the kernel runs on the same torch stream: bitwise the untrimmed launch). +_FFI_DEVICES: dict[int, Any] = {} + + +def _ffi_device(device_index: int) -> Any: + dev = _FFI_DEVICES.get(int(device_index)) + if dev is None: + import tvm_ffi # noqa: PLC0415 - the launch path's dependency, imported at the first trimmed launch + + dev = _FFI_DEVICES[int(device_index)] = tvm_ffi.device( + f"cuda:{int(device_index)}" + ) + return dev + + +def _raw_stream(device_index: int) -> int: + get = getattr(torch._C, "_cuda_getCurrentRawStream", None) + if get is not None: + return int(get(int(device_index))) + return int(torch.cuda.current_stream(int(device_index)).cuda_stream) + + +def _ffi_stream(device_index: int) -> Any: + """The FFI stream context of the trimmed launch: the device's current torch stream, as a raw handle.""" + import tvm_ffi # noqa: PLC0415 + + return tvm_ffi.use_raw_stream(_ffi_device(device_index), _raw_stream(device_index)) + + +_ARCH_BY_INDEX: dict[int, str] = {} + + +def _arch_of(device_index: int) -> str: + """``cake_jit.device_arch`` cached per device index (a device's capability never changes within a process).""" + arch = _ARCH_BY_INDEX.get(int(device_index)) + if arch is None: + arch = _ARCH_BY_INDEX[int(device_index)] = cake_jit.device_arch( + int(device_index) + ) + return arch + + +def _local_record(p: Any, stage: str, path: str) -> tuple[str, tuple[int, int, int]]: + """``(record, grid)`` of local stage ``stage`` on ``path`` of plan ``p``: the plan's selection, or the direct WY + record (:data:`WY_DIRECT_RECORDS`) on its own grid rule where the plan fuses the WY representation into a state + recurrence the KCP does not run.""" + if stage in WY_DIRECT_RECORDS and not p.dispatched(stage, path): + record = WY_DIRECT_RECORDS[stage] + grid = cake_jit.launch_grid( + record, + num_chunks_pad=p.NC_pad, + num_seqs=p.N, + num_heads=p.H, + sm_count=p.sm_count, + ) + return record, grid + return p.record(stage, path), p.grid(stage, path) + + +def kcp_local_grid(stage: str, p: Any) -> tuple[int, int, int]: + """The launch grid of the KCP-owned local record of ``stage`` (:data:`KCP_LOCAL_RECORDS`) from the plan's + scalars: the output program is persistent over the (chunk, head) items, ``min(NC_pad * H, sm_count)`` CTAs; the + ``dA_qk`` / ``dv1`` program launches one CTA per chunk pair and head, ``(NC_pad / 2, H, 1)``.""" + if stage == "fwd_o": + return (max(1, min(int(p.NC_pad) * int(p.H), int(p.sm_count))), 1, 1) + if stage == "dav": + return (int(p.NC_pad) // 2, int(p.H), 1) + raise ValueError(f"no KCP-owned local record for stage {stage!r}") + + +def _scratch_tables(local: Any) -> dict[str, torch.Tensor]: + """The table set the intra record that computes the chunk layout from ``cu_seqlens`` writes (the existing + backend's tables, its block table included), on a scratch buffer held by the arena: the KCP's CPU-built tables + are never rewritten by a kernel and no KCP program reads this set (:func:`cake_kcp_plan.local_scratch_table_bytes` + counts it).""" + tables = local.__dict__.get(SCRATCH_TABLES_ATTR) + if tables is None: + p = local.plan + buffer = torch.zeros((p.table_words(),), dtype=torch.int32, device=local.device) + tables = { + name: buffer[off : off + words] for name, off, words in p.table_layout() + } + setattr(local, SCRATCH_TABLES_ATTR, tables) + return tables + + +def _launch_record_trim( + local: Any, name: str, grid: tuple[int, int, int], values: dict[str, Any] +) -> None: + """Trimmed host path: ``cake_jit.run_module(name, values, grid=grid, device_index=local.device_index)`` with the + launch inlined -- the record's argument plan (cached by cake_jit), the arena's cached arch, the raw-stream FFI + context. Same record, same grid, same positional arguments, same stream as ``local.run``.""" + jit = cake_jit + record = jit.MODULES[name] + args = [] + for logical_name, axis in jit._argument_plan(name): + if logical_name is None: + args.append(int(grid[axis])) + else: + try: + args.append(values[logical_name]) + except KeyError: + raise ValueError( + f"record {name} needs logical argument {logical_name!r}" + ) from None + module = jit.load(name, local.arch) + with _ffi_stream(local.device_index): + getattr(module, record["ffi_entry"])(*args) + + def walk_device_cache_config(device_index: int, config: int | None) -> dict[str, Any]: """Set (``config`` not ``None``) and read back the device-wide cache-configuration preference of device ``device_index`` through the CUDA runtime; returns the evidence record (``rc`` of the set, the read-back value).""" @@ -1059,6 +1683,24 @@ def __init__(self, jit: Any) -> None: self.unavailable = ( 0 # capped launches issued without a loadable CUDA runtime (no bracket) ) + self.node_attrs = ( + 0 # captured tape launches whose kernel node got the carveout attribute + ) + self.node_attr_log: dict[ + tuple[str, int], list[dict[str, Any]] + ] = {} # (program, device) -> evidence of the first node attributes + + def run_record( + self, + record: str, + values: dict[str, Any], + *, + grid: tuple[int, int, int], + device_index: int, + ) -> None: + """A KCP-owned local record (:data:`KCP_LOCAL_RECORDS`) with logical ``values`` on ``grid``: the registry's + ``run_module`` (the record maps its kernel parameters onto the logical vocabulary).""" + self.jit.run_module(record, values, grid=grid, device_index=device_index) def set_cache_config(self, program: str, device_index: int, config: int) -> None: """One side of a capped launch's bracket. The first brackets of a (program, device) go through @@ -1107,6 +1749,27 @@ def launch( call, "scratch", None ) # the call's record of its cache-config form if _capturing(): # never inside a stream capture; the capture's preference is _graph_cap_bracket's + form = getattr( + getattr(call, "opts", None), "graph_cap", GRAPH_CAP_DEFAULT + ) + if ( + form in GRAPH_NODE_CAPS and _cudart_graph_available() is not None + ): # the cap as the node's launch attribute + self._launch( + program, + abi, + grid=grid, + device_index=device_index, + precision_mode=precision_mode, + transfer_walk=transfer_walk, + name=name, + args=args, + cols_per_cta=cols_per_cta, + ) + self.cap_captured_node( + program, device_index, call, immediate=(form == "node_attr") + ) + return self.capture_skips += 1 CACHE_CONFIG_STATS["capture_skips"] += 1 self._launch( @@ -1169,6 +1832,67 @@ def launch( cols_per_cta=cols_per_cta, ) + def cap_captured_node( + self, program: str, device_index: int, call: Any, *, immediate: bool + ) -> None: + """Kernel-node cap: the kernel node just captured for ``program`` gets the one-CTA-per-SM carveout attribute -- now + (``graph_cap="node_attr"``) or when the capture has ended (``"node_attr_post"``: the handle is kept on the call for + :meth:`apply_node_caps`). The first four records per (program, device) are kept as evidence (``node_attr_log``). + Degrade contract: when :func:`capture_kernel_node` cannot name the node (a missing entry point, a + capture-info reply without exactly one kernel-node dependency, a non-kernel node), the launch just issued stays uncapped + and the capture goes on -- counted as a capture skip and in ``CACHE_CONFIG_STATS["node_attr_degraded"]``, one + ``RuntimeWarning`` per process, the reason recorded as ``call.scratch["node_attr_degraded"]`` (the call then records + ``cache_config="unavailable"``); never an eager fallback. A failing set-attribute call keeps raising.""" + try: + info = capture_kernel_node(device_index) + except RuntimeError as exc: # the launch stays uncapped, the capture goes on (never an eager fallback) + self.capture_skips += 1 + CACHE_CONFIG_STATS["capture_skips"] += 1 + CACHE_CONFIG_STATS["node_attr_degraded"] += 1 + scratch = getattr(call, "scratch", None) + if scratch is not None: + scratch["node_attr_degraded"] = str(exc)[:200] + if not _CUDART_CACHE.get("node_attr_degraded_warned"): + _CUDART_CACHE["node_attr_degraded_warned"] = True + warnings.warn( + f"graph_cap node_attr: {exc}: the captured tape launch stays uncapped (the call records cache_config='unavailable')", + RuntimeWarning, + stacklevel=2, + ) + return + if immediate: + info.update(set_kernel_node_carveout(info["node"])) + self.node_attrs += 1 + CACHE_CONFIG_STATS["node_attr_launches"] += 1 + log = self.node_attr_log.setdefault((program, int(device_index)), []) + if len(log) < 4: + log.append({**info, "program": program, "immediate": bool(immediate)}) + scratch = getattr(call, "scratch", None) + if scratch is not None: + scratch["node_attr_launches"] = int( + scratch.get("node_attr_launches", 0) + ) + (1 if immediate else 0) + if not immediate: + scratch.setdefault("node_attr_pending", []).append( + (program, int(device_index), int(info["node"])) + ) + + def apply_node_caps(self, call: Any) -> int: + """``graph_cap="node_attr_post"``: set the carveout attribute on every node kept by :meth:`cap_captured_node` (after the + capture ended, before the graph is instantiated); returns the number of nodes set.""" + pending = list(call.scratch.pop("node_attr_pending", [])) + for program, device_index, node in pending: + rec = set_kernel_node_carveout(node) + self.node_attrs += 1 + CACHE_CONFIG_STATS["node_attr_launches"] += 1 + for entry in self.node_attr_log.get((program, int(device_index)), []): + if entry.get("node") == node: + entry.update(rec) + call.scratch["node_attr_launches"] = int( + call.scratch.get("node_attr_launches", 0) + ) + len(pending) + return len(pending) + def _launch( self, program: str, @@ -1245,6 +1969,18 @@ def get_backend() -> Backend: # Tables of one (ctx, H) problem, derived from the plan module's head plan (``ctx.heads(H)``) +@dataclasses.dataclass(frozen=True) +class _IndexedMeta: + """Trimmed host path: the precomputed part of an :class:`Indexed` (see :meth:`_Tables.indexed_meta`).""" + + index: tuple[int, ...] + tiles: tuple[tuple[int, ...], ...] + H: int + addr: tuple[int, int, int, int] + device: str + dev_index: torch.Tensor + + @dataclasses.dataclass(frozen=True) class Indexed: """A compact state pool as a walk kernel reads or writes it: ``pool`` (a flat tensor of ``128 x 128`` tiles), @@ -1255,6 +1991,7 @@ class Indexed: pool: torch.Tensor index: tuple[int, ...] tiles: tuple[tuple[int, ...], ...] + meta: Any = None # the tables' precomputed _IndexedMeta (addressing + device index table) of the trimmed host path @property def n_slots(self) -> int: @@ -1262,6 +1999,9 @@ def n_slots(self) -> int: def addressing(self, H: int) -> tuple[int, int, int, int]: """``(base_tile, hg, n, stride)`` of the kernels' addressing, ``ValueError`` when the table has no such form.""" + H = int(H) + if self.meta is not None and self.meta.H == H: + return self.meta.addr n_slots = len(self.tiles) if n_slots == 0: return 0, int(H), 0, 1 @@ -1332,6 +2072,55 @@ def check_alignment(self) -> None: f"Chains: {name} holds an offset that is not a multiple of 4 floats" ) + def device_tables(self, device: str) -> dict[str, torch.Tensor]: + """Trimmed host path: the six device tables of the K4 launch (``chain_start`` int32; ``elem_m`` / ``elem_c`` / ``elem_y`` / + ``x0_index`` / ``y_final`` int64; an empty element table holds one 0 so the kernel has a valid pointer -- the + same tensors :func:`cake_kcp_jit._chain_tables` builds), cached on the chain set per device. The tables are + frozen once the launches start (``_Tables`` builds them once per ``(ctx, H)``).""" + cache = self.__dict__.setdefault("_dev_tables", {}) + t = cache.get(device) + if t is None: + pad = lambda table: table if table else [0] # noqa: E731 + t = cache[device] = { + "chain_start": torch.tensor( + list(self.chain_start), dtype=torch.int32, device=device + ), + "elem_m": torch.tensor( + pad(list(self.elem_m)), dtype=torch.int64, device=device + ), + "elem_c": torch.tensor( + pad(list(self.elem_c)), dtype=torch.int64, device=device + ), + "elem_y": torch.tensor( + pad(list(self.elem_y)), dtype=torch.int64, device=device + ), + "x0_index": torch.tensor( + list(self.x0_index), dtype=torch.int64, device=device + ), + "y_final": torch.tensor( + list(self.y_final), dtype=torch.int64, device=device + ), + } + return t + + @classmethod + def concat(cls, *parts: Chains) -> Chains: + """The chains of ``parts`` in order as one launch table (same ``ncols``; element and chain tables appended, the + CSR starts shifted) -- the fused carry families of ``merge_fuse="carry"``. Every chain keeps its own elements, initial state and + final slot, so the launch computes exactly what the separate launches compute.""" + out = cls(parts[0].ncols if parts else D) + for c in parts: + if c.ncols != out.ncols: + raise ValueError(f"Chains.concat: ncols {c.ncols} != {out.ncols}") + base = len(out.elem_m) + out.elem_m += c.elem_m + out.elem_c += c.elem_c + out.elem_y += c.elem_y + out.chain_start += [base + s for s in c.chain_start[1:]] + out.x0_index += c.x0_index + out.y_final += c.y_final + return out + def as_dict(self) -> dict[str, Any]: return { "ncols": self.ncols, @@ -1386,7 +2175,7 @@ def __init__( self.H = H = int(H) self.fused_walks = int( fused_walks - ) # lever L-B: one base pointer per tape phase + ) # fused walks: one base pointer per tape phase if walk_geometry not in WALK_GEOMETRIES: raise ValueError( f"walk_geometry must be one of {WALK_GEOMETRIES}, got {walk_geometry!r}" @@ -1454,7 +2243,7 @@ def segs(fs, ps): self.walk_m_b1 = ( segs([f for f in ins if f.crossing_out], [0]) if P == 1 else self.walk_all ) # K3 M-only (b1) - # fused_walks (L-B): one launch per tape phase -- every live segment after wait(c2) / wait(b2) (schedule + # fused_walks: one launch per tape phase -- every live segment after wait(c2) / wait(b2) (schedule # wait_first), or the crossing-in (crossing-out) fragments' segments when the document-start (document-last) # walk stays under the collective (schedule overlap) self.walk_live = segs(starts + ins, range(P)) @@ -1479,10 +2268,10 @@ def lpt(lst): self.dcarry_index = tuple(int(i) for i in hp.dcarry_index) self.h0_index_fused = tuple( int(i) for i in hp.h0_index_fused - ) # L-B: the carry plane (P > 1) / S_in (P = 1) + ) # fused walks: the carry plane (P > 1) / S_in (P = 1) self.dht_index_fused = tuple( int(i) for i in hp.dht_index_fused - ) # L-B: the carry plane (P > 1) / dS_out (P = 1) + ) # fused walks: the carry plane (P > 1) / dS_out (P = 1) self.out_index = tuple(int(i) for i in hp.out_index) self.in_index = tuple( in_slot[s // P] if (s % P == 0 and frags[s // P].crossing_in) else -1 @@ -1521,7 +2310,7 @@ def lpt(lst): n_carry = self.n_frag * (P - 1) if P > 1 else 0 if ( self.fused_walks and P > 1 - ): # L-B: + one slot per crossing fragment for the S_in / dS_out copies + ): # fused walks: + one slot per crossing fragment for the S_in / dS_out copies n_carry += max(self.n_in, self.n_out) self.carry_tiles = tuple( tuple(c * H + h for h in range(H)) for c in range(n_carry) @@ -1639,6 +2428,24 @@ def lpt(lst): "tm": 1, "ws": 1, }, + # merge_fuse="carry": the two carry families of a phase as one launch (chains of carry_start then + # carry_in; dcarry_end then dcarry_cross); the start / end chains read no x0 (x0_index -1), the others S_in / dS_out + "carry_both": { + "m": "stage_m", + "c": "stage_b", + "x0": "S_in", + "y": "carry", + "tm": 0, + "ws": 1, + }, + "dcarry_both": { + "m": "stage_m", + "c": "stage_b", + "x0": "b2_recv", + "y": "carry", + "tm": 1, + "ws": 1, + }, } # the b1 pool bound of a family depends on the layout form it addresses _B1_BOUND = { @@ -1665,7 +2472,7 @@ def bound(name: str, tiles, limit_elems: int) -> None: bound("b1 d tiles (d form)", self.b1_d_tiles["d"], n["b1_send_d"]) bound("b1 d tiles (dm form)", self.b1_d_tiles["dm"], n["b1_send_dm"]) bound("b1 M tiles", self.b1_m_tiles, n["b1_send_dm"]) - # the L-B tables address the fused carry plane (n_frag (P - 1) + max(n_in, n_out) slots; allocated only when + # the fused-walk tables address the fused carry plane (n_frag (P - 1) + max(n_in, n_out) slots; allocated only when # fused_walks is set, inert otherwise) at P > 1 and the S_in / dS_out pools at P = 1 n_fused = self.n_frag * (self.P - 1) + max(self.n_in, self.n_out) fused_pool = ( @@ -1726,7 +2533,7 @@ def carry(f, i: int, h: int) -> int: ) extra = ( self.n_frag * (P - 1) - ) # L-B: the extra carry slots hold the S_in (dS_out) copies of the fused form + ) # fused walks: the extra carry slots hold the S_in (dS_out) copies of the fused form def copy_slot( f, h: int, slot: int @@ -1798,6 +2605,9 @@ def copy_slot( copy_slot(f, h, out_slot[f.slot]), ) ch[name].add(elems, x0=x0, y_final=carry(f, 0, h)) + # merge_fuse="carry": the two carry families of a phase as one launch + ch["carry_both"] = Chains.concat(ch["carry_start"], ch["carry_in"]) + ch["dcarry_both"] = Chains.concat(ch["dcarry_end"], ch["dcarry_cross"]) for name, c in ch.items(): c.check_alignment() self._check_chain(name, c) @@ -1917,8 +2727,21 @@ def split_elems(self, coll: str) -> tuple[list[int], list[int]]: ) return list(self.splits[f"{coll}_send"]), list(self.splits[f"{coll}_recv"]) - def device_tiles(self, tiles, device: torch.device) -> torch.Tensor: - """The flat int64 device copy of a tile table ``[slot][h]`` (cached per device).""" + def device_tiles( + self, tiles, device: torch.device, *, by_identity: bool = False + ) -> torch.Tensor: + """The flat int64 device copy of a tile table ``[slot][h]`` (cached per device). ``by_identity``: the + cache key is the table object itself (one of this table set's tuples, pinned by the entry) instead of its + content -- no per-call conversion of the nested rows.""" + if by_identity: + cache = self._id_cache("_device_tiles_id") + id_key = (id(tiles), str(device)) + hit = cache.get(id_key) + if hit is not None and hit[0] is tiles: + return hit[1] + t = self.device_tiles(tiles, device) + cache[id_key] = (tiles, t) + return t key = (tuple(tuple(int(t) for t in row) for row in tiles), str(device)) cache = self.__dict__.setdefault("_device_tiles", {}) t = cache.get(key) @@ -1928,6 +2751,74 @@ def device_tiles(self, tiles, device: torch.device) -> torch.Tensor: cache[key] = t return t + @functools.cached_property + def walk_all_p0(self) -> list[int]: + """Trimmed host path: the document-first segments of ``walk_all`` (``s % P == 0``) as one list object per table set: the + trimmed backward passes it (not a per-call comprehension) to the identity-keyed device-table cache.""" + return [s for s in self.walk_all if s % self.P == 0] + + @functools.cached_property + def walk_all_rest(self) -> list[int]: + """Trimmed host path: the other segments of ``walk_all`` (``s % P != 0``), one list object per table set (see ``walk_all_p0``).""" + return [s for s in self.walk_all if s % self.P != 0] + + ID_CACHE_MAX = 256 # entries per identity-keyed cache; a caller passing fresh objects cannot grow one without bound + + def _id_cache(self, name: str) -> dict: + """Trimmed host path: one identity-keyed cache of this table set (``dev_list`` / ``indexed_meta`` / ``device_tiles``), emptied + when it holds ID_CACHE_MAX entries (the table set's own objects number far fewer; the bound only limits a misuse).""" + cache = self.__dict__.setdefault(name, {}) + if len(cache) >= self.ID_CACHE_MAX: + cache.clear() + return cache + + def dev_list( + self, + values: Sequence[int], + device: torch.device | str, + dtype: torch.dtype = torch.int32, + ) -> torch.Tensor: + """Trimmed host path: the device copy of one of this table set's host int lists (a walk list, ``hp.seg_chunk_start``), cached + by the list object's identity (the entry pins the list; the table set's lists are never mutated). A list object + not owned by the table set misses on every call (a device copy each time): callers pass the set's own lists.""" + cache = self._id_cache("_dev_lists") + key = (id(values), str(device), dtype) + hit = cache.get(key) + if hit is not None and hit[0] is values: + return hit[1] + t = torch.tensor([int(v) for v in values], dtype=dtype, device=device) + cache[key] = (values, t) + return t + + def indexed_meta( + self, + index: Sequence[int], + tiles: Sequence[Sequence[int]], + H: int, + device: torch.device | str, + ) -> _IndexedMeta: + """Trimmed host path: the per-``(index, tiles, H, device)`` part of an :class:`Indexed` pool view -- the integer tuples, the + two-integer addressing and the int32 device index table -- computed once per table set (the lists are this + set's; the entry pins them). ``KcpCall.indexed`` builds its ``Indexed`` from it on the trimmed host path.""" + cache = self._id_cache("_indexed_meta") + key = (id(index), id(tiles), int(H), str(device)) + hit = cache.get(key) + if hit is not None and hit[0] is index and hit[1] is tiles: + return hit[2] + index_t = tuple(int(i) for i in index) + tiles_t = tuple(tuple(int(t) for t in row) for row in tiles) + addr = Indexed(torch.empty(0), index_t, tiles_t).addressing(int(H)) + meta = _IndexedMeta( + index=index_t, + tiles=tiles_t, + H=int(H), + addr=addr, + device=str(device), + dev_index=torch.tensor(list(index_t), dtype=torch.int32, device=device), + ) + cache[key] = (index, tiles, meta) + return meta + def pool_sizes(self, backward: bool) -> dict[str, int]: names = [ "c1_send", @@ -2014,6 +2905,12 @@ def __init__( self._pools: dict[str, torch.Tensor] = {} self._tables: dict[tuple[Any, int, int, str], _Tables] = {} self.local: Any | None = None + self._last_ensure: tuple[Any, ...] | None = ( + None # (ctx, H, device, backward, fused_walks, walk_geometry, tb) + ) + self._slices: dict[ + tuple[str, int], tuple[torch.Tensor, torch.Tensor] + ] = {} # pool[:n] views per (name, n) if ctx is not None: if H is None: raise ValueError("KcpWorkspace: H is required with ctx") @@ -2043,11 +2940,19 @@ def ensure( backward: bool = True, fused_walks: int = FUSED_WALKS_DEFAULT, walk_geometry: str = WALK_GEOMETRY_DEFAULT, + host_trim: bool = False, ) -> _Tables: - """Size the pools (and the local arena) for ``(ctx, H)``; returns the tables.""" + """Size the pools (and the local arena) for ``(ctx, H)``; returns the tables. ``host_trim``: a call + with the same ``(ctx, H, device, backward, fused_walks, walk_geometry)`` as the last ensure returns its tables + at once (the pools only grow and the arena is replaced only by a different plan, so nothing is to do).""" device = ( torch.device(device) if device is not None else torch.device(ctx.device) ) + sig = (int(H), device, bool(backward), int(fused_walks), str(walk_geometry)) + if host_trim: + last = self._last_ensure + if last is not None and last[0] is ctx and last[1] == sig: + return last[2] if self.device is None: self.device = device elif self.device != device: @@ -2055,14 +2960,25 @@ def ensure( f"KcpWorkspace: lives on {self.device}, the call is on {device}" ) tb = self.tables(ctx, H, fused_walks=fused_walks, walk_geometry=walk_geometry) + replaced = False for name, n in tb.pool_sizes(backward).items(): cur = self._pools.get(name) if cur is None or cur.numel() < n: self._pools[name] = _alloc((n,), self.state_dtype, device) + replaced = True if name == "eye": self._pools[name].view(D, D).copy_( torch.eye(D, dtype=self.state_dtype, device=device) ) + if ( + replaced and self._slices + ): # a superseded pool must not live on in the cached [:n] views + for key in [ + k + for k, (t, _) in self._slices.items() + if self._pools.get(k[0]) is not t + ]: + del self._slices[key] if self.with_local and ctx.T_r > 0 and torch.device(device).type == "cuda": p = self.local.plan if self.local is not None else None if ( @@ -2078,9 +2994,13 @@ def ensure( backward=backward, save_mode=self.save_mode, ) + if host_trim: + self._last_ensure = (ctx, sig, tb) return tb - def pool(self, name: str, n: int | None = None) -> torch.Tensor: + def pool( + self, name: str, n: int | None = None, *, cached: bool = False + ) -> torch.Tensor: t = self._pools.get(name) if t is None: if n not in (None, 0) or self.device is None: @@ -2088,7 +3008,18 @@ def pool(self, name: str, n: int | None = None) -> torch.Tensor: f"KcpWorkspace: pool {name!r} was not sized for this call (ensure() first)" ) t = self._pools[name] = _alloc((0,), self.state_dtype, self.device) - return t if n is None else t[:n] + if n is None: + return t + if ( + cached + ): # the [:n] view of an unchanged pool tensor is the same view call after call + hit = self._slices.get((name, n)) + if hit is not None and hit[0] is t: + return hit[1] + v = t[:n] + self._slices[(name, n)] = (t, v) + return v + return t[:n] def has(self, name: str) -> bool: return name in self._pools @@ -2096,7 +3027,14 @@ def has(self, name: str) -> bool: @property def bytes(self) -> int: kcp = sum(t.numel() * t.element_size() for t in self._pools.values()) - return kcp + (int(self.local.bytes) if self.local is not None else 0) + local = 0 + if ( + self.local is not None + ): # the arena and the intra record's scratch table set where it writes one + local = int(self.local.bytes) + cake_kcp_plan.local_scratch_table_bytes( + self.local.plan + ) + return kcp + local @property def pool_bytes(self) -> dict[str, int]: @@ -2110,16 +3048,14 @@ def pool_bytes(self) -> dict[str, int]: @dataclasses.dataclass(frozen=True) class KcpSaved: """The saved context of one forward (``boundary_mode="saved"``, the design note §16): ``S_in`` (the fresh c2 - output), the owner's c1 receive tensor (``save_transfer``), ``M_seg`` (``save_transfer`` at ``P > 1``), the - token-row ``Aqk`` / ``Akk`` of the existing version-1 set (kernel path), and the scalars of the call. The int32 - tables are rebuilt from ``ctx`` (droppable, §16 item 3).""" + output), the owner's c1 receive tensor (``save_transfer``), ``M_seg`` (``save_transfer`` at ``P > 1``) and the + scalars of the call. The int32 tables are rebuilt from ``ctx`` (droppable, §16 item 3). No token-row ``Aqk`` / + ``Akk``: the backward recomputes the internal rows it reads in both modes (version 2; version 1 held them).""" version: int S_in: torch.Tensor c1_recv: torch.Tensor | None M_seg: torch.Tensor | None - Aqk: torch.Tensor | None - Akk: torch.Tensor | None T_r: int H: int n_frag: int @@ -2134,11 +3070,11 @@ class KcpSaved: ctx: Any transfer_walk: str = "tf32_hilo" - _OPTIONAL = ("c1_recv", "M_seg", "Aqk", "Akk") + _OPTIONAL = ("c1_recv", "M_seg") def tensors(self) -> tuple[torch.Tensor, ...]: return (self.S_in,) + tuple( - t for t in (self.c1_recv, self.M_seg, self.Aqk, self.Akk) if t is not None + t for t in (self.c1_recv, self.M_seg) if t is not None ) def flags(self) -> tuple[bool, ...]: @@ -2184,7 +3120,7 @@ def bytes(self) -> int: @dataclasses.dataclass class KcpOptions: - """The extra keywords of the entry points (the design note §10.1): ``save_transfer`` (R6: keep the owner's c1 + """The extra keywords of the entry points (the design note §10.1): ``save_transfer`` (keep the owner's c1 receive tensor / ``M_seg`` for the backward; ``None`` = the default = resolve: :data:`cake_kcp_plan.SAVE_TRANSFER_DEFAULT` = ``True`` since an earlier revision for the a2a exchange, ``False`` for ``exchange="ring"``, whose path receives no c1 tensor; an explicit ``False`` = the earlier recompute form, which re-sends @@ -2195,7 +3131,11 @@ class KcpOptions: collective issue order and assert it), ``graph`` (``"off"`` shipped | ``"replay"``: keyed CUDA-graph capture / replay of the per-rank call -- see "CUDA-graph replay of a call" below for the validity key, the pointer-stability requirement, the private pools' memory, the eager fallback and :func:`graph_cache_clear`), - ``graph_cap`` (``"none"`` | ``"capture_pref"``: the device cache-config preference around a capture; + ``graph_cap`` (the cache-config form of a captured call: ``"node_attr"`` shipped + -- the tape one-CTA-per-SM cap as a launch attribute of the captured kernel nodes; ``"none"`` + = the previous form of record (no device preference), ``"capture_pref"`` (rejected) and ``"node_attr_post"`` stay selectable; + ``_kcp_test_knobs`` only), ``host_trim`` (``1`` shipped: the + trimmed host path of the eager call -- the same launches in the same order, bitwise; ``0`` = the previous host path; ``_kcp_test_knobs`` only), ``precision_mode`` (K3's static program; set only through ``_kcp_test_knobs``).""" save_transfer: bool | None = ( @@ -2210,14 +3150,16 @@ class KcpOptions: "tf32_hilo" # K3 program family; set only through _kcp_test_knobs ) p_rule: str | None = ( - None # plan P rule named by _kcp_test_knobs (the wave rules); None = the context's; must equal ctx.p_rule + None # plan P rule named by _kcp_test_knobs (the wave rules); None = the context's (P_RULE_DEFAULT "wave_fill1_p1tie"); must equal ctx.p_rule ) - schedule: str = SCHEDULE_DEFAULT # L-A: "wait_first" (shipped) issues the merges and tape walks after each wait; "overlap" via _kcp_test_knobs - fused_walks: int = FUSED_WALKS_DEFAULT # L-B (shipped 1): one tape launch per phase from the carry plane (+ the S_in / dS_out copies); 0 via _kcp_test_knobs + schedule: str = SCHEDULE_DEFAULT # "wait_first" (shipped) issues the merges and tape walks after each wait; "overlap" via _kcp_test_knobs + fused_walks: int = FUSED_WALKS_DEFAULT # shipped 1: one tape launch per phase from the carry plane (+ the S_in / dS_out copies); 0 via _kcp_test_knobs cols_rule: str = "fixed" # D1: "by_chains" = c8 up to COLS_RULE_CHAINS state chains, c16 beyond; _kcp_test_knobs only walk_geometry: str = WALK_GEOMETRY_DEFAULT # walk geometry: shipped "lpt_tape_one_cta_per_sm"; "shipped" (the earlier geometry) and the other WALK_GEOMETRIES via _kcp_test_knobs only graph: str = GRAPH_DEFAULT # graph replay: public keyword; "replay" captures a call once per validity key and replays it (eager on a key miss); also reachable through _kcp_test_knobs for the A/B record - graph_cap: str = GRAPH_CAP_DEFAULT # the cache-config form of a captured call ("none" of record | "capture_pref" rejected); _kcp_test_knobs only + graph_cap: str = GRAPH_CAP_DEFAULT # the cache-config form of a captured call ("node_attr" shipped | "none" the previous form of record | "capture_pref" rejected | "node_attr_post"); _kcp_test_knobs only + host_trim: int = HOST_TRIM_DEFAULT # shipped 1 = the trimmed host path of the eager call (same launches, bitwise); 0 = the previous path, _kcp_test_knobs only + merge_fuse: str = MERGE_FUSE_DEFAULT # "carry" = carry_start + carry_in (dcarry_end + dcarry_cross) in one K4 launch; _kcp_test_knobs only def check(self) -> KcpOptions: if self.exchange not in EXCHANGES: @@ -2234,7 +3176,7 @@ def check(self) -> KcpOptions: ) if int(self.head_groups) != 1: raise NotImplementedError( - "head_groups > 1 (lever L3) is not part of Phase 1" + "head_groups > 1 is not supported on the KCP route" ) if ( self.save_transfer is None @@ -2282,6 +3224,15 @@ def _options(kwargs: dict[str, Any]) -> KcpOptions: through the plan module's validator (``precision_mode``) and, for the test knob ``transfer_walk`` (K3's ``bf16x3`` / ``bf16x3_hw64`` candidates of ``build.test.json``), through the same environment gate here, so the public API can never select a non-zero ``precision_mode`` or a non-shipped transfer-walk program.""" + pre = kwargs.get("_kcp_options") + if ( + pre is not None + ): # the options validated by chunk_kda_kcp on the caller's frame, passed through the autograd node + if not isinstance(pre, KcpOptions) or len(kwargs) != 1: + raise TypeError( + "_kcp_options is the validated KcpOptions of the entry point and travels alone" + ) + return pre kwargs = dict(kwargs) raw = kwargs.pop("_kcp_test_knobs", None) transfer_walk = TRANSFER_WALKS[0] @@ -2299,7 +3250,7 @@ def _options(kwargs: dict[str, Any]) -> KcpOptions: p2: dict[str, Any] = {} if isinstance(raw, dict) and any( k in raw for k in P2_KNOBS - ): # the earlier levers: the same gate + ): # the earlier knobs: the same gate if os.environ.get(cake_kcp_plan.TEST_KNOBS_ENV) != "1": raise ValueError( f"_kcp_test_knobs is a private test hook: it is accepted only with {cake_kcp_plan.TEST_KNOBS_ENV}=1 in the environment" @@ -2308,7 +3259,9 @@ def _options(kwargs: dict[str, Any]) -> KcpOptions: for key in P2_KNOBS: if key in raw: p2[key] = ( - int(raw.pop(key)) if key == "fused_walks" else str(raw.pop(key)) + int(raw.pop(key)) + if key in ("fused_walks", "host_trim") + else str(raw.pop(key)) ) if ( "graph" in p2 @@ -2321,7 +3274,7 @@ def _options(kwargs: dict[str, Any]) -> KcpOptions: knobs = cake_kcp_plan.validate_test_knobs(raw) p_rule = ( knobs["p_rule"] if isinstance(raw, dict) and "p_rule" in raw else None - ) # lever L-C: the context carries the rule + ) # plan rule: the context carries the rule unknown = sorted(set(kwargs) - set(_OPTION_KEYS)) if unknown: raise TypeError( @@ -2420,7 +3373,7 @@ def __init__( raise ValueError( f"_kcp_test_knobs p_rule {opts.p_rule!r} differs from the context's plan rule {getattr(ctx, 'p_rule', cake_kcp_plan.P_RULE_DEFAULT)!r}: give the knob to build_kcp_context" ) - self.schedule = str(opts.schedule) # L-A + self.schedule = str(opts.schedule) self.cols_rule = str(opts.cols_rule) # D1 if int(tb.fused_walks) != int(opts.fused_walks): raise ValueError( @@ -2431,11 +3384,18 @@ def __init__( raise ValueError( f"tables built with walk_geometry={tb.walk_geometry!r}, the call asks for {self.walk_geometry!r}" ) + # 1 = the trimmed host path (shipped), 0 = the previous path (same launches, bitwise) + self.host_trim = int(opts.host_trim) + self.merge_fuse = str(opts.merge_fuse) # fused carry merges # -- buffers def buf(self, name: str) -> torch.Tensor: t = self.bufs.get(name) - return t if t is not None else self.ws.pool(name, self.tb.sizes[name]) + return ( + t + if t is not None + else self.ws.pool(name, self.tb.sizes[name], cached=self.host_trim == 1) + ) def fresh(self, name: str) -> torch.Tensor: """A fresh per-call tensor standing in for pool ``name`` (the saved-set tensors of SPEC-4).""" @@ -2446,6 +3406,9 @@ def fresh(self, name: str) -> torch.Tensor: def indexed( self, pool: str, index: Sequence[int], tiles: Sequence[Sequence[int]] ) -> Indexed: + if self.host_trim: # the tables' precomputed tuples / addressing / device index table (same values) + meta = self.tb.indexed_meta(index, tiles, self.H, self.device) + return Indexed(self.buf(pool), meta.index, meta.tiles, meta) return Indexed( self.buf(pool), tuple(int(i) for i in index), @@ -2464,7 +3427,10 @@ def walk(self, name: str, walk: Sequence[int], **args: Any) -> None: if len(walk) == 0: return self.counters["walks"] += 1 - self.kernel(name, walk=[int(s) for s in walk], **args) + # the trimmed path passes the tables' own list (its identity keys the cached device table) + self.kernel( + name, walk=walk if self.host_trim else [int(s) for s in walk], **args + ) def merge( self, @@ -2500,11 +3466,14 @@ def transpose_tiles(self, src: str, src_tiles, dst: str, dst_tiles) -> None: return self.counters["launches"] += 1 self.sequence.append("kcp_transpose_tiles") + by_id = ( + self.host_trim == 1 + ) # the tile tables are this table set's tuples -> identity-keyed cache self.backend.transpose_tiles( src=self.buf(src), - src_tiles=self.tb.device_tiles(src_tiles, self.device), + src_tiles=self.tb.device_tiles(src_tiles, self.device, by_identity=by_id), dst=self.buf(dst), - dst_tiles=self.tb.device_tiles(dst_tiles, self.device), + dst_tiles=self.tb.device_tiles(dst_tiles, self.device, by_identity=by_id), ) def stage(self, name: str, **extra: Any) -> None: @@ -2535,8 +3504,20 @@ def a2a( raise ValueError( f"{tag}: split lists must have one entry per rank ({self.world})" ) - send = self.buf(send_name)[: sum(send_splits)] - recv = self.buf(recv_name)[: sum(recv_splits)] + if self.host_trim: # cached [:n] views of the pools (a fresh per-call tensor is sliced as today) + send = ( + self.ws.pool(send_name, sum(send_splits), cached=True) + if send_name not in self.bufs + else self.buf(send_name)[: sum(send_splits)] + ) + recv = ( + self.ws.pool(recv_name, sum(recv_splits), cached=True) + if recv_name not in self.bufs + else self.buf(recv_name)[: sum(recv_splits)] + ) + else: + send = self.buf(send_name)[: sum(send_splits)] + recv = self.buf(recv_name)[: sum(recv_splits)] self.counters["collectives"] += 1 self._sequence(tag) if self.world == 1 or not (dist.is_available() and dist.is_initialized()): @@ -2603,6 +3584,9 @@ def _record(kind: str, call: KcpCall) -> None: "fused_walks": call.tb.fused_walks, "cols_rule": call.cols_rule, # the earlier knobs "walk_geometry": call.walk_geometry, # walk-launch geometry + # the host path of this call (1 trimmed = shipped / 0 previous; same launches) + "host_trim": call.host_trim, + "merge_fuse": call.merge_fuse, # the fused carry merges (test knob) "merges": list( call.merge_log ), # D1 evidence: (family, ncols, n_chains, explicit cols_per_cta or None = registry) @@ -2615,17 +3599,23 @@ def _record(kind: str, call: KcpCall) -> None: "graph_fallback": None, # the graph mode and how this call ran "cache_config": call.scratch.get( "cache_config", "none" - ), # "bracket" (eager capped launches) / "capture_pref" / "none" / "unavailable" + ), # "bracket" (eager capped launches) / "node_attr" (the replay form of record) / "node_attr_post" / "capture_pref" / "none" / "unavailable" (no loadable CUDA runtime, or a degraded cap) + "cache_config_degraded": call.scratch.get( + "node_attr_degraded" + ), # why a captured tape launch stayed uncapped (None otherwise) "cache_config_bracket": call.scratch.get( "cache_config_bracket" ), # the capture-wide preference's set / reset evidence (arm "capture_pref") + "cache_config_nodes": call.scratch.get( + "node_attr_launches" + ), # kernel-node cap forms: captured tape launches whose kernel node carries the carveout attribute } CACHE_CONFIG_STATS[LAST_CALL[kind]["cache_config"]] += 1 # ---------------------------------------------------------------------------------------------------------------- -# Launch sequences (the design note §3.5 forward, §6.5 backward; one compute stream, waits at first consumption, R7; -# shipped: schedule wait_first + fused tape walks, the R7 overlaps only under the overlap knob) +# Launch sequences (the design note §3.5 forward, §6.5 backward; one compute stream, waits at first consumption; +# shipped: schedule wait_first + fused tape walks, the overlapped issue order only under the overlap knob) # ---------------------------------------------------------------------------------------------------------------- @@ -2998,7 +3988,7 @@ def _summaries(call: KcpCall) -> None: def _fwd_tapes( call: KcpCall, *, out: torch.Tensor | None, carry: Indexed | None, wait_first: bool ) -> None: - """The forward tape walk(s) after ``S_in`` is complete. ``fused_walks`` (L-B): one launch from one pool -- the + """The forward tape walk(s) after ``S_in`` is complete. ``fused_walks``: one launch from one pool -- the carry plane at ``P > 1`` (``h0_index_fused``: carries, the ``S_in`` copies of ``carry_in``, ``-1`` zero states), ``S_in`` itself at ``P = 1`` -- over every live segment (``wait_first``) or the crossing-in fragments' segments (the document-start walk ran under c1); otherwise the three launches of the design note §3.5 (the document-start @@ -3021,7 +4011,7 @@ def _fwd_tapes( if wait_first: call.walk( "kcp_fwdh_tape", tb.walk_start, h0=carry, out=out - ) # document starts: after wait(c2) (L-A) + ) # document starts: after wait(c2) call.walk("kcp_fwdh_tape", tb.walk_in_p0, h0=h0_in, out=out) if tb.P > 1: call.walk("kcp_fwdh_tape", tb.walk_in_p1, h0=carry, out=out) @@ -3045,7 +4035,7 @@ def _adj_tapes(call: KcpCall, *, dcarry: Indexed | None, wait_first: bool) -> No if wait_first: call.walk( "kcp_dhu_tape", tb.walk_end, dht=dcarry - ) # document ends: after wait(b2) (L-A) + ) # document ends: after wait(b2) call.walk("kcp_dhu_tape", tb.walk_cross_last, dht=dht_cross) if tb.P > 1: call.walk("kcp_dhu_tape", tb.walk_cross_rest, dht=dcarry) @@ -3054,7 +4044,7 @@ def _adj_tapes(call: KcpCall, *, dcarry: Indexed | None, wait_first: bool) -> No def _forward(call: KcpCall, *, save: bool, replay: bool) -> None: """The design note §3.5. ``save``: the saved set is kept (fresh ``S_in``; with ``save_transfer`` also the c1 receive tensor and ``M_seg``); ``replay``: the forward of a replay-mode backward (no ``fwd_o``, nothing kept). - ``schedule="wait_first"`` (L-A; the shipped default since an earlier revision): ``c1 -> wait -> fwd merge -> c2 -> wait -> carry merges -> tape walk(s)`` + ``schedule="wait_first"`` (the shipped default since an earlier revision): ``c1 -> wait -> fwd merge -> c2 -> wait -> carry merges -> tape walk(s)`` -- no merge or walk is issued while a collective is in flight (the R7 overlap of the document-start walk under c1 is dropped); the merges read only ``c1_recv`` / the staging planes and the tapes read no collective buffer, so every published bit is the same.""" @@ -3066,9 +4056,13 @@ def _forward(call: KcpCall, *, save: bool, replay: bool) -> None: call.fresh("c1_recv") if tb.P > 1: call.fresh("stage_m") # M_seg^T of the segmented fragments (§16 item 5b) - call.stage("fwd_prep") - call.stage("fwd_intra", write_ext=1 if (save and not replay) else 0, write_int=1) - call.stage("wy") + if replay: # the replay's prep (the GLM prep's rows, the internal dout copy), then the internal Aqk / Akk rows + call.stage("bwd_prep", saved=None) + call.stage("fwd_intra") + call.stage("wy") + else: # the intra program writes the GLM prep's rows itself (write_prep = 1); no token-row copies are saved + call.stage("fwd_intra") + call.stage("wy_fwd") _summaries(call) carry = call.indexed("carry", tb.carry_index, tb.carry_tiles) if tb.P > 1 else None if ( @@ -3094,11 +4088,14 @@ def _forward(call: KcpCall, *, save: bool, replay: bool) -> None: w2 = call.a2a("c2", "c2_send", "S_in", *tb.split_elems("c2")) call.wait(w2) if tb.P > 1: - if wait_first: - call.merge("carry_start") - call.merge( - "carry_in" - ) # with fused_walks also the S_in copies into the extra carry slots (write_states) + if wait_first and call.merge_fuse == "carry": + call.merge("carry_both") # carry_start + carry_in in one launch + else: + if wait_first: + call.merge("carry_start") + call.merge( + "carry_in" + ) # with fused_walks also the S_in copies into the extra carry slots (write_states) _fwd_tapes(call, out=call.out, carry=carry, wait_first=wait_first) if not replay: call.stage("fwd_o") @@ -3108,7 +4105,7 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: """The design note §6.5. ``saved=None`` is the replay mode (the forward sequence without ``fwd_o`` recomputes ``S_in``, the tapes and the owner's c1 receive tensor); otherwise the tapes come from K1 ``_tape`` over the saved ``S_in`` and the owner's adjoint reads ``M_f`` from b1 (``save_transfer=False``) or ``M_f^T`` from the saved c1 - receive tensor. ``schedule="wait_first"`` (L-A; the shipped default since an earlier revision): ``b1 -> wait -> adjoint merge -> b2 -> wait -> dcarry merges -> + receive tensor. ``schedule="wait_first"`` (the shipped default since an earlier revision): ``b1 -> wait -> adjoint merge -> b2 -> wait -> dcarry merges -> adjoint tape walk(s)`` (the R7 overlap of the document-last walk under b2 is dropped).""" tb = call.tb wait_first = call.schedule == "wait_first" @@ -3125,6 +4122,9 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: if saved.M_seg is not None: call.bufs["stage_m"] = saved.M_seg call.stage("bwd_prep", saved=saved) + call.stage( + "fwd_intra" + ) # the internal Aqk / Akk rows the WY, dav and gradient programs read call.stage("wy") if ( tb.P > 1 @@ -3144,11 +4144,14 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: use_u=0, write_transposed=0, ) - call.merge("carry_start") - call.merge("carry_in") + if call.merge_fuse == "carry": + call.merge("carry_both") # one launch for both carry families + else: + call.merge("carry_start") + call.merge("carry_in") if ( tb.fused_walks and tb.P > 1 - ): # L-B: the whole tape replay from the carry plane in one launch + ): # fused walks: the whole tape replay from the carry plane in one launch call.walk( "kcp_fwdh_tape", tb.walk_live, @@ -3156,7 +4159,10 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: out=None, ) else: - p0 = [s for s in tb.walk_all if s % tb.P == 0] + # the trimmed arm passes the table set's own sublists (a fresh comprehension per call would miss the + # identity-keyed device-table cache on every call: one blocking host-to-device copy per backward) + trim = bool(call.host_trim) + p0 = tb.walk_all_p0 if trim else [s for s in tb.walk_all if s % tb.P == 0] call.walk( "kcp_fwdh_tape", p0, @@ -3164,9 +4170,14 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: out=None, ) if tb.P > 1: + rest = ( + tb.walk_all_rest + if trim + else [s for s in tb.walk_all if s % tb.P != 0] + ) call.walk( "kcp_fwdh_tape", - [s for s in tb.walk_all if s % tb.P != 0], + rest, h0=call.indexed("carry", tb.carry_index, tb.carry_tiles), out=None, ) @@ -3229,11 +4240,14 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: ) # document ends: under b2 (R7) call.wait(wb2) if tb.P > 1: - if wait_first: - call.merge("dcarry_end") - call.merge( - "dcarry_cross" - ) # with fused_walks also the dS_out copies into the extra carry slots + if wait_first and call.merge_fuse == "carry": + call.merge("dcarry_both") # dcarry_end + dcarry_cross in one launch + else: + if wait_first: + call.merge("dcarry_end") + call.merge( + "dcarry_cross" + ) # with fused_walks also the dS_out copies into the extra carry slots _adj_tapes(call, dcarry=dcarry, wait_first=wait_first) call.stage("dqkg") call.stage("glm_intra") @@ -3260,11 +4274,16 @@ def _backward(call: KcpCall, saved: KcpSaved | None) -> None: # # Cache-config preference under replay (a probe showed cudaDeviceSetCacheConfig inside a capture # invalidates it): the per-launch bracket of the capped tape programs (_FiPrograms.launch) is never issued inside a stream -# capture. The knob ``graph_cap`` selects the captured call's form: "capture_pref" sets PreferL1 once before the capture -# of a key and resets PreferNone right after the capture ends (never around the replays), "none" (the default of record) -# touches no preference. Every call records the form as LAST_CALL[kind]["cache_config"] ("bracket" for an eager call -# under a capped geometry, "capture_pref" / "none" for a captured call and its replays, "unavailable" without a loadable -# CUDA runtime); CACHE_CONFIG_STATS counts them. +# capture. The knob ``graph_cap`` selects the captured call's form: "node_attr" (the default of record) puts the one-CTA-per-SM +# carveout on each captured tape kernel node as a launch +# attribute right after its launch is captured, so that the replayed graph itself carries the cap; "node_attr_post" sets the +# same attribute on the kept node handles after the capture ends, before the instantiation; "capture_pref" sets PreferL1 once +# before the capture of a key and resets PreferNone right after the capture ends (never around the replays); "none" (the +# previous default of record) touches no preference. Degrade: a missing entry point or a capture-info +# reply without exactly one kernel-node dependency leaves that launch uncapped (one RuntimeWarning per process; the call records +# "unavailable" + cache_config_degraded), never an eager fallback. Every call records the form as LAST_CALL[kind]["cache_config"] +# ("bracket" for an eager call under a capped geometry, "node_attr" / "node_attr_post" / "capture_pref" / "none" for a captured +# call and its replays, "unavailable" without a loadable CUDA runtime or degraded); CACHE_CONFIG_STATS counts them. GRAPH_STATS: dict[str, int] = {"eager": 0, "capture": 0, "replay": 0, "fallback": 0} GRAPH_CACHE_MAX = 64 # keys kept per process; the oldest entry goes when a new key arrives beyond this @@ -3349,10 +4368,22 @@ def _graph_cap_bracket(call: KcpCall) -> tuple[str, int | None]: """The device cache-config preference around a capture of ``call``: ``("capture_pref", WALK_CACHE_CONFIG_CAP)`` under ``graph_cap="capture_pref"`` with a capped geometry and a loadable CUDA runtime (PreferL1 for the whole capture, PreferNone right after it ends), ``("unavailable", None)`` when no libcudart soname loads, ``("none", None)`` - otherwise (``graph_cap="none"``, or a geometry without capped programs).""" - if call.opts.graph_cap != "capture_pref" or not WALK_CAP_PROGRAMS.get( - call.walk_geometry - ): + otherwise (``graph_cap="none"``, or a geometry without capped programs); ``("node_attr", None)`` / ``("node_attr_post", None)`` + under the kernel-node forms (``"node_attr"`` shipped: no device preference, the cap is a launch + attribute of the captured nodes; ``("unavailable", None)`` without the capture-graph entry points).""" + if not WALK_CAP_PROGRAMS.get(call.walk_geometry): + return "none", None + form = call.opts.graph_cap + if ( + form in GRAPH_NODE_CAPS + ): # the carveout attribute on the captured tape kernel nodes; no device preference + if _cudart_graph_available() is None or ( + form == "node_attr_post" + and not hasattr(torch.cuda.CUDAGraph, "instantiate") + ): + return "unavailable", None + return form, None + if form != "capture_pref": return "none", None if _cudart_available() is None: return "unavailable", None @@ -3363,10 +4394,13 @@ def _graph_capture(kind: str, entry: _GraphEntry, call: KcpCall, body: Any) -> A """Capture ``body()`` (the call's allocations, launches and collectives) into a CUDA graph on the capture's side stream, replay it once and keep it for the key; returns the body's result, or ``None`` when the capture failed (recorded in the entry; the key then runs eagerly for the rest of the process). The cache-config preference of the - capture is :func:`_graph_cap_bracket`'s: set before the capture begins, reset right after it ends, never inside it - (the call records the form and the set / reset evidence).""" - g = torch.cuda.CUDAGraph() + capture is :func:`_graph_cap_bracket`'s: set before the capture begins, reset right after it ends, never inside it; + the call records the form and the set / reset evidence.""" cache_config, pref = _graph_cap_bracket(call) + post = ( + cache_config == "node_attr_post" + ) # the node attributes go on after the capture, before the instantiation + g = torch.cuda.CUDAGraph(keep_graph=True) if post else torch.cuda.CUDAGraph() dev = ( call.device.index if call.device.index is not None @@ -3382,11 +4416,24 @@ def _graph_capture(kind: str, entry: _GraphEntry, call: KcpCall, body: Any) -> A finally: if pref is not None: bracket.append(walk_device_cache_config(dev, WALK_CACHE_CONFIG_DEFAULT)) + if post: + programs_of = getattr(call.backend, "programs", None) + programs = programs_of() if callable(programs_of) else None + call.scratch["node_attr_applied"] = ( + programs.apply_node_caps(call) + if isinstance(programs, _FiPrograms) + else 0 + ) + g.instantiate() except Exception as exc: # noqa: BLE001 - the failure is recorded; the caller runs the eager path entry.disabled = f"{type(exc).__name__}: {str(exc)[:200]}" GRAPH_STATS["fallback"] += 1 return None g.replay() + if call.scratch.get( + "node_attr_degraded" + ): # a tape launch of this capture stayed uncapped + cache_config = "unavailable" call.scratch["cache_config"] = cache_config call.scratch["cache_config_bracket"] = bracket or None _record(kind, call) @@ -3438,7 +4485,7 @@ def _prepare( raise ValueError(f"save_mode must be one of {SAVE_MODES}, got {save_mode!r}") if save_mode != "light": raise NotImplementedError( - "save_mode='tape' is not part of Phase 1 of the KCP route (light only)" + "save_mode='tape' is not supported on the KCP route (light only)" ) q, k, v, g, beta, T, H, batched = _operands( q, k, v, g, beta, ctx, dtypes=backend.dtypes @@ -3466,6 +4513,7 @@ def _prepare( backward=backward, fused_walks=opts.fused_walks, walk_geometry=opts.walk_geometry, + host_trim=bool(opts.host_trim), ) ops = {"q": q, "k": k, "v": v, "g": g, "beta": beta} if dout is not None: @@ -3582,8 +4630,6 @@ def _fwd_body( M_seg=call.bufs.get("stage_m") if (opts.save_transfer and tb.P > 1) else None, - Aqk=call.scratch.get("Aqk"), - Akk=call.scratch.get("Akk"), T_r=call.T_r, H=call.H, n_frag=tb.n_frag, @@ -3673,11 +4719,7 @@ def chunk_kda_kcp_bwd( opts = _options(options) entry = None if opts.graph != "off": - kept = ( - () - if saved is None - else (saved.S_in, saved.c1_recv, saved.M_seg, saved.Aqk, saved.Akk) - ) + kept = () if saved is None else (saved.S_in, saved.c1_recv, saved.M_seg) entry = _graph_entry( "bwd", ctx, @@ -3838,11 +4880,13 @@ def chunk_kda_kcp( **options: Any, ) -> torch.Tensor: """The differentiable KCP forward (the KCP contract §3.2 signature): ``out`` like ``v``; gradients for ``q, k, v, g, beta``.""" - _options(options) # validate on the caller's frame + opts = _options(options) # validate on the caller's frame if boundary_mode not in BOUNDARY_MODES: raise ValueError( f"boundary_mode must be one of {BOUNDARY_MODES}, got {boundary_mode!r}" ) + if opts.host_trim: # the forward and the backward reuse the validated options instead of re-validating the dict + options = {"_kcp_options": opts} return ChunkKDAKcpFunction.apply( q, k, v, g, beta, ctx, scale, save_mode, boundary_mode, workspace, dict(options) ) @@ -3858,6 +4902,8 @@ def chunk_kda_kcp( "TRANSFER_WALKS", "WALK_GEOMETRIES", "WALK_GEOMETRY_DEFAULT", + "HOST_TRIMS", + "HOST_TRIM_DEFAULT", "WALK_CAP_PROGRAMS", "WALK_CACHE_CONFIG_CAP", "WALK_CACHE_CONFIG_DEFAULT", @@ -3872,6 +4918,11 @@ def chunk_kda_kcp( "GRAPH_DEFAULT", "GRAPH_CAPS", "GRAPH_CAP_DEFAULT", + "GRAPH_NODE_CAPS", + "GRAPH_NODE_CARVEOUT", + "capture_kernel_node", + "graph_entry_points", + "set_kernel_node_carveout", "GRAPH_STATS", "CACHE_CONFIG_STATS", "graph_cache_clear", diff --git a/flashinfer/experimental/cake_kda_train/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_train/cake_kcp_jit.py index 8d7453925d3..96d6f06a057 100644 --- a/flashinfer/experimental/cake_kda_train/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_train/cake_kcp_jit.py @@ -31,7 +31,7 @@ the library from a prebuilt ``aot/`` file when a builder compiled it and otherwise through FlashInfer's JIT from the same sources and flags. -The programs are the six launches of the KCP composition +The programs are the eight launches of the KCP composition (:data:`PROGRAMS`, the names :mod:`.cake_kcp` passes to :func:`run_program`): * ``kcp_fwdh_offset`` / ``kcp_fwdh_tape`` -- the forward state recurrence @@ -48,7 +48,16 @@ (``grid_rule`` ``kcp_chains``: ``n_chains * (ncols / cols_per_cta)`` CTAs); two variants by the chain's column count (``select_rule`` ``columns``: ``cols_per_cta`` 8 for 128-column state chains, 16 for 256-column map - chains, the faster build of each on GB300). + chains, the faster build of each on GB300); +* ``kcp_fwd_o`` -- the chunk output stage of the forward at the rank's + shape: ``out`` from the ``h`` / ``v_new`` tapes the forward walk writes and + the intra stage's ``qg`` / ``Aqk`` rows, persistent CTAs over the (chunk, + head) tiles (``grid_rule`` ``kcp_sm_bounded``: ``min(num_chunks_pad * + num_heads, ctas_per_sm * sm_count)`` CTAs, at least one); +* ``kcp_dav`` -- the output backward at the rank's shape: ``dAqk`` and the + first value gradient ``dv1`` (read by the adjoint walk) from ``do_int``, the + ``v_new`` tape and the ``Aqk`` rows, one CTA per chunk pair and head + (``grid_rule`` ``kcp_pairs``: ``(num_chunks_pad / 2, num_heads, 1)``). Compact state pools (the walks' ``h0`` / ``dht`` inputs and ``out`` outputs) are flat FP32 tensors of ``128 x 128`` tiles addressed by an int32 slot table @@ -101,7 +110,8 @@ # The launches of the KCP composition, by the names :mod:`.cake_kcp` uses # (``cake_kcp.KCP_PROGRAMS``). The first five are walks over a segment list; -# the last composes chains. +# the sixth composes chains; the last two are the chunk stages of the forward +# output and the output backward at the rank's shape. PROGRAMS = ( "kcp_fwdh_offset", "kcp_fwdh_tape", @@ -109,6 +119,8 @@ "kcp_dhu_tape", "kcp_transfer_walk", "kcp_affine_merge", + "kcp_fwd_o", + "kcp_dav", ) WALK_PROGRAMS = PROGRAMS[:5] # Elements of one FP32 state tile [128, 128]. @@ -136,7 +148,13 @@ # tile pools ``m`` / ``c`` / ``x0`` / ``y`` and its chain tables # ``chain_start`` (int32 ``[n_chains + 1]``), ``elem_m`` / ``elem_c`` / # ``elem_y`` (int64 element offsets per element), ``x0_index`` / ``y_final`` -# (int64 ``[n_chains]``). +# (int64 ``[n_chains]``). The chunk stages' names: ``h2d`` the ``[NC * H * +# 128, 128]`` view of the ``h`` tape, ``aqk_int`` the BF16 internal ``Aqk`` +# rows ``[R, H, 64]`` and ``dAqk`` their FP32 gradient, ``out`` (bound by +# ``kcp_fwd_o``) the call's BF16 output ``[T, H, 128]``, and the int32 chunk +# tables ``chunk_bos`` / ``chunk_len`` (``[NC_pad]``) / ``seq_chunk_start`` +# (``[num_seqs + 1]``, the chunk table of the rank's sequences, not the KCP +# segment table ``seg_chunk_start``). LOGICAL_TENSORS = ( "w", "kg", @@ -149,6 +167,12 @@ "v_new", "dh", "dv2", + "h2d", + "aqk_int", + "dAqk", + "chunk_bos", + "chunk_len", + "seq_chunk_start", "seg_chunk_start", "walk_seg", "h0", @@ -618,6 +642,90 @@ "select_rule": {"kind": "columns", "ncols": 256}, "variant": "c16", }, + "kcp_fwd_o": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "qg_tma"], + ["buffer", "h_tma"], + ["buffer", "aqk_tma"], + ["buffer", "vn_tma"], + ["buffer", "o_out"], + ["buffer", "chunk_bos"], + ["buffer", "chunk_len"], + ["buffer", "seq_chunk_start"], + ["parameter", "num_seqs"], + ["parameter", "num_heads"], + ["parameter", "scale"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": {"ctas_per_sm": 1, "kind": "kcp_sm_bounded"}, + "kernel_symbol": "kernel_cake_kda_train_1214bfc28e50985bfa0c", + "launch": { + "block": (192, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 222208, + "use_pdl": False, + }, + "library": "cake_kda_train_1214bfc28e50985bfa0c", + "logical": { + "aqk_tma": "aqk_int", + "chunk_bos": "chunk_bos", + "chunk_len": "chunk_len", + "h_tma": "h2d", + "num_heads": "num_heads", + "num_seqs": "num_seqs", + "o_out": "out", + "qg_tma": "qg", + "scale": "scale", + "seq_chunk_start": "seq_chunk_start", + "vn_tma": "v_new", + }, + "program": "kcp_fwd_o", + "select_rule": {"kind": "always"}, + "variant": None, + }, + "kcp_dav": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "do_tma"], + ["buffer", "vnew_tma"], + ["buffer", "aqk_tma"], + ["buffer", "dAqk"], + ["buffer", "dv1"], + ["parameter", "num_heads"], + ["parameter", "scale"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": {"kind": "kcp_pairs"}, + "kernel_symbol": "kernel_cake_kda_train_f227dd21f1f3cb464b29", + "launch": { + "block": (192, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 99328, + "use_pdl": False, + }, + "library": "cake_kda_train_f227dd21f1f3cb464b29", + "logical": { + "aqk_tma": "aqk_int", + "dAqk": "dAqk", + "do_tma": "do_int", + "dv1": "dv1", + "num_heads": "num_heads", + "scale": "scale", + "vnew_tma": "v_new", + }, + "program": "kcp_dav", + "select_rule": {"kind": "always"}, + "variant": None, + }, } # The test-only record set: the transfer walk built with the static knob @@ -907,7 +1015,14 @@ def select_record(program: str, *, ncols: int | None = None) -> str: def launch_grid( - name: str, *, n_walk: int = 0, num_heads: int = 0, n_chains: int = 0, ncols: int = 0 + name: str, + *, + n_walk: int = 0, + num_heads: int = 0, + n_chains: int = 0, + ncols: int = 0, + num_chunks_pad: int = 0, + sm_count: int = 0, ) -> tuple[int, int, int]: """The launch grid of record ``name`` under its recorded grid rule. @@ -916,7 +1031,11 @@ def launch_grid( half-width test-only walk; the host lists the segments longest chain first); ``kcp_chains`` launches ``(n_chains * (ncols / cols_per_cta), 1, 1)`` (one CTA per chain and column block; ``ncols`` is a multiple of the record's - ``cols_per_cta``). + ``cols_per_cta``); ``kcp_sm_bounded`` the persistent grid ``(max(1, + min(num_chunks_pad * num_heads, ctas_per_sm * sm_count)), 1, 1)`` over the + (chunk, head) tiles of the rank's padded chunk count; ``kcp_pairs`` + ``(num_chunks_pad / 2, num_heads, 1)`` (one CTA per chunk pair and head; + ``num_chunks_pad`` is even). """ record = MODULES.get(name) or TEST_MODULES.get(name) if record is None: @@ -936,6 +1055,21 @@ def launch_grid( f"record {name}: n_chains={n_chains} must be >= 1 and ncols={ncols} a positive multiple of {cols}" ) return (int(n_chains) * (int(ncols) // cols), 1, 1) + if kind == "kcp_sm_bounded": + if int(num_chunks_pad) < 1 or int(num_heads) < 1 or int(sm_count) < 1: + raise ValueError( + f"record {name}: num_chunks_pad={num_chunks_pad}, num_heads={num_heads} and " + f"sm_count={sm_count} must be >= 1" + ) + bound = int(rule["ctas_per_sm"]) * int(sm_count) + return (max(1, min(int(num_chunks_pad) * int(num_heads), bound)), 1, 1) + if kind == "kcp_pairs": + if int(num_chunks_pad) < 1 or int(num_heads) < 1 or int(num_chunks_pad) % 2: + raise ValueError( + f"record {name}: num_chunks_pad={num_chunks_pad} must be a positive even count and " + f"num_heads={num_heads} >= 1" + ) + return (int(num_chunks_pad) // 2, int(num_heads), 1) raise ValueError(f"unknown Cake KDA KCP grid rule: {kind!r}") @@ -989,6 +1123,37 @@ def _argument_plan(name: str) -> tuple[tuple[str | None, int], ...]: return tuple(plan) +_ARCH_BY_INDEX: dict[int, str] = {} +_FFI_DEVICES: dict[int, Any] = {} + + +def _arch_of(device_index: int) -> str: + """:func:`device_arch` cached per device index (no capability query per launch).""" + arch = _ARCH_BY_INDEX.get(int(device_index)) + if arch is None: + arch = _ARCH_BY_INDEX[int(device_index)] = device_arch(int(device_index)) + return arch + + +def _ffi_stream(device_index: int) -> Any: + """Trimmed host path: the FFI stream context of a trimmed launch -- the device's current torch stream as a raw handle + (``use_raw_stream``), the stream ``tvm_ffi.use_torch_stream()`` would set without building a torch.cuda.Stream.""" + import torch # noqa: PLC0415 + + dev = _FFI_DEVICES.get(int(device_index)) + if dev is None: + dev = _FFI_DEVICES[int(device_index)] = tvm_ffi.device( + f"cuda:{int(device_index)}" + ) + get = getattr(torch._C, "_cuda_getCurrentRawStream", None) + raw = ( + int(get(int(device_index))) + if get is not None + else int(torch.cuda.current_stream(int(device_index)).cuda_stream) + ) + return tvm_ffi.use_raw_stream(dev, raw) + + def run_module( name: str, values: Mapping[str, Any], @@ -996,6 +1161,7 @@ def run_module( grid: tuple[int, int, int], device_index: int, module: Any | None = None, + trim: bool = False, ) -> None: """Launch record ``name`` with logical argument ``values`` on ``grid``. @@ -1016,6 +1182,12 @@ def run_module( raise ValueError( f"record {name} needs logical argument {logical_name!r}" ) from None + if trim: # same module, same arguments, same stream; the arch and the FFI device are cached + if module is None: + module = load(name, _arch_of(device_index)) + with _ffi_stream(device_index): + getattr(module, record["ffi_entry"])(*args) + return if module is None: module = load(name, device_arch(device_index)) with tvm_ffi.use_torch_stream(): @@ -1221,7 +1393,7 @@ def _flat_pool(pool: Any, what: str) -> Any: def _pool_arguments( - ix: Any, *, what: str, H: int, nseg: int, device: str + ix: Any, *, what: str, H: int, nseg: int, device: str, trim: bool = False ) -> tuple[Any, Any, int, int, int | None]: """``(pool view, slot table, hg, n, tile_stride)`` of a compact pool as the kernels address it. @@ -1238,6 +1410,18 @@ def _pool_arguments( 1, None, ) + meta = getattr(ix, "meta", None) if trim else None + H = int(H) + if meta is not None and meta.H == H and meta.device == device: + # the tables' precomputed addressing and device index table (cake_kcp._Tables.indexed_meta) + base_tile, hg, n, stride = meta.addr + if len(meta.index) != nseg: + raise ValueError( + f"{what}: the slot table has {len(meta.index)} entries, the problem {nseg} segments" + ) + pool = _flat_pool(ix.pool, what) + view = pool[int(base_tile) * TILE :] if int(base_tile) else pool + return view, meta.dev_index, int(hg), int(n), int(stride) base_tile, hg, n, stride = ix.addressing(H) index = tuple(int(i) for i in ix.index) if len(index) != nseg: @@ -1260,8 +1444,10 @@ def _local_values(call: Any) -> Mapping[str, Any]: def _walk(name: str, args: Mapping[str, Any], call: Any, device_index: int) -> None: - walk = tuple(int(s) for s in args["walk"]) - if not walk: + trim = int(getattr(call, "host_trim", 0)) == 1 + walk_list = args["walk"] + n_walk = len(walk_list) + if not n_walk: return local = _local_values(call) H, nseg = int(call.H), int(call.tb.NSEG) @@ -1273,21 +1459,37 @@ def _walk(name: str, args: Mapping[str, Any], call: Any, device_index: int) -> N f"{name}: the local stage values carry no {key!r} (needed by the walk)" ) values[key] = local[key] - values.update( - num_heads=H, - num_seqs=nseg, - seg_chunk_start=_device_table( - tuple(int(x) for x in call.tb.hp.seg_chunk_start), "int32", device - ), - walk_seg=_device_table(walk, "int32", device), - ) + if ( + trim + ): # the table set's identity-cached device copies of its own lists (same integers) + values.update( + num_heads=H, + num_seqs=nseg, + seg_chunk_start=call.tb.dev_list(call.tb.hp.seg_chunk_start, device), + walk_seg=call.tb.dev_list(walk_list, device), + ) + else: + walk = tuple(int(s) for s in walk_list) + values.update( + num_heads=H, + num_seqs=nseg, + seg_chunk_start=_device_table( + tuple(int(x) for x in call.tb.hp.seg_chunk_start), "int32", device + ), + walk_seg=_device_table(walk, "int32", device), + ) in_pool, out_pool = _WALK_POOLS[name] stride: int | None = None for key in (in_pool, out_pool): if key is None: continue pool, index, hg, n, pool_stride = _pool_arguments( - args.get(key), what=f"{name}.{key}", H=H, nseg=nseg, device=device + args.get(key), + what=f"{name}.{key}", + H=H, + nseg=nseg, + device=device, + trim=trim, ) values.update( {key: pool, f"{key}_index": index, f"{key}_hg": hg, f"{key}_n": n} @@ -1320,9 +1522,10 @@ def _walk(name: str, args: Mapping[str, Any], call: Any, device_index: int) -> N run_module( record, values, - grid=launch_grid(record, n_walk=len(walk), num_heads=H), + grid=launch_grid(record, n_walk=n_walk, num_heads=H), device_index=device_index, module=module, + trim=trim, ) @@ -1386,17 +1589,26 @@ def _merge(args: Mapping[str, Any], call: Any, device_index: int) -> None: "write_states": int(args["write_states"]), "cols_per_cta": int(MODULES[record]["grid_rule"]["cols_per_cta"]), } - values.update(_chain_tables(chains, device)) + trim = int(getattr(call, "host_trim", 0)) == 1 + if trim and hasattr( + chains, "device_tables" + ): # the chain set's own cached device tables (same integers) + values.update(chains.device_tables(device)) + else: + values.update(_chain_tables(chains, device)) run_module( record, values, grid=launch_grid(record, n_chains=n_chains, ncols=ncols), device_index=device_index, + trim=trim, ) def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> None: - """Launch KCP program ``name`` (one of :data:`PROGRAMS`) with the composition's arguments ``args``. + """Launch KCP program ``name`` (a walk or the merge of :data:`PROGRAMS`) with the composition's arguments + ``args``. The chunk stages ``kcp_fwd_o`` / ``kcp_dav`` take no composition arguments: the host launches + them through :func:`run_module` with their logical values on the grid of :func:`launch_grid`. ``args`` is the vocabulary of ``cake_kcp.Backend.run``: for a walk the segment list ``walk`` (a launch over no segment is skipped), the compact pools ``h0`` / ``dht`` / ``out`` as ``Indexed`` (or ``None``), the transfer @@ -1416,6 +1628,10 @@ def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> Non _merge(args, call, device_index) elif name in WALK_PROGRAMS: _walk(name, args, call, device_index) + elif name in PROGRAMS: + raise ValueError( + f"KCP program {name!r} is a chunk stage: launch it through run_module with its logical values" + ) else: raise ValueError(f"unknown KCP program {name!r}; the programs are {PROGRAMS}") diff --git a/flashinfer/experimental/cake_kda_train/cake_kcp_plan.py b/flashinfer/experimental/cake_kda_train/cake_kcp_plan.py index 4e430859b8d..e615a140f8e 100644 --- a/flashinfer/experimental/cake_kda_train/cake_kcp_plan.py +++ b/flashinfer/experimental/cake_kda_train/cake_kcp_plan.py @@ -52,25 +52,36 @@ segment table) is derived on demand by ``ctx.heads(H)`` and cached. Segment plan rule. ``kcp_plan(n_frag, NC_max_frag, H, sm_count)`` picks the -segments per fragment ``P`` (``NSEG = n_frag * P``) with the shipped ``"wave_fill1"`` -rule (the one-wave fill on top of lever L-C, adopted in this revision): on a rank +segments per fragment ``P`` (``NSEG = n_frag * P``) with the shipped +``"wave_fill1_p1tie"`` rule (the one-wave fill rule plus the P = 1 tie line, on top of the +earlier ``"wave"`` rule): on a rank with at most ``WAVE_N_FRAG_MAX = 3`` fragments (N_MAX) whose longest fragment exceeds the chain cap (``NC_max_frag >= WAVE_MIN_NC_MAX_FRAG = 257`` chunks, i.e. ``P_chain >= 2``), the argmin over ``P in [P_chain, min(64, NC_max_frag)]`` of the equal-segment wave cost ``ceil(n_frag P H / SM) * ceil(NC_max_frag / P) * t_walk + (P - 1) * t_compose`` -(constants ``WAVE_RULE_10NS``, smallest ``P`` on ties; the earlier ``"wave"`` rule -of lever L-C); on a one-fragment rank below the +(constants ``WAVE_RULE_10NS``, smallest ``P`` on ties; the earlier ``"wave"`` rule); on a +one-fragment rank below the cap (``0 < NC_max_frag < 257``) the one-wave fill ``clamp(floor(SM / H), P_chain, min(64, NC_max_frag))`` -- the largest ``P`` whose ``P H`` CTAs run one CTA per SM instead of the occ ``ceil`` (the one-wave fill rule, :data:`WAVE_FILL_N_FRAG_MAX`; a rank with ``H > SM`` keeps the occ ``P``); every other rank (two or three fragments below the cap, more than three fragments, or an empty longest fragment) keeps the -earlier ``"occ"`` rule ``clamp(max(P_occ, P_chain), 1, ...)``. ``"wave"`` (the -earlier rule), ``"occ"`` and ``"wave_fill"`` (the fill on every rank with at most -three fragments; implemented, not adopted) stay selectable through the private -test knob ``_kcp_test_knobs={"p_rule": ...}`` (``CAKE_KCP_TEST_KNOBS=1``) for the -A/B record. Every rule is a pure function of (fragment table, H, SM count). +earlier ``"occ"`` rule ``clamp(max(P_occ, P_chain), 1, ...)`` (the occupancy rule); and -- the +P = 1 tie line -- on a two- or three-fragment rank below the cap whose ``P = 1`` +walks run one wave (``n_frag H <= SM``) with an LPT makespan over the rank's +ACTUAL fragment chunk counts (:func:`p1_tie_steps`, the wave model's waves x steps +generalised to unequal fragments) not above the makespan at the base rule's ``P``, +``P = 1`` (one chain per fragment, no segment composition; :data:`P1_TIE_BASE`). +``"wave_fill1"`` (the previous rule), ``"wave"`` (the earlier rule), ``"occ"`` and +``"wave_fill"`` (the fill on every rank with at most three fragments; implemented, +not adopted) stay selectable through the private test knob +``_kcp_test_knobs={"p_rule": ...}`` (``CAKE_KCP_TEST_KNOBS=1``) for the A/B +record. Every rule is a pure function of (fragment table, H, SM count). A ``P`` +change alters the segment recomposition (not bitwise in general); +for the tie line the bitwise proof set found every changed gate row +bitwise identical between ``P = 2`` and ``P = 1`` on their own operands (the +composition is exact there), re-proven on the merged tree. Exchange layouts (one ``all_to_all_single`` per direction; ``s`` = one FP32 ``[128, 128]`` state of one head = 16384 elements; owner ``j`` holds heads @@ -109,25 +120,26 @@ "wave", "wave_fill1", "wave_fill", -) # kcp_plan P rules: "wave_fill1" = the shipped rule = "wave" + the one-wave -# fill on one-fragment ranks below the chain cap; "wave" = the earlier rule (lever L-C); "occ" = the design note §7.2 rule; -# "wave_fill" = the fill on every rank with n_frag <= 3 (the second knob value; implemented, not adopted) -# Shipped P rule (the one-wave fill on top of the earlier wave rule of lever L-C): -# the public API plans with it; "wave", "occ" and "wave_fill" stay selectable through the private test -# knob ``_kcp_test_knobs={"p_rule": ...}`` for the A/B record. Every rule is a pure function of -# (fragment table, H, SM count), so G3 holds under any of them. -P_RULE_DEFAULT = "wave_fill1" + "wave_fill1_p1tie", +) # kcp_plan P rules: "wave_fill1_p1tie" = the shipped rule = "wave_fill1" + P = 1 +# on two- / three-fragment ranks below the chain cap whose P = 1 LPT walk makespan over the actual fragment chunk counts is <= the +# base rule's (p1_tie_steps; one wave: n_frag H <= SM); "wave_fill1" = the previous shipped rule = "wave" +# + the one-wave fill on one-fragment ranks below the chain cap; "wave" = the earlier rule; "occ" = the design note §7.2 +# rule; "wave_fill" = the fill on every rank with n_frag <= 3 (the second knob value; implemented, not adopted) +# Shipped P rule (the one-wave fill rule plus the P = 1 tie line, on top of the earlier +# wave rule): the public API plans with it; "wave_fill1", "wave", "occ" and +# "wave_fill" stay selectable through the private test knob ``_kcp_test_knobs={"p_rule": ...}`` for the A/B +# record. Every rule is a pure function of (fragment table, H, SM count), so G3 holds under any of them. +P_RULE_DEFAULT = "wave_fill1_p1tie" # shipped; "wave_fill1" = the previous default, via _kcp_test_knobs only # The shipped default of ``save_transfer`` (the design note §6.2) -- the merge owner # keeps its c1 receive tensor (and the home rank its M_seg at P > 1) for the backward instead of re-sending M_f^T in b1. # Default of the byte queries below and of ``cake_kcp.KcpOptions``; ``False`` (the earlier recompute form) stays # selectable through the ``save_transfer`` argument; ``kcp_saved_bytes`` reports the delta. SAVE_TRANSFER_DEFAULT = True -# Lever L-C wave-aware P rule: per-CTA-step constants at one CTA per SM on GB300, H = 64, in 10 ns -# units (exact integer arithmetic keeps the argmin and its ties identical on every rank). Sources (P2A / M1 generated files): -# t_tape 2.85 us = P2A timeline.md "Q1 CP4 C B1" rank 3 #11 / #12 (64 CTAs x 244 steps 0.706 ms -> 2.893; 192 CTAs = 2 waves x 244 -# steps 1.378 ms -> 2.824); t_xfer 3.5 us = M1 k3_micro.md "t_step us (one chain, per wave)" tf32_hilo z-q6 3.515 / c-cp8 3.317; -# t_offset 1.2 us = timeline.md rank 3 #5 (0.581 ms / (2 waves x 244 steps) = 1.19); t_compose 20.4 us = M1 refit.md P_bal table -# "t_compose us/step" 20.389 (K4 map-mode c16 step). +# Wave-aware P rule: per-CTA-step constants at one CTA per SM on GB300, H = 64, in 10 ns units (exact integer arithmetic +# keeps the argmin and its ties identical on every rank), measured on the CP4 / CP8 profiles and the walk / merge +# micro-benchmarks: t_tape 2.85 us (the K1 tape step), t_xfer 3.5 us (the K3 transfer step), t_offset 1.2 us (the K1 +# offset step), t_compose 20.4 us (the K4 map-mode step at 16 columns). WAVE_RULE_10NS = {"tape": 285, "xfer": 350, "offset": 120, "compose": 2040} # N_MAX (a plan constant, part of the pure function of the fragment table, H and SM count): the # wave rule applies to a rank with at most this many fragments; ranks with more keep P_occ (the measured @@ -138,8 +150,8 @@ # fragments of <= 256 chunks (Q1 CP16: 244) the extra segments' compose / staging cost exceeds the shorter walks # (measured CP16 C +5.0 % B1 / +6.4 % B3, CP16 Z +7.3 % B3 at the unconditioned wave P). WAVE_MIN_NC_MAX_FRAG = MAX_CHAIN_STEPS + 1 -# The one-wave fill rules: the "one-wave -# fill" rules -- "wave_fill1" is the shipped default since this revision, "wave_fill" stays a test knob value. Below +# The one-wave fill rules: "wave_fill1_p1tie" is the shipped default ("wave_fill1" plus the P = 1 tie line; "wave_fill1" +# was the previous default), "wave_fill" stays a test knob value. Below # the chain cap (NC_max_frag < WAVE_MIN_NC_MAX_FRAG) the earlier "wave" rule keeps the occ P = ceil(SM / (n_frag H)), which on the one-fragment CP16 C ranks (NC 222-256, H = 64) launches 192 CTAs on # 152 SMs (40 SMs host two CTAs: 148-172 step-times for 74-86-step chains); the fill rule takes the LARGEST P whose # n_frag P H CTAs run one CTA per SM, floor(SM / (n_frag H)) (= 2 -> 128 CTAs, one wave of 111-128 steps, the floor's @@ -148,9 +160,31 @@ # n_frag <= WAVE_N_FRAG_MAX (two-fragment H = 64 ranks move P 2 -> 1). Above the cap both are the shipped wave argmin. # At H <= 4 on 152 SMs floor(SM / H) equals ceil(SM / H) (38; 64-cap at H = 1), so every H <= 4 row keeps its program. WAVE_FILL_N_FRAG_MAX = {"wave_fill1": 1, "wave_fill": WAVE_N_FRAG_MAX} +# The P = 1 tie line of the shipped rule: +# below the chain cap a two-fragment H = 64 rank +# keeps the occ P = ceil(152 / 128) = 2 (256 CTAs: two waves of ceil(NC / 2) steps) although the wave model scores P = 1 +# (128 CTAs, one wave of NC steps) at the same or fewer walk step-times (111 vs 2 x 56 = 112 on Q2 CP16 Z) and P = 2 adds the +# segment-composition K4 launches (map before c1, carry_in / carry_start after c2, tmap / dcarry in the backward) plus the +# backward's offset re-walk. "wave_fill1_p1tie" = "wave_fill1", then P := 1 on a rank with 2 <= n_frag <= WAVE_N_FRAG_MAX, +# 0 < NC_max_frag < WAVE_MIN_NC_MAX_FRAG, P > 1 and n_frag H <= SM (P = 1 runs one wave) when the LPT makespan of the walks +# at P = 1 over the rank's actual fragment chunk counts is <= the makespan at the base rule's P (p1_tie_steps: the wave +# model's waves x steps generalised to unequal fragments through lpt_makespan; equal fragments reproduce it exactly). The +# compose term then favours P = 1 as well. Unequal fragments (CP16 C two-fragment ranks, e.g. [170, 59] chunks) keep P = 2: +# their P = 1 makespan is the long fragment alone. At H <= 4 the base P fills the SMs with short chains, so no H <= 4 rank +# changes. A P change alters the segment recomposition and is not bitwise in general; +# for THIS rule the bitwise proof set found every changed gate row (the CP16 Z, +# CP8 Z and CP4 D / Z rows at H = 64) bitwise identical between P = 2 and P = 1 on their own operands -- the +# segment composition is exact on them -- so the adoption needed no FP64 re-judge; the proof was re-run on the merged tree. +P1_TIE_BASE = { + "wave_fill1_p1tie": "wave_fill1" +} # p1tie rule -> the rule whose P it starts from ALIGN = 256 # every table starts on a 256 B boundary (cake_backend.ALIGN) LAYOUTS = ("contiguous", "zigzag", "zigzag_per_document") EXPLICIT_LAYOUT = "explicit" +# The int32 tables of the KCP saved set, in buffer order: the four chunk tables of the existing chunk-parallel stages and +# the boundary kernels' segment table (``[NSEG + 1]``, their ``seq_chunk_start`` argument). The existing backend's own +# fifth table is the block table of its segmented recurrences (``cake_backend.TABLE_NAMES[4]``, ``seg_table``), which +# the KCP route does not launch. TABLE_NAMES = ( "chunk_bos", "chunk_len", @@ -546,7 +580,8 @@ def kcp_partition( # ---------------------------------------------------------------------------------------------------------------- -# CPU-built tables of the existing chunk-parallel stages (the ``table`` / ``seg_table`` programs' contract) +# CPU-built tables: the four chunk tables of the existing chunk-parallel stages (the ``table`` program's contract) and the +# boundary kernels' segment table def chunk_geometry(T: int, num_seqs: int) -> tuple[int, int, int]: @@ -593,7 +628,8 @@ def chunk_tables(cu_seqlens_local: tuple[int, ...]) -> dict[str, list[int]]: def seg_chunk_start(seq_chunk_start: list[int], P: int) -> list[int]: """``seg_chunk_start[d * P + p] = seq_chunk_start[d] + (p * n_d) // P``, ``[N * P] = NC``: segment ``p`` of local - sequence ``d`` is pseudo-sequence ``d * P + p`` (bitwise the ``kda_chunk_train_seg_table_v1`` output).""" + sequence ``d`` is pseudo-sequence ``d * P + p`` -- the ``[NSEG + 1]`` segment table the boundary kernels (K1-K4) take + as their ``seq_chunk_start`` argument (the form of the earlier segment-table program of the chunk-parallel stages).""" N = len(seq_chunk_start) - 1 P = int(P) if P < 1: @@ -609,8 +645,11 @@ def seg_chunk_start(seq_chunk_start: list[int], P: int) -> list[int]: def table_layout(NC_pad: int, N: int, NSEG: int) -> tuple[tuple[str, int, int], ...]: - """``(name, word offset, words)`` of the five tables inside one int32 buffer, each 256 B aligned (the existing - ``cake_backend._table_layout`` with the segment table always present on the KCP route).""" + """``(name, word offset, words)`` of the five tables inside one int32 buffer, each 256 B aligned: the four chunk + tables as the existing backend lays them out (``cake_backend._table_layout`` without a segment table), then the + boundary kernels' ``[NSEG + 1]`` segment table in the same form, always present on the KCP route. The existing + backend's fifth table (``seg_table``, ``NSEG + 2 * N + 1`` words: the block table of its segmented recurrences) is + another object: the KCP route never launches those recurrences and never hands this buffer to them.""" words_per_align = ALIGN // 4 tables = [ ("chunk_bos", NC_pad), @@ -692,7 +731,7 @@ def needed( def wave_rule_cost_10ns( n_frag: int, NC_max_frag: int, H: int, sm_count: int, P: int ) -> int: - """Lever L-C cost model of ``P`` segments per fragment (10 ns units, exact integers; the rule of + """Cost model of ``P`` segments per fragment (10 ns units, exact integers; the rule of record): the three walks (K1 offset, K3 transfer, K1 tape; one CTA per SM, equal segments) each take ``ceil(n_frag P H / SM)`` waves of ``ceil(NC_max_frag / P)`` CTA-steps, plus ``P - 1`` K4 compose steps.""" waves = _ceil_div(n_frag * P * H, sm_count) @@ -751,6 +790,52 @@ def wave_rule_lpt_P(nc_frags: "list[int]", H: int, sm_count: int) -> int: return best_P +def p1_tie_steps(nc_frags: "list[int]", H: int, sm_count: int, P: int) -> int: + """CTA step-times (LPT makespan at one CTA per SM) of one equal-segment walk of ``P`` segments per fragment over a + rank's fragments of ``nc_frags`` chunks: ``P * H`` chains of ``ceil(NC_f / P)`` steps per non-empty fragment ``f`` on + ``sm_count`` SMs (:func:`lpt_makespan`). Equal fragments reproduce the wave model's ``ceil(n_frag P H / SM) x + ceil(NC / P)`` exactly; unequal fragments are charged their actual makespan (the "wave_fill1_p1tie" tie test).""" + P, H = int(P), int(H) + chains = [(_ceil_div(int(nc), P), P * H) for nc in nc_frags if int(nc) > 0] + return lpt_makespan(chains, int(sm_count)) if chains else 0 + + +def _p1_tie( + n_frag: int, + NC_max_frag: int, + H: int, + sm_count: int, + P: int, + nc_frags: "list[int] | None", +) -> int: + """The "wave_fill1_p1tie" line: ``1`` when the rank has 2..WAVE_N_FRAG_MAX fragments below the chain cap, the base + rule took ``P > 1``, ``P = 1`` runs one wave (``n_frag H <= SM``) and the P = 1 walk makespan over the fragment chunk + counts does not exceed the base ``P``'s; otherwise the base ``P``. ``nc_frags`` = the rank's fragment chunk counts + (``None``: ``n_frag`` equal fragments of ``NC_max_frag`` chunks, the pure (n_frag, NC_max) form of the examples).""" + if nc_frags is None: + frags = [NC_max_frag] * n_frag + else: + frags = [int(x) for x in nc_frags] + if ( + len(frags) != n_frag + or max(frags, default=0) != NC_max_frag + or min(frags, default=0) < 0 + ): + raise ValueError( + f"kcp_plan: nc_frags={frags} must list n_frag={n_frag} fragment chunk counts >= 0 with max NC_max_frag={NC_max_frag}" + ) + if not ( + 2 <= n_frag <= WAVE_N_FRAG_MAX + and 0 < NC_max_frag < WAVE_MIN_NC_MAX_FRAG + and P > 1 + and n_frag * H <= sm_count + ): + return P + if p1_tie_steps(frags, H, sm_count, 1) <= p1_tie_steps(frags, H, sm_count, P): + return 1 + return P + + def kcp_plan( n_frag: int, NC_max_frag: int, @@ -759,9 +844,11 @@ def kcp_plan( *, P_bal: int = P_BAL_INACTIVE, p_rule: str = P_RULE_DEFAULT, + nc_frags: "list[int] | None" = None, ) -> KcpSegmentPlan: - """The KCP plan rule (``p_rule="wave_fill1"``, the shipped default since this revision: the one-wave fill - on the earlier ``"wave"`` rule of lever L-C). The ``"wave"`` rule:: + """The KCP plan rule (``p_rule="wave_fill1_p1tie"``, the shipped default: the one-wave fill on the earlier + ``"wave"`` rule plus the P = 1 tie line; ``"wave_fill1"`` = the previous default, the fill without the tie + line). The ``"wave"`` rule:: P_occ = ceil(SM / (n_frag * H)) # one transfer-walk CTA per SM P_chain = ceil(NC_max_frag / 256) # no walk chain longer than 256 steps @@ -776,11 +863,11 @@ def kcp_plan( walks at one CTA per SM plus the K4 compose chain. The wave line applies only to ranks with ``n_frag <= WAVE_N_FRAG_MAX`` (3) whose longest fragment exceeds the chain cap (``NC_max_frag >= WAVE_MIN_NC_MAX_FRAG`` = 257 chunks, i.e. ``P_chain >= 2``); every other rank keeps the occ ``P`` - (many-fragment D ranks and short fragments measured slower at every ``P > 1``). ``p_rule="occ"`` (the round-0 - §7.2 rule, selectable only through the ``p_rule`` test knob for the A/B record) uses the ``clamp`` line on every + (many-fragment D ranks and short fragments measured slower at every ``P > 1``). ``p_rule="occ"`` (the earlier + §7.2 occupancy rule, selectable only through the ``p_rule`` test knob for the A/B record) uses the ``clamp`` line on every rank. ``P_occ`` / ``P_chain`` / ``P_bal`` are reported under every rule. - The one-wave fill rule (``"wave_fill1"`` shipped, ``"wave_fill"`` test knob only; :data:`WAVE_FILL_N_FRAG_MAX`):: + The one-wave fill rule (``"wave_fill1"`` the previous shipped rule, ``"wave_fill"`` test knob only; :data:`WAVE_FILL_N_FRAG_MAX`):: P_fill = floor(SM / (n_frag * H)) P = clamp(P_fill, P_chain, min(64, NC_max_frag)) if p_rule in ("wave_fill1", "wave_fill") and @@ -791,6 +878,19 @@ def kcp_plan( i.e. below the chain cap the largest segment count that still runs one CTA per SM (``P_bal`` does not enter); ``p_rule="wave"`` (the earlier rule, selectable through the test knob) is the first block alone. + The P = 1 tie rule ``"wave_fill1_p1tie"`` (the shipped rule):: + + P = 1 if p_rule == "wave_fill1_p1tie" and 2 <= n_frag <= WAVE_N_FRAG_MAX and 0 < NC_max_frag < 257 and + P_wave_fill1 > 1 and n_frag H <= SM and + p1_tie_steps(nc_frags, H, SM, 1) <= p1_tie_steps(nc_frags, H, SM, P_wave_fill1) + = the "wave_fill1" rule above otherwise + + with ``nc_frags`` the rank's fragment chunk counts (``None`` = ``n_frag`` equal fragments of ``NC_max_frag`` chunks; + :func:`_build_head_plan` passes the actual counts) and :func:`p1_tie_steps` the LPT makespan of one equal-segment + walk at one CTA per SM -- the wave model's ``waves x steps`` for equal fragments. The segment-composition K4 + launches and the backward's offset re-walk exist only at ``P > 1``; the one-wave P = 1 walk of 128 CTAs runs at a + higher per-step cost than the two-wave 256-CTA walk, which the GPU A/B measures. + A pure function of ``(fragment table, H, SM count)``. A rank without fragments has ``P = 1``, ``NSEG = 0``; when every fragment is empty (``NC_max_frag = 0``) the lower bound wins: ``P = 1``. """ @@ -819,7 +919,10 @@ def kcp_plan( ) p_occ = _ceil_div(sm_count, n_frag * H) p_chain = max(1, _ceil_div(NC_max_frag, MAX_CHAIN_STEPS)) - fill_n_frag_max = WAVE_FILL_N_FRAG_MAX.get(p_rule) + base_rule = P1_TIE_BASE.get( + p_rule, p_rule + ) # "wave_fill1_p1tie" starts from the "wave_fill1" P + fill_n_frag_max = WAVE_FILL_N_FRAG_MAX.get(base_rule) p_fill = sm_count // (n_frag * H) if ( fill_n_frag_max is not None @@ -829,7 +932,7 @@ def kcp_plan( ): P = max(p_chain, min(MAX_SEGMENTS, NC_max_frag, p_fill)) elif ( - p_rule != "occ" + base_rule != "occ" and n_frag <= WAVE_N_FRAG_MAX and NC_max_frag >= WAVE_MIN_NC_MAX_FRAG ): @@ -841,6 +944,12 @@ def kcp_plan( else: upper = min(MAX_SEGMENTS, NC_max_frag, max(p_chain, int(P_bal))) P = max(1, min(max(p_occ, p_chain), upper)) + if p_rule in P1_TIE_BASE: + P = _p1_tie(n_frag, NC_max_frag, H, sm_count, P, nc_frags) + elif nc_frags is not None: + _p1_tie( + n_frag, NC_max_frag, H, sm_count, 1, nc_frags + ) # validation only: the other rules ignore the counts return KcpSegmentPlan( P=P, NSEG=n_frag * P, @@ -1086,7 +1195,7 @@ class KcpHeadPlan: crossing-out fragment -> ``dS_out`` slot; segments ``< P - 1`` -> adjoint carry slot), ``out_index`` (``P = 1``: the c1 send element of a crossing-out fragment, else ``-1``; ``P > 1``: the staging slot ``= segment``), ``walk_seg_offset`` (segments the offset walks visit) and ``walk_seg_tape`` (all). ``h0_index_fused`` / - ``dht_index_fused`` are the one-pool form of the same tables for lever L-B (``fused_walks``; the shipped + ``dht_index_fused`` are the one-pool form of the same tables for the fused walks (``fused_walks``; the shipped form since an earlier revision; ``fused_walks=0`` through the test knob): at ``P > 1`` every segment's initial state (final adjoint) sits in the carry plane -- segment 0 of a crossing-in fragment in the extra slot ``n_frag (P - 1) + in_slot`` that the ``carry_in`` merge fills with its ``S_in`` tiles (``write_states``), the last segment of a crossing-out fragment in ``n_frag (P - 1) + out_slot`` @@ -1313,8 +1422,11 @@ def b2_send_elem(dst: int, out_slot: int) -> int: ) # -- segment plan and the K1 / K2 / K3 segment index tables - nc_max_frag = max((_ceil_div(f.length, CHUNK) for f in mine), default=0) - seg = kcp_plan(ctx.n_frag, nc_max_frag, H, ctx.sm_count, p_rule=ctx.p_rule) + nc_frags = [_ceil_div(f.length, CHUNK) for f in mine] + nc_max_frag = max(nc_frags, default=0) + seg = kcp_plan( + ctx.n_frag, nc_max_frag, H, ctx.sm_count, p_rule=ctx.p_rule, nc_frags=nc_frags + ) P, NSEG = seg.P, seg.NSEG seg_table = ( tuple(seg_chunk_start(list(ctx.seq_chunk_start), P)) if ctx.n_frag else (0,) @@ -1344,7 +1456,7 @@ def b2_send_elem(dst: int, out_slot: int) -> int: out_index.append(s) walk_offset.append(s) walk_tape = list(range(NSEG)) - # the one-pool tables of lever L-B (fused_walks; the shipped form since an earlier revision, fused_walks=0 via the knob) + # the one-pool tables of the fused walks (fused_walks; the shipped form since an earlier revision, fused_walks=0 via the knob) extra = ctx.n_frag * (P - 1) h0_fused, dht_fused = [], [] for f in mine: @@ -1448,7 +1560,7 @@ def b2_send_elem(dst: int, out_slot: int) -> int: numel["carry"] = (ctx.n_frag * (P - 1) * H * TILE) if P > 1 else 0 numel["carry_fused"] = ( ((extra + max(n_in, n_out)) * H * TILE) if P > 1 else 0 - ) # lever L-B's carry plane + ) # the fused walks' carry plane dev = ctx.device i32, i64 = torch.int32, torch.int64 @@ -1620,7 +1732,7 @@ class KcpContext: tables: dict[str, torch.Tensor] halo: KcpHaloPlan halo_tables: dict[str, torch.Tensor] - p_rule: str = P_RULE_DEFAULT # plan rule: "wave_fill1" shipped; "wave" / "occ" / "wave_fill" via the test knob + p_rule: str = P_RULE_DEFAULT # plan rule: "wave_fill1_p1tie" shipped; "wave_fill1" (previous) / "wave" / "occ" / "wave_fill" via the test knob sequence: KcpSequence = dataclasses.field(default_factory=KcpSequence, repr=False) _head_plans: dict[int, KcpHeadPlan] = dataclasses.field( default_factory=dict, repr=False @@ -1739,9 +1851,11 @@ def build_kcp_context( Out-of-range problems are rejected here, before any launch: ``n_frag > 65536`` (the layout table's document search), int32 overflows of the kernels' offsets (``heads(H)``). ``_kcp_test_knobs`` is the private test hook (SPEC-11; accepted only with ``CAKE_KCP_TEST_KNOBS=1``): its ``p_rule`` selects the plan's ``P`` rule for this - context (:data:`P_RULE_DEFAULT` = ``"wave_fill1"`` without the knob; ``"wave"`` = the earlier rule, ``"occ"`` = the - earlier rule, ``"wave_fill"`` = the unadopted fill on every rank with at most three fragments, kept for the A/B - record); the public API never selects a rule other than the default. + context (:data:`P_RULE_DEFAULT` = ``"wave_fill1_p1tie"`` without the knob; + ``"wave_fill1"`` = the previous rule (the fill without the tie line), ``"wave"`` = the equal-segment wave rule + above the chain cap alone, ``"occ"`` = the occupancy rule, ``"wave_fill"`` = the unadopted fill on every rank with + at most three fragments, kept for the A/B record); the + public API never selects a rule other than the default. """ if not isinstance(part, KcpPartition): raise TypeError( @@ -1900,7 +2014,7 @@ def kcp_saved_bytes( ) -> int: """Saved-set bytes of one layer on this rank (the KCP contract §3.2 ``kcp_saved_bytes(ctx, ...)``):: - 256 * T_r * H + tables_bytes(NC_pad, n_frag, NSEG) # Aqk / Akk on the token rows + the int32 tables + tables_bytes(NC_pad, n_frag, NSEG) # the int32 tables + n_in * s # S_in, s = H * 128 * 128 * 4 + [save_transfer] * N_out * 2 * s * Hg_r / H # the owner's whole c1 receive tensor + [save_transfer] * [P > 1] * NSEG * s # M_seg of the segmented fragments (home rank) @@ -1910,6 +2024,9 @@ def kcp_saved_bytes( ``save_transfer`` defaults to :data:`SAVE_TRANSFER_DEFAULT` (``True`` since an earlier revision); the two ``[save_transfer]`` terms are the delta against the earlier recompute form (``save_transfer=False``). + The set carries no token-row ``Aqk`` / ``Akk`` (the ``256 * T_r * H`` term of its earlier version): the backward + recomputes the internal rows in both modes. + The retained ``S_in`` and c1 tensors are the padded receive buffers (``s`` -> ``n_owners Hg`` tiles, ``Hg_r`` -> ``Hg_pad_r``); identical to the formula for ``H = 64``, and the padding is reported by ``kcp_comm_bytes(...)["padding"]``. ``M_seg`` stays the dense copy of §9 (``NSEG s``). @@ -1919,7 +2036,7 @@ def kcp_saved_bytes( hp = ctx.heads(H) s = H * STATE_TILE_BYTES s_pad = hp.n_owners * hp.Hg * STATE_TILE_BYTES # = s without padding - total = 256 * ctx.T_r * H + tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) + total = tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) total += ctx.n_in * s_pad if save_transfer: total += ctx.N_out * 2 * hp.Hg_pad_by_rank[ctx.rank] * STATE_TILE_BYTES @@ -2070,6 +2187,32 @@ def kcp_arena_layout( return out +# The paths of the existing backend's plan the KCP's local stages take their records from: the forward without a +# saved set (its records read the internal ``Aqk`` / ``Akk`` rows the intra program writes; the output program is +# the KCP's own) and the checkpoint replay for both backward modes (its prep keeps the internal ``dout`` copy K2 and +# the KCP's ``dav`` read; its WY and fused-gradient records read the recomputed internal rows). +LOCAL_PATH_FWD = "fwd_nosave" +LOCAL_PATH_BWD = "replay" + + +def intra_writes_tables(plan, path: str = LOCAL_PATH_FWD) -> bool: + """Whether the existing backend's intra record of ``path`` (a :class:`cake_backend.Plan`) computes the chunk + layout from ``cu_seqlens`` and writes the tables itself (its logical vocabulary binds ``cu_seqlens``): the kernel + backend then hands it a scratch table set so the KCP's CPU-built tables are never rewritten by a kernel.""" + from . import ( + cake_jit, + ) # the registry literal; imported here so the metadata module stays torch-only + + record = plan.record("fwd_intra", path) + return "cu_seqlens" in cake_jit.MODULES[record]["logical"].values() + + +def local_scratch_table_bytes(plan) -> int: + """Bytes of the scratch table set of :func:`intra_writes_tables` on the arena of ``plan`` (its table buffer: the + four chunk tables and the block table, 256 B aligned), zero where the intra record reads the tables instead.""" + return int(plan.table_words()) * 4 if intra_writes_tables(plan) else 0 + + def kcp_workspace_bytes( ctx: KcpContext, *, @@ -2080,7 +2223,8 @@ def kcp_workspace_bytes( C: int = 8192, ) -> int: """Caller-held workspace bytes of one rank: the existing backend's local arena for ``(T_r, n_frag, H)`` - (``kda_train_workspace_bytes``; zero on a rank without tokens) plus the KCP arena of transient buffers + (``kda_train_workspace_bytes``; zero on a rank without tokens) with the intra record's scratch table set where + it writes one (:func:`local_scratch_table_bytes`), plus the KCP arena of transient buffers (:func:`kcp_arena_layout`; CONSISTENCY-6, SPEC-4: saved-set tensors are counted by :func:`kcp_saved_bytes`).""" _check_save_mode(save_mode) local = 0 @@ -2099,6 +2243,16 @@ def kcp_workspace_bytes( sm_count=ctx.sm_count, ) ) + # the scratch table set where the intra record writes one (the rule is architecture-independent; the + # buffer size follows the same segment-count convention as the arena bytes above) + if intra_writes_tables( + cake_backend.plan(ctx.T_r, ctx.n_frag, int(H), ctx.sm_count) + ): + local += int( + cake_backend.tables_bytes( + ctx.T_r, ctx.n_frag, int(H), sm_count=ctx.sm_count + ) + ) arena = kcp_arena_layout( ctx, H=H, @@ -2116,11 +2270,12 @@ def kcp_workspace_bytes( def validate_test_knobs(knobs) -> dict: """Normalize ``_kcp_test_knobs``: ``None`` -> the public defaults ``{"precision_mode": 0, "p_rule": - P_RULE_DEFAULT}`` (``"wave_fill1"``); a dict is accepted only when the environment variable ``CAKE_KCP_TEST_KNOBS=1`` + P_RULE_DEFAULT}`` (``"wave_fill1_p1tie"``); a dict is accepted only when the environment variable ``CAKE_KCP_TEST_KNOBS=1`` is set, holds only known keys and selects a known ``precision_mode`` (0 shipped TF32 hi/lo; 1 the explicit three-MMA form of the R1 bitwise test; 2 plain TF32, the optional G2 cross-check) and a known ``p_rule`` - (:data:`P_RULES`; ``"wave"`` = the earlier rule, ``"occ"`` = the occupancy rule, ``"wave_fill"`` = the unadopted - fill, read by :func:`build_kcp_context`). The public API can + (:data:`P_RULES`; ``"wave_fill1"`` = the previous rule (the fill without the tie line), ``"wave"`` = the equal-segment + wave rule above the chain cap alone, ``"occ"`` = the occupancy rule, ``"wave_fill"`` = the unadopted fill; read by + :func:`build_kcp_context`). The public API can therefore never select a non-zero mode or a P rule other than the shipped one.""" if knobs is None: return {"precision_mode": 0, "p_rule": P_RULE_DEFAULT} @@ -2198,11 +2353,13 @@ def validate_test_knobs(knobs) -> dict: "table_words", "tables_bytes", "validate_test_knobs", + "P1_TIE_BASE", "WAVE_FILL_N_FRAG_MAX", "WAVE_MIN_NC_MAX_FRAG", "WAVE_N_FRAG_MAX", "WAVE_RULE_10NS", "lpt_makespan", + "p1_tie_steps", "wave_rule_cost_10ns", "wave_rule_cost_us", "wave_rule_lpt_cost_10ns", diff --git a/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu new file mode 100644 index 00000000000..5757360f91c --- /dev/null +++ b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_binding.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_train_1214bfc28e50985bfa0c, for the device source that defines it. Do not edit. +#include +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_kda_train_1214bfc28e50985bfa0c(const __grid_constant__ CUtensorMap qg_tma, const __grid_constant__ CUtensorMap h_tma, const __grid_constant__ CUtensorMap aqk_tma, const __grid_constant__ CUtensorMap vn_tma, __nv_bfloat16* __restrict__ o_out, int* __restrict__ chunk_bos, int* __restrict__ chunk_len, int* __restrict__ seq_chunk_start, int num_seqs, int num_heads, float scale); + +namespace cake_binding_kernel_cake_kda_train_1214bfc28e50985bfa0c { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_kda_train_1214bfc28e50985bfa0c"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +// The product of the extents above the trailing `trailing` dims. +int64_t NumelAbove(const TensorView& t, int trailing) { + int64_t product = 1; + for (int axis = 0; axis < t.ndim() - trailing; ++axis) product *= t.size(axis); + return product; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_qg_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'qg_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "qg_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "qg_tma"), Extent(d2, "qg_tma"), Extent(outer2, "qg_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "qg_tma"), Stride(((d1 * d2) * 2), "qg_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'qg_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'qg_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'qg_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of qg_tma"); + return map; +} + +CUtensorMap EncodeTma_h_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'h_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "h_tma"); + int64_t d1 = t.size(-1); + int64_t outer1 = NumelAbove(t, 1); + check_dense_leading_dims(t, 1, "h_tma"); + int64_t s2 = t.stride(-2) * 2; + cuuint64_t global_dim[2] = {Extent(d1, "h_tma"), Extent(outer1, "h_tma")}; + cuuint64_t global_strides[1] = {Stride(s2, "h_tma")}; + cuuint32_t box_dim[2] = {64u, 128u}; + cuuint32_t elem_strides[2] = {1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1], ValueError) + << KERNEL_NAME << ": the box of descriptor 'h_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'h_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'h_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of h_tma"); + return map; +} + +CUtensorMap EncodeTma_aqk_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'aqk_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "aqk_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "aqk_tma"), Extent(d2, "aqk_tma"), Extent(outer2, "aqk_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "aqk_tma"), Stride(((d1 * d2) * 2), "aqk_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'aqk_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'aqk_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'aqk_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of aqk_tma"); + return map; +} + +CUtensorMap EncodeTma_vn_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'vn_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "vn_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "vn_tma"), Extent(d2, "vn_tma"), Extent(outer2, "vn_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "vn_tma"), Stride(((d1 * d2) * 2), "vn_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'vn_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'vn_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'vn_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of vn_tma"); + return map; +} + +void Run(TensorView a_qg_tma, TensorView a_h_tma, TensorView a_aqk_tma, TensorView a_vn_tma, TensorView a_o_out, TensorView a_chunk_bos, TensorView a_chunk_len, TensorView a_seq_chunk_start, int64_t a_num_seqs, int64_t a_num_heads, double a_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_qg_tma.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_qg_tma, "qg_tma"); + check_dtype(a_qg_tma, DLDataType{kDLBfloat, 16, 1}, "qg_tma"); + check_cuda_tensor(a_h_tma, "h_tma"); + check_dtype(a_h_tma, DLDataType{kDLBfloat, 16, 1}, "h_tma"); + check_cuda_tensor(a_aqk_tma, "aqk_tma"); + check_dtype(a_aqk_tma, DLDataType{kDLBfloat, 16, 1}, "aqk_tma"); + check_cuda_tensor(a_vn_tma, "vn_tma"); + check_dtype(a_vn_tma, DLDataType{kDLBfloat, 16, 1}, "vn_tma"); + check_cuda_tensor(a_o_out, "o_out"); + check_dtype(a_o_out, DLDataType{kDLBfloat, 16, 1}, "o_out"); + check_cuda_tensor(a_chunk_bos, "chunk_bos"); + check_dtype(a_chunk_bos, DLDataType{kDLInt, 32, 1}, "chunk_bos"); + check_cuda_tensor(a_chunk_len, "chunk_len"); + check_dtype(a_chunk_len, DLDataType{kDLInt, 32, 1}, "chunk_len"); + check_cuda_tensor(a_seq_chunk_start, "seq_chunk_start"); + check_dtype(a_seq_chunk_start, DLDataType{kDLInt, 32, 1}, "seq_chunk_start"); + TVM_FFI_CHECK(a_num_seqs >= -2147483648LL && a_num_seqs <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_seqs' = " << a_num_seqs + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_num_heads >= -2147483648LL && a_num_heads <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_heads' = " << a_num_heads + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_qg_tma, "qg_tma"); + check_same_device(a_h_tma, a_qg_tma, "h_tma", "qg_tma"); + check_contiguous(a_aqk_tma, "aqk_tma"); + check_same_device(a_aqk_tma, a_qg_tma, "aqk_tma", "qg_tma"); + check_contiguous(a_vn_tma, "vn_tma"); + check_same_device(a_vn_tma, a_qg_tma, "vn_tma", "qg_tma"); + check_contiguous(a_o_out, "o_out"); + check_same_device(a_o_out, a_qg_tma, "o_out", "qg_tma"); + check_contiguous(a_chunk_bos, "chunk_bos"); + check_same_device(a_chunk_bos, a_qg_tma, "chunk_bos", "qg_tma"); + check_contiguous(a_chunk_len, "chunk_len"); + check_same_device(a_chunk_len, a_qg_tma, "chunk_len", "qg_tma"); + check_contiguous(a_seq_chunk_start, "seq_chunk_start"); + check_same_device(a_seq_chunk_start, a_qg_tma, "seq_chunk_start", "qg_tma"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_train_1214bfc28e50985bfa0c), cudaFuncAttributeMaxDynamicSharedMemorySize, 222208), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_qg_tma = EncodeTma_qg_tma(a_qg_tma); + CUtensorMap v_h_tma = EncodeTma_h_tma(a_h_tma); + CUtensorMap v_aqk_tma = EncodeTma_aqk_tma(a_aqk_tma); + CUtensorMap v_vn_tma = EncodeTma_vn_tma(a_vn_tma); + void* v_o_out = DataPtr(a_o_out); + void* v_chunk_bos = DataPtr(a_chunk_bos); + void* v_chunk_len = DataPtr(a_chunk_len); + void* v_seq_chunk_start = DataPtr(a_seq_chunk_start); + int32_t v_num_seqs = static_cast(a_num_seqs); + int32_t v_num_heads = static_cast(a_num_heads); + float v_scale = static_cast(a_scale); + void* kernel_args[] = {&v_qg_tma, &v_h_tma, &v_aqk_tma, &v_vn_tma, &v_o_out, &v_chunk_bos, &v_chunk_len, &v_seq_chunk_start, &v_num_seqs, &v_num_heads, &v_scale}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(192, 1, 1); + config.dynamicSmemBytes = 222208; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_train_1214bfc28e50985bfa0c), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_kda_train_1214bfc28e50985bfa0c + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_train_1214bfc28e50985bfa0c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu new file mode 100644 index 00000000000..e3d9f769741 --- /dev/null +++ b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_1214bfc28e50985bfa0c_kernel.cu @@ -0,0 +1,646 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_ACC_QH_OFFSET 0 +#define TMEM_ACC_AV_OFFSET 256 +#define NUM_RING_STAGES 3 +#define NUM_ACCS_STAGES 2 +#define SMEM_QG_S_OFF 1024 +#define SMEM_QG_S_STAGE_BYTES 16384 +#define SMEM_QG_S_STRIDE 73728 +#define SMEM_H_S_OFF 17408 +#define SMEM_H_S_STAGE_BYTES 32768 +#define SMEM_H_S_STRIDE 73728 +#define SMEM_AQK_P_OFF 50176 +#define SMEM_AQK_P_STAGE_BYTES 8192 +#define SMEM_AQK_P_STRIDE 73728 +#define SMEM_VN_S_OFF 58368 +#define SMEM_VN_S_STAGE_BYTES 16384 +#define SMEM_VN_S_STRIDE 73728 +#define SMEM_QG_K_OFF 1024 +#define SMEM_QG_K_STAGE_BYTES 16384 +#define SMEM_QG_K_STRIDE 73728 +#define SMEM_H_MN_OFF 17408 +#define SMEM_H_MN_STAGE_BYTES 32768 +#define SMEM_H_MN_STRIDE 73728 +#define SMEM_AQK_K_OFF 50176 +#define SMEM_AQK_K_STAGE_BYTES 8192 +#define SMEM_AQK_K_STRIDE 73728 +#define SMEM_VN_MN_OFF 58368 +#define SMEM_VN_MN_STAGE_BYTES 16384 +#define SMEM_VN_MN_STRIDE 73728 +#define SMEM_TOTAL 222208 +#define THREADS 192 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + + + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(192, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_train_1214bfc28e50985bfa0c(const __grid_constant__ CUtensorMap qg_tma, const __grid_constant__ CUtensorMap h_tma, const __grid_constant__ CUtensorMap aqk_tma, const __grid_constant__ CUtensorMap vn_tma, __nv_bfloat16* __restrict__ o_out, int* __restrict__ chunk_bos, int* __restrict__ chunk_len, int* __restrict__ seq_chunk_start, int num_seqs, int num_heads, float scale) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define tiles_full_addr (mbar_base + 0) + #define tiles_empty_addr (mbar_base + 24) + #define acc_full_addr (mbar_base + 48) + #define acc_empty_addr (mbar_base + 64) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* qg_s = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int qg_s_addr = smem + 1024; + __nv_bfloat16* h_s = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int h_s_addr = smem + 17408; + __nv_bfloat16* aqk_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int aqk_p_addr = smem + 50176; + __nv_bfloat16* vn_s = reinterpret_cast<__nv_bfloat16*>(smem_raw + 58368); + const int vn_s_addr = smem + 58368; + __nv_bfloat16* qg_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int qg_k_addr = smem + 1024; + __nv_bfloat16* h_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int h_mn_addr = smem + 17408; + __nv_bfloat16* aqk_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int aqk_k_addr = smem + 50176; + __nv_bfloat16* vn_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 58368); + const int vn_mn_addr = smem + 58368; + + // Mbarrier init (4 pipeline groups, 0 ordered-sequence groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'ring' --- + // tiles_full: 3 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 1); + mbarrier_init(smem + 16, 1); + // tiles_empty: 3 barriers, init_count=1 + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // --- pipeline 'accs' --- + // acc_full: 2 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // acc_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 64, 128); + mbarrier_init(smem + 72, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_acc_qh = taddr; + const int tmem_acc_av = taddr + 256; + + // ---- Role: epilogue ---- + if (warp <= 3) { + { // epilogue_main + unsigned int epi_buf = 0; + int nc_e = seq_chunk_start[num_seqs]; + int num_tiles_e = nc_e * num_heads; + int wq = warp; + int lane_half = lane >> 4; + int t_row = wq * 16 + (lane & 15); + int row_base = wq * 32 << 16; + unsigned int _phase_acc_full = 0; + #pragma unroll 1 + for (unsigned int tile = bid; tile < num_tiles_e; tile += num_bids) { + int t_i = (int)tile; + int chunk = nc_e - 1 - t_i / num_heads; + int head = t_i - t_i / num_heads * num_heads; + int ext0 = chunk_bos[chunk]; + int clen = chunk_len[chunk]; + long long obase = ((long long)(ext0 + t_row) * (long long)num_heads + (long long)head) * 128; + int buf_col = (int)epi_buf * 128; + mbarrier_wait(acc_full_addr + (epi_buf) * 8, _phase_acc_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int col = lane_half * 16; + float _tmem_load_0[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])) + : "r"(taddr + (unsigned int)row_base + (unsigned int)buf_col)); + float _tmem_load_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])) + : "r"(taddr + 256 + (unsigned int)row_base + (unsigned int)buf_col)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float ov[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + ov[j] = __fmaf_rn(_tmem_load_0[j], scale, _tmem_load_1[j]); + } + if (t_row < clen) { + { + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(ov[0 + 0], ov[0 + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(ov[0 + 2], ov[0 + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(ov[0 + 4], ov[0 + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(ov[0 + 6], ov[0 + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(ov[0 + 8], ov[0 + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(ov[0 + 10], ov[0 + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(ov[0 + 12], ov[0 + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(ov[0 + 14], ov[0 + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(o_out + (obase + (long long)col)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + int col_0 = 32 + lane_half * 16; + float _tmem_load_2[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])) + : "r"(taddr + (unsigned int)row_base + (unsigned int)(buf_col + 32))); + float _tmem_load_3[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_3[0])), "=r"(*reinterpret_cast(&_tmem_load_3[1])), "=r"(*reinterpret_cast(&_tmem_load_3[2])), "=r"(*reinterpret_cast(&_tmem_load_3[3])), "=r"(*reinterpret_cast(&_tmem_load_3[4])), "=r"(*reinterpret_cast(&_tmem_load_3[5])), "=r"(*reinterpret_cast(&_tmem_load_3[6])), "=r"(*reinterpret_cast(&_tmem_load_3[7])), "=r"(*reinterpret_cast(&_tmem_load_3[8])), "=r"(*reinterpret_cast(&_tmem_load_3[9])), "=r"(*reinterpret_cast(&_tmem_load_3[10])), "=r"(*reinterpret_cast(&_tmem_load_3[11])), "=r"(*reinterpret_cast(&_tmem_load_3[12])), "=r"(*reinterpret_cast(&_tmem_load_3[13])), "=r"(*reinterpret_cast(&_tmem_load_3[14])), "=r"(*reinterpret_cast(&_tmem_load_3[15])) + : "r"(taddr + 256 + (unsigned int)row_base + (unsigned int)(buf_col + 32))); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float ov_1[16]; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ov_1[j_1] = __fmaf_rn(_tmem_load_2[j_1], scale, _tmem_load_3[j_1]); + } + if (t_row < clen) { + { + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(ov_1[0 + 0], ov_1[0 + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(ov_1[0 + 2], ov_1[0 + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(ov_1[0 + 4], ov_1[0 + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(ov_1[0 + 6], ov_1[0 + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(ov_1[0 + 8], ov_1[0 + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(ov_1[0 + 10], ov_1[0 + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(ov_1[0 + 12], ov_1[0 + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(ov_1[0 + 14], ov_1[0 + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(o_out + (obase + (long long)col_0)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + int col_2 = 64 + lane_half * 16; + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + (unsigned int)row_base + (unsigned int)(buf_col + 64))); + float _tmem_load_5[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_5[0])), "=r"(*reinterpret_cast(&_tmem_load_5[1])), "=r"(*reinterpret_cast(&_tmem_load_5[2])), "=r"(*reinterpret_cast(&_tmem_load_5[3])), "=r"(*reinterpret_cast(&_tmem_load_5[4])), "=r"(*reinterpret_cast(&_tmem_load_5[5])), "=r"(*reinterpret_cast(&_tmem_load_5[6])), "=r"(*reinterpret_cast(&_tmem_load_5[7])), "=r"(*reinterpret_cast(&_tmem_load_5[8])), "=r"(*reinterpret_cast(&_tmem_load_5[9])), "=r"(*reinterpret_cast(&_tmem_load_5[10])), "=r"(*reinterpret_cast(&_tmem_load_5[11])), "=r"(*reinterpret_cast(&_tmem_load_5[12])), "=r"(*reinterpret_cast(&_tmem_load_5[13])), "=r"(*reinterpret_cast(&_tmem_load_5[14])), "=r"(*reinterpret_cast(&_tmem_load_5[15])) + : "r"(taddr + 256 + (unsigned int)row_base + (unsigned int)(buf_col + 64))); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float ov_3[16]; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + ov_3[j_2] = __fmaf_rn(_tmem_load_4[j_2], scale, _tmem_load_5[j_2]); + } + if (t_row < clen) { + { + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(ov_3[0 + 0], ov_3[0 + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(ov_3[0 + 2], ov_3[0 + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(ov_3[0 + 4], ov_3[0 + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(ov_3[0 + 6], ov_3[0 + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(ov_3[0 + 8], ov_3[0 + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(ov_3[0 + 10], ov_3[0 + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(ov_3[0 + 12], ov_3[0 + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(ov_3[0 + 14], ov_3[0 + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(o_out + (obase + (long long)col_2)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + int col_4 = 96 + lane_half * 16; + float _tmem_load_6[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_6[0])), "=r"(*reinterpret_cast(&_tmem_load_6[1])), "=r"(*reinterpret_cast(&_tmem_load_6[2])), "=r"(*reinterpret_cast(&_tmem_load_6[3])), "=r"(*reinterpret_cast(&_tmem_load_6[4])), "=r"(*reinterpret_cast(&_tmem_load_6[5])), "=r"(*reinterpret_cast(&_tmem_load_6[6])), "=r"(*reinterpret_cast(&_tmem_load_6[7])), "=r"(*reinterpret_cast(&_tmem_load_6[8])), "=r"(*reinterpret_cast(&_tmem_load_6[9])), "=r"(*reinterpret_cast(&_tmem_load_6[10])), "=r"(*reinterpret_cast(&_tmem_load_6[11])), "=r"(*reinterpret_cast(&_tmem_load_6[12])), "=r"(*reinterpret_cast(&_tmem_load_6[13])), "=r"(*reinterpret_cast(&_tmem_load_6[14])), "=r"(*reinterpret_cast(&_tmem_load_6[15])) + : "r"(taddr + (unsigned int)row_base + (unsigned int)(buf_col + 96))); + float _tmem_load_7[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&_tmem_load_7[0])), "=r"(*reinterpret_cast(&_tmem_load_7[1])), "=r"(*reinterpret_cast(&_tmem_load_7[2])), "=r"(*reinterpret_cast(&_tmem_load_7[3])), "=r"(*reinterpret_cast(&_tmem_load_7[4])), "=r"(*reinterpret_cast(&_tmem_load_7[5])), "=r"(*reinterpret_cast(&_tmem_load_7[6])), "=r"(*reinterpret_cast(&_tmem_load_7[7])), "=r"(*reinterpret_cast(&_tmem_load_7[8])), "=r"(*reinterpret_cast(&_tmem_load_7[9])), "=r"(*reinterpret_cast(&_tmem_load_7[10])), "=r"(*reinterpret_cast(&_tmem_load_7[11])), "=r"(*reinterpret_cast(&_tmem_load_7[12])), "=r"(*reinterpret_cast(&_tmem_load_7[13])), "=r"(*reinterpret_cast(&_tmem_load_7[14])), "=r"(*reinterpret_cast(&_tmem_load_7[15])) + : "r"(taddr + 256 + (unsigned int)row_base + (unsigned int)(buf_col + 96))); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float ov_5[16]; + #pragma unroll + for (int j_3 = 0; j_3 < 16; j_3++) { + ov_5[j_3] = __fmaf_rn(_tmem_load_6[j_3], scale, _tmem_load_7[j_3]); + } + if (t_row < clen) { + { + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(ov_5[0 + 0], ov_5[0 + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(ov_5[0 + 2], ov_5[0 + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(ov_5[0 + 4], ov_5[0 + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(ov_5[0 + 6], ov_5[0 + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(ov_5[0 + 8], ov_5[0 + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(ov_5[0 + 10], ov_5[0 + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(ov_5[0 + 12], ov_5[0 + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(ov_5[0 + 14], ov_5[0 + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(o_out + (obase + (long long)col_4)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(acc_empty_addr + (epi_buf) * 8); + epi_buf += 1; + if (epi_buf == 2) { epi_buf = 0; _phase_acc_full ^= 1; } + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + if (warp == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: load ---- + if (warp == 4) { + { // load_main + unsigned int ld_stage = 0; + int nc_l = seq_chunk_start[num_seqs]; + int num_tiles_l = nc_l * num_heads; + unsigned int _phase_tiles_empty = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_1 = bid; tile_1 < num_tiles_l; tile_1 += num_bids) { + int t_i_1 = (int)tile_1; + int chunk_1 = nc_l - 1 - t_i_1 / num_heads; + int head_1 = t_i_1 - t_i_1 / num_heads * num_heads; + int row0 = chunk_1 * 64; + int hrow0 = (chunk_1 * num_heads + head_1) * 128; + int ext0_1 = chunk_bos[chunk_1]; + mbarrier_wait(tiles_empty_addr + (ld_stage) * 8, _phase_tiles_empty); + mbarrier_arrive_expect_tx(tiles_full_addr + (ld_stage) * 8, 73728); + #pragma unroll + for (int seg = 0; seg < 2; seg++) { + int c0 = seg * 64; + tma_3d_gmem2smem(qg_s_addr + ld_stage * 73728 + (unsigned int)(seg * 8192), (&qg_tma), c0, head_1, row0, tiles_full_addr + (ld_stage) * 8); + tma_3d_gmem2smem(vn_s_addr + ld_stage * 73728 + (unsigned int)(seg * 8192), (&vn_tma), c0, head_1, row0, tiles_full_addr + (ld_stage) * 8); + tma_2d_gmem2smem(h_s_addr + ld_stage * 73728 + (unsigned int)(seg * 16384), (&h_tma), c0, hrow0, tiles_full_addr + (ld_stage) * 8); + } + tma_3d_gmem2smem(aqk_p_addr + ld_stage * 73728, (&aqk_tma), 0, head_1, row0, tiles_full_addr + (ld_stage) * 8); + ld_stage += 1; + if (ld_stage == 3) { ld_stage = 0; _phase_tiles_empty ^= 1; } + } + } + } + } + // ---- Role: mma ---- + if (warp == 5) { + { // mma_main + unsigned int mma_stage = 0; + unsigned int mma_buf = 0; + int nc_m = seq_chunk_start[num_seqs]; + int num_tiles_m = nc_m * num_heads; + unsigned int _phase_tiles_full = 0; + unsigned int _phase_acc_empty = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int _tile = bid; _tile < num_tiles_m; _tile += num_bids) { + mbarrier_wait(tiles_full_addr + (mma_stage) * 8, _phase_tiles_full); + mbarrier_wait(acc_empty_addr + (mma_buf) * 8, _phase_acc_empty); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = (((qg_k_addr) >> 4) & 0x7FFF) + (mma_stage) * 4608; + int _smem_desc_lo_1 = ((((h_mn_addr) >> 4) & 0x7FFF) | 0x4000000) + (mma_stage) * 4608; +#else + int _smem_desc_lo_0 = (((qg_k_addr) >> 4) & 0x3FFF) + (mma_stage) * 4608; + int _smem_desc_lo_1 = ((((h_mn_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_stage) * 4608; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"((tmem_acc_qh + (mma_buf * 128))), "r"(0) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = (((aqk_k_addr) >> 4) & 0x7FFF) + (mma_stage) * 4608; + int _smem_desc_lo_3 = ((((vn_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (mma_stage) * 4608; +#else + int _smem_desc_lo_2 = (((aqk_k_addr) >> 4) & 0x3FFF) + (mma_stage) * 4608; + int _smem_desc_lo_3 = ((((vn_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (mma_stage) * 4608; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"((tmem_acc_av + (mma_buf * 128))), "r"(0) : "memory"); + tcgen05_commit(tiles_empty_addr + (mma_stage) * 8); + tcgen05_commit(acc_full_addr + (mma_buf) * 8); + mma_stage += 1; + if (mma_stage == 3) { mma_stage = 0; _phase_tiles_full ^= 1; } + mma_buf += 1; + if (mma_buf == 2) { mma_buf = 0; _phase_acc_empty ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu new file mode 100644 index 00000000000..a9a8a573e74 --- /dev/null +++ b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_binding.cu @@ -0,0 +1,215 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_train_f227dd21f1f3cb464b29, for the device source that defines it. Do not edit. +#include +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_kda_train_f227dd21f1f3cb464b29(const __grid_constant__ CUtensorMap do_tma, const __grid_constant__ CUtensorMap vnew_tma, const __grid_constant__ CUtensorMap aqk_tma, float* __restrict__ dAqk, __nv_bfloat16* __restrict__ dv1, int num_heads, float scale); + +namespace cake_binding_kernel_cake_kda_train_f227dd21f1f3cb464b29 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_kda_train_f227dd21f1f3cb464b29"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +// The product of the extents above the trailing `trailing` dims. +int64_t NumelAbove(const TensorView& t, int trailing) { + int64_t product = 1; + for (int axis = 0; axis < t.ndim() - trailing; ++axis) product *= t.size(axis); + return product; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_do_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'do_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "do_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "do_tma"), Extent(d2, "do_tma"), Extent(outer2, "do_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "do_tma"), Stride(((d1 * d2) * 2), "do_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 128u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'do_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'do_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'do_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of do_tma"); + return map; +} + +CUtensorMap EncodeTma_vnew_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'vnew_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "vnew_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "vnew_tma"), Extent(d2, "vnew_tma"), Extent(outer2, "vnew_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "vnew_tma"), Stride(((d1 * d2) * 2), "vnew_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 128u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'vnew_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'vnew_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'vnew_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of vnew_tma"); + return map; +} + +CUtensorMap EncodeTma_aqk_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'aqk_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "aqk_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "aqk_tma"), Extent(d2, "aqk_tma"), Extent(outer2, "aqk_tma")}; + cuuint64_t global_strides[2] = {Stride((d1 * 2), "aqk_tma"), Stride(((d1 * d2) * 2), "aqk_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) + << KERNEL_NAME << ": the box of descriptor 'aqk_tma' exceeds its source extents"; + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'aqk_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'aqk_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of aqk_tma"); + return map; +} + +void Run(TensorView a_do_tma, TensorView a_vnew_tma, TensorView a_aqk_tma, TensorView a_dAqk, TensorView a_dv1, int64_t a_num_heads, double a_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_do_tma.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_do_tma, "do_tma"); + check_dtype(a_do_tma, DLDataType{kDLBfloat, 16, 1}, "do_tma"); + check_cuda_tensor(a_vnew_tma, "vnew_tma"); + check_dtype(a_vnew_tma, DLDataType{kDLBfloat, 16, 1}, "vnew_tma"); + check_cuda_tensor(a_aqk_tma, "aqk_tma"); + check_dtype(a_aqk_tma, DLDataType{kDLBfloat, 16, 1}, "aqk_tma"); + check_cuda_tensor(a_dAqk, "dAqk"); + check_dtype(a_dAqk, DLDataType{kDLFloat, 32, 1}, "dAqk"); + check_cuda_tensor(a_dv1, "dv1"); + check_dtype(a_dv1, DLDataType{kDLBfloat, 16, 1}, "dv1"); + TVM_FFI_CHECK(a_num_heads >= -2147483648LL && a_num_heads <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_heads' = " << a_num_heads + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_do_tma, "do_tma"); + check_contiguous(a_vnew_tma, "vnew_tma"); + check_same_device(a_vnew_tma, a_do_tma, "vnew_tma", "do_tma"); + check_contiguous(a_aqk_tma, "aqk_tma"); + check_same_device(a_aqk_tma, a_do_tma, "aqk_tma", "do_tma"); + check_contiguous(a_dAqk, "dAqk"); + check_same_device(a_dAqk, a_do_tma, "dAqk", "do_tma"); + check_contiguous(a_dv1, "dv1"); + check_same_device(a_dv1, a_do_tma, "dv1", "do_tma"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_train_f227dd21f1f3cb464b29), cudaFuncAttributeMaxDynamicSharedMemorySize, 99328), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_do_tma = EncodeTma_do_tma(a_do_tma); + CUtensorMap v_vnew_tma = EncodeTma_vnew_tma(a_vnew_tma); + CUtensorMap v_aqk_tma = EncodeTma_aqk_tma(a_aqk_tma); + void* v_dAqk = DataPtr(a_dAqk); + void* v_dv1 = DataPtr(a_dv1); + int32_t v_num_heads = static_cast(a_num_heads); + float v_scale = static_cast(a_scale); + void* kernel_args[] = {&v_do_tma, &v_vnew_tma, &v_aqk_tma, &v_dAqk, &v_dv1, &v_num_heads, &v_scale}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(192, 1, 1); + config.dynamicSmemBytes = 99328; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_train_f227dd21f1f3cb464b29), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_kda_train_f227dd21f1f3cb464b29 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_train_f227dd21f1f3cb464b29::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu new file mode 100644 index 00000000000..f527a5d6b18 --- /dev/null +++ b/flashinfer/experimental/cake_kda_train/csrc/cake_kda_train_f227dd21f1f3cb464b29_kernel.cu @@ -0,0 +1,521 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_ACC1_OFFSET 0 +#define TMEM_ACC2_OFFSET 128 +#define NUM_ONE_STAGES 1 +#define SMEM_DO_PANEL_OFF 1024 +#define SMEM_DO_PANEL_STAGE_BYTES 16384 +#define SMEM_DO_PANEL_STRIDE 16384 +#define SMEM_VNEW_PANEL_OFF 33792 +#define SMEM_VNEW_PANEL_STAGE_BYTES 16384 +#define SMEM_VNEW_PANEL_STRIDE 16384 +#define SMEM_AQK_PANEL_OFF 66560 +#define SMEM_AQK_PANEL_STAGE_BYTES 8192 +#define SMEM_AQK_PANEL_STRIDE 8192 +#define SMEM_DO_KMAJ_OFF 1024 +#define SMEM_DO_KMAJ_STAGE_BYTES 32768 +#define SMEM_DO_KMAJ_STRIDE 32768 +#define SMEM_VNEW_KMAJ_OFF 33792 +#define SMEM_VNEW_KMAJ_STAGE_BYTES 32768 +#define SMEM_VNEW_KMAJ_STRIDE 32768 +#define SMEM_DO_MN_OFF 1024 +#define SMEM_DO_MN_STAGE_BYTES 32768 +#define SMEM_DO_MN_STRIDE 32768 +#define SMEM_AQK_MN_OFF 66560 +#define SMEM_AQK_MN_STAGE_BYTES 32768 +#define SMEM_AQK_MN_STRIDE 32768 +#define SMEM_TOTAL 99328 +#define THREADS 192 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +extern "C" { + +__global__ __launch_bounds__(192, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_train_f227dd21f1f3cb464b29(const __grid_constant__ CUtensorMap do_tma, const __grid_constant__ CUtensorMap vnew_tma, const __grid_constant__ CUtensorMap aqk_tma, float* __restrict__ dAqk, __nv_bfloat16* __restrict__ dv1, int num_heads, float scale) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define tiles_full_addr (mbar_base + 0) + #define zero_done_addr (mbar_base + 8) + #define acc_full_addr (mbar_base + 16) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* do_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int do_panel_addr = smem + 1024; + __nv_bfloat16* vnew_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int vnew_panel_addr = smem + 33792; + __nv_bfloat16* aqk_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int aqk_panel_addr = smem + 66560; + __nv_bfloat16* do_kmaj = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int do_kmaj_addr = smem + 1024; + __nv_bfloat16* vnew_kmaj = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int vnew_kmaj_addr = smem + 33792; + __nv_bfloat16* do_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int do_mn_addr = smem + 1024; + __nv_bfloat16* aqk_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int aqk_mn_addr = smem + 66560; + + // Mbarrier init (3 pipeline groups, 0 ordered-sequence groups, 3 barriers) + // Mbarriers at smem_raw[0..24) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'one' --- + // tiles_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // zero_done: 1 barriers, init_count=128 + mbarrier_init(smem + 8, 128); + // acc_full: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 24); + if (warp == 0) { + int _tmem_hold = smem + 24; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_acc1 = taddr; + const int tmem_acc2 = taddr + 128; + + // ---- Role: epilogue ---- + if (warp <= 3) { + { // epilogue_main + int pair = blockIdx.x; + int head = blockIdx.y; + int row0 = pair * 128; + const int epi_warp = warp; + const int epi_tid = epi_warp * 32 + lane; + unsigned int zeros[4]; + zeros[0] = 0; + zeros[1] = 0; + zeros[2] = 0; + zeros[3] = 0; + #pragma unroll + for (int blk = 1; blk < 3; blk++) { + int base = aqk_panel_addr + (unsigned int)(blk * 8192); + #pragma unroll + for (int j = 0; j < 4; j++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(base + (j * 128 + epi_tid) * 16), "r"(zeros[0]), "r"(zeros[1]), "r"(zeros[2]), "r"(zeros[3]) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(zero_done_addr); + unsigned int _phase_acc_full_0 = 0; + mbarrier_wait(acc_full_addr, _phase_acc_full_0); + _phase_acc_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int row_base = epi_warp * 32; + int row = row_base + lane; + int chunk = row / 64; + int t_local = row - chunk * 64; + int grow = row0 + row; + int dA_base = (grow * num_heads + head) * 64; + #pragma unroll + for (int seg = 0; seg < 4; seg++) { + float _tmem_load_0[16]; + tmem_ld_x16(&_tmem_load_0[0], taddr + (unsigned int)(row_base << 16) + (unsigned int)(chunk * 64) + (unsigned int)(seg * 16)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int half = 0; half < 2; half++) { + float out8[8]; + #pragma unroll + for (int j_1 = 0; j_1 < 8; j_1++) { + int s_local = seg * 16 + half * 8 + j_1; + float val = _tmem_load_0[half * 8 + j_1] * scale; + if (s_local > t_local) { + val = 0.0f; + } + out8[j_1] = val; + } + { + unsigned _stv8_0_0 = __float_as_uint(out8[0 + 0]); + unsigned _stv8_0_1 = __float_as_uint(out8[0 + 1]); + unsigned _stv8_0_2 = __float_as_uint(out8[0 + 2]); + unsigned _stv8_0_3 = __float_as_uint(out8[0 + 3]); + unsigned _stv8_0_4 = __float_as_uint(out8[0 + 4]); + unsigned _stv8_0_5 = __float_as_uint(out8[0 + 5]); + unsigned _stv8_0_6 = __float_as_uint(out8[0 + 6]); + unsigned _stv8_0_7 = __float_as_uint(out8[0 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(dAqk + (dA_base + seg * 16 + half * 8) + (0))), "r"(_stv8_0_0), "r"(_stv8_0_1), "r"(_stv8_0_2), "r"(_stv8_0_3), "r"(_stv8_0_4), "r"(_stv8_0_5), "r"(_stv8_0_6), "r"(_stv8_0_7) : "memory"); + } + } + } + long long dv_base = ((long long)grow * (long long)num_heads + (long long)head) * 128; + #pragma unroll + for (int seg_1 = 0; seg_1 < 8; seg_1++) { + float _tmem_load_1[16]; + tmem_ld_x16(&_tmem_load_1[0], taddr + (unsigned int)(row_base << 16) + 128 + (unsigned int)(seg_1 * 16)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + { + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_1[0 + 0], _tmem_load_1[0 + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_1[0 + 2], _tmem_load_1[0 + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_1[0 + 4], _tmem_load_1[0 + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_1[0 + 6], _tmem_load_1[0 + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_1[0 + 8], _tmem_load_1[0 + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_1[0 + 10], _tmem_load_1[0 + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_1[0 + 12], _tmem_load_1[0 + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_1[0 + 14], _tmem_load_1[0 + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(dv1 + (dv_base + (long long)(seg_1 * 16))))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("barrier.sync 1, 128;" ::: "memory"); + if (warp == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(256)); + } + } + } + // ---- Role: load ---- + if (warp == 4) { + { // load_main + int pair_1 = blockIdx.x; + int head_1 = blockIdx.y; + int row0_1 = pair_1 * 128; + if (elect_sync()) { + mbarrier_arrive_expect_tx(tiles_full_addr, 81920); + #pragma unroll + for (int seg_2 = 0; seg_2 < 2; seg_2++) { + tma_3d_gmem2smem(do_panel_addr + (unsigned int)(seg_2 * 16384), (&do_tma), seg_2 * 64, head_1, row0_1, tiles_full_addr); + tma_3d_gmem2smem(vnew_panel_addr + (unsigned int)(seg_2 * 16384), (&vnew_tma), seg_2 * 64, head_1, row0_1, tiles_full_addr); + } + tma_3d_gmem2smem(aqk_panel_addr, (&aqk_tma), 0, head_1, row0_1, tiles_full_addr); + tma_3d_gmem2smem(aqk_panel_addr + 24576, (&aqk_tma), 0, head_1, row0_1 + 64, tiles_full_addr); + } + } + } + // ---- Role: mma ---- + if (warp == 5) { + { // mma_main + unsigned int _phase_tiles_full_0 = 0; + unsigned int _phase_zero_done_0 = 0; + if (elect_sync()) { + mbarrier_wait(tiles_full_addr, _phase_tiles_full_0); + _phase_tiles_full_0 ^= 1; + mbarrier_wait(zero_done_addr, _phase_zero_done_0); + _phase_zero_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = ((do_kmaj_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_1 = ((vnew_kmaj_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_0 = ((do_kmaj_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_1 = ((vnew_kmaj_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 1018;\n\t" + "add.u32 blo, blo, 1018;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc1), "r"(0) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = (((aqk_mn_addr) >> 4) & 0x7FFF) | 0x4000000; + int _smem_desc_lo_3 = (((do_mn_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_2 = (((aqk_mn_addr) >> 4) & 0x3FFF) | 0x4000000; + int _smem_desc_lo_3 = (((do_mn_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136414352;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc2), "r"(0) : "memory"); + tcgen05_commit(acc_full_addr); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/flashinfer/kda_train.py b/flashinfer/kda_train.py index 6f4defbba08..e269cd4425f 100644 --- a/flashinfer/kda_train.py +++ b/flashinfer/kda_train.py @@ -294,7 +294,8 @@ def kcp_partition( The number of ranks of the context-parallel group. layout : str ``"contiguous"`` (equal token shares in order), ``"zigzag"`` or ``"zigzag_per_document"`` (the load-balanced - layouts); ``fragments`` or ``cut_points`` give an explicit layout instead. + layouts); ``fragments`` gives an explicit ragged layout instead, and ``cut_points`` (per document the + ``2 * world_size + 1`` cut offsets) applies to ``layout="zigzag_per_document"`` only. Returns ------- diff --git a/tests/experimental/test_cake_kda_train_kcp.py b/tests/experimental/test_cake_kda_train_kcp.py index b486504222f..3150b880a0c 100644 --- a/tests/experimental/test_cake_kda_train_kcp.py +++ b/tests/experimental/test_cake_kda_train_kcp.py @@ -15,17 +15,20 @@ CPU tests of the KCP metadata layer (``cake_kcp_plan``): partitions against the rows of record, exchange splits and chain tables, head ownership, the plan rule, the byte queries, the private test knobs; the GPU parity test -of the CPU-built chunk / segment tables against the ``table`` / ``seg_table`` programs (``-k gpu``); and the FP64 +of the CPU-built chunk tables against the ``table`` program (``-k gpu``); and the FP64 fake backend against the unpartitioned FP64 reference (``-k fake``). Rows, reference partitions and operands come from the reference harness (``kcp_shapes``, ``kcp_inputs``) and the reference module ``reference_kcp`` when they are on ``PYTHONPATH``; the harness-backed tests are skipped otherwise. """ +import ctypes +import gc import json import os import time import types import warnings +import weakref import pytest import torch @@ -35,6 +38,7 @@ kp = cake_package("cake_kda_train", "cake_kcp_plan") cake_backend = cake_package("cake_kda_train", "cake_backend") cake_jit = cake_package("cake_kda_train", "cake_jit") +cake_kcp_jit = cake_package("cake_kda_train", "cake_kcp_jit") try: # the reference harness (rows of record and the reference partitions) import kcp_shapes @@ -451,8 +455,7 @@ def test_head_ownership_rows(H, W) -> None: ) assert ( kp.kcp_saved_bytes(ctx, H=H, with_conv=False, save_transfer=False) - == 256 * ctx.T_r * H - + kp.tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) + == kp.tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) + ctx.n_in * n_owners * Hg * kp.STATE_TILE_BYTES ) if ( @@ -571,20 +574,31 @@ def test_zero_token_rank_keeps_chains_and_splits() -> None: def test_kcp_plan_examples() -> None: - """Gate 3: the shipped default rule is wave_fill1 (the one-wave fill on the earlier wave rule of - lever L-C), and the §7.2 examples at 152 SMs hold under the earlier occ rule, now selected explicitly.""" - assert kp.P_RULE_DEFAULT == "wave_fill1" and kp.P_RULE_DEFAULT in kp.P_RULES + """Gate 3: the shipped default rule is wave_fill1_p1tie (the one-wave fill rule plus the P = 1 tie line, on top of the + earlier wave rule), and the §7.2 examples at 152 SMs hold + under the earlier occ rule, now selected explicitly.""" + assert kp.P_RULE_DEFAULT == "wave_fill1_p1tie" and kp.P_RULE_DEFAULT in kp.P_RULES dflt = kp.kcp_plan(1, 256, 64, SM_COUNT) - assert dflt == kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="wave_fill1") + assert dflt == kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="wave_fill1_p1tie") + fill1 = kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="wave_fill1") + assert (dflt.P, dflt.NSEG) == ( + fill1.P, + fill1.NSEG, + ) # one fragment: the P = 1 tie line does not apply # 256 chunks is AT the chain cap, not above it: the wave line does not apply; the earlier wave # rule gives the occ P = 3 (192 CTAs on 152 SMs), the shipped one-wave fill the largest one-CTA-per-SM P = 2 wave = kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="wave") assert (wave.p_rule, wave.P, wave.P_occ, wave.P_chain) == ("wave", 3, 3, 1) assert wave.P == kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="occ").P - assert (dflt.p_rule, dflt.P, dflt.P_occ, dflt.P_chain) == ("wave_fill1", 2, 3, 1) + assert (dflt.p_rule, dflt.P, dflt.P_occ, dflt.P_chain) == ( + "wave_fill1_p1tie", + 2, + 3, + 1, + ) # the first NC above the cap: P_chain = 2 and the wave line applies under the shipped rule as under wave above = kp.kcp_plan(1, 257, 64, SM_COUNT) - assert (above.p_rule, above.P_chain) == ("wave_fill1", 2) and above.P >= 2 + assert (above.p_rule, above.P_chain) == ("wave_fill1_p1tie", 2) and above.P >= 2 assert above.P == kp.kcp_plan(1, 257, 64, SM_COUNT, p_rule="wave").P assert kp.kcp_plan(12, 64, 64, SM_COUNT, p_rule="occ").P == 1 # D / CP16 Q3 assert kp.kcp_plan(1, 256, 64, SM_COUNT, p_rule="occ").P == 3 # C / CP16 @@ -619,7 +633,7 @@ def test_kcp_plan_examples() -> None: ) -# Lever L-C. Rule of record: on ranks with n_frag <= WAVE_N_FRAG_MAX = 3 the "wave" +# The plan rule. Rule of record: on ranks with n_frag <= WAVE_N_FRAG_MAX = 3 the "wave" # rule is argmin_P over P in [P_chain, min(64, NC_max)] (smallest P on ties) of # ceil(n_frag P H / SM) x ceil(NC_max / P) x t_walk + (P - 1) x t_compose, t_walk = 2.85 + 3.5 + 1.2 us, t_compose = 20.4 us # (10-ns integers); ranks with more fragments, and ranks whose longest fragment is <= 256 chunks (P_chain = 1), keep the occ P. Information form (b) (superseded): the same @@ -798,7 +812,7 @@ def test_lpt_makespan_matches_the_job_by_job_greedy() -> None: def test_kcp_plan_wave_rule_examples() -> None: - """Lever L-C at 152 SMs, H = 64: the rule of record equals the hand-derived argmin on ranks with <= 3 fragments, + """The plan rule at 152 SMs, H = 64: the rule of record equals the hand-derived argmin on ranks with <= 3 fragments, the occ P on guarded ranks; the information form (b) equals its hand values and the oracle; occ is untouched.""" walk = sum(kp.WAVE_RULE_10NS[k] for k in ("tape", "xfer", "offset")) assert (walk, kp.WAVE_RULE_10NS["compose"], kp.WAVE_N_FRAG_MAX) == ( @@ -929,7 +943,7 @@ def test_kcp_plan_wave_rule_examples() -> None: assert (dflt_488.P, dflt_488.NSEG, dflt_488.p_rule) == ( wave_488.P, wave_488.NSEG, - "wave_fill1", + "wave_fill1_p1tie", ) # above the chain cap the shipped default (wave_fill1) is the wave rule assert kp.kcp_plan(1, 488, 64, SM_COUNT, p_rule="occ").P == 3 with pytest.raises(ValueError, match="p_rule"): @@ -1035,10 +1049,14 @@ def test_kcp_plan_wave_rule_on_q1_c(monkeypatch, W) -> None: sm_count=SM_COUNT, _kcp_test_knobs={"p_rule": "wave"}, ) - assert (occ.p_rule, wave.p_rule, dflt.p_rule) == ("occ", "wave", "wave_fill1") + assert (occ.p_rule, wave.p_rule, dflt.p_rule) == ( + "occ", + "wave", + "wave_fill1_p1tie", + ) assert (dflt.heads(64).P, dflt.describe(64)["p_rule"]) == ( expect_dflt[r], - "wave_fill1", + "wave_fill1_p1tie", ) assert ( dflt.heads(64).P @@ -1215,7 +1233,13 @@ def test_kcp_plan_wave_fill_examples() -> None: H = 64 ranks move 2 -> 1 under wave_fill only, every rank above the chain cap and every H <= 4 rank keeps its program, the public default is untouched.""" SM = SM_COUNT - assert kp.P_RULES == ("occ", "wave", "wave_fill1", "wave_fill") + assert kp.P_RULES == ( + "occ", + "wave", + "wave_fill1", + "wave_fill", + "wave_fill1_p1tie", + ) assert kp.WAVE_FILL_N_FRAG_MAX == {"wave_fill1": 1, "wave_fill": 3} for nc in (222, 244, 256): assert kp.kcp_plan(1, nc, 64, SM, p_rule="occ").P == 3 @@ -1311,7 +1335,9 @@ def test_kcp_plan_wave_fill_examples() -> None: 2, "wave_fill1", ) - assert kp.kcp_plan(1, 244, 64, SM).P == 2 and kp.P_RULE_DEFAULT == "wave_fill1" + assert ( + kp.kcp_plan(1, 244, 64, SM).P == 2 and kp.P_RULE_DEFAULT == "wave_fill1_p1tie" + ) with pytest.raises(ValueError, match="p_rule"): kp.kcp_plan(1, 244, 64, SM, p_rule="fill") _evidence("plan_wave_fill_examples", {"pass": True}) @@ -1321,7 +1347,7 @@ def test_wave_fill_knobs_are_gated(monkeypatch) -> None: monkeypatch.delenv(kp.TEST_KNOBS_ENV, raising=False) with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): kp.validate_test_knobs({"p_rule": "wave_fill1"}) - assert kp.validate_test_knobs(None)["p_rule"] == "wave_fill1" + assert kp.validate_test_knobs(None)["p_rule"] == "wave_fill1_p1tie" monkeypatch.setenv(kp.TEST_KNOBS_ENV, "1") for rule in ("wave_fill1", "wave_fill"): assert kp.validate_test_knobs({"p_rule": rule}) == { @@ -1440,6 +1466,535 @@ def test_kcp_plan_wave_fill_lane_rows(monkeypatch, case) -> None: ) +# --------------------------------------------------------------------------- the P = 1 tie rule and the fused carry merges + + +def _p1tie_ref(n_frag: int, nc_frags, H: int, sm: int) -> int: + """Oracle of the "wave_fill1_p1tie" rule: the "wave_fill1" P, then 1 on a rank with 2..3 fragments below + the chain cap whose base P > 1 runs more CTAs than one wave of P = 1 (n_frag H <= SM) when the LPT makespan of one + equal-segment walk at P = 1 over the fragment chunk counts is <= the base P's (the wave model's waves x steps on equal + fragments).""" + nc_frags = [int(x) for x in nc_frags] + nc_max = max(nc_frags, default=0) + base = kp.kcp_plan(n_frag, nc_max, H, sm, p_rule="wave_fill1").P + if ( + 2 <= n_frag <= kp.WAVE_N_FRAG_MAX + and 0 < nc_max < kp.WAVE_MIN_NC_MAX_FRAG + and base > 1 + and n_frag * H <= sm + ): + + def steps(P: int) -> int: + return kp.lpt_makespan( + [(-(-nc // P), P * H) for nc in nc_frags if nc > 0], sm + ) + + if steps(1) <= steps(base): + return 1 + return base + + +def test_kcp_plan_p1tie_examples() -> None: + """The P = 1 tie rule at 152 SMs: equal two-fragment H = 64 ranks below the chain cap (the CP16 / CP8 + Z ranks of record) move P 2 -> 1; unequal two-fragment ranks (CP16 C ranks cut by a document boundary), three-fragment + ranks, one-fragment ranks, ranks above the cap and every H <= 4 rank keep the shipped program; nc_frags is validated; + the public default is untouched.""" + SM = SM_COUNT + assert kp.P1_TIE_BASE == {"wave_fill1_p1tie": "wave_fill1"} + assert "wave_fill1_p1tie" in kp.P_RULES and kp.P_RULE_DEFAULT == "wave_fill1_p1tie" + for nc in (56, 111, 118, 122, 125, 128, 200, 256): + assert kp.kcp_plan(2, nc, 64, SM, p_rule="wave_fill1").P == 2 + assert kp.kcp_plan(2, nc, 64, SM, p_rule="wave_fill1_p1tie").P == 1 + assert ( + kp.kcp_plan(2, nc, 64, SM, p_rule="wave_fill1_p1tie", nc_frags=[nc, nc]).P + == 1 + ) + assert kp.p1_tie_steps([nc, nc], 64, SM, 1) == nc + assert kp.p1_tie_steps([nc, nc], 64, SM, 2) == 2 * -(-nc // 2) >= nc + assert _p1tie_ref(2, [nc, nc], 64, SM) == 1 + # equal fragments reproduce the wave model's waves x steps exactly + for n_frag, nc, P in ( + (2, 111, 2), + (2, 111, 1), + (3, 122, 1), + (1, 244, 2), + (1, 244, 3), + (2, 170, 2), + (2, 60, 19), + ): + waves = -(-(n_frag * P * 64) // SM) + assert kp.p1_tie_steps([nc] * n_frag, 64, SM, P) == waves * -(-nc // P) + # unequal two-fragment ranks: P = 1 walks the long fragment alone -> the base P stays + for frags in ([170, 59], [200, 30], [60, 62], [130, 100], [111, 50]): + nc = max(frags) + assert ( + kp.kcp_plan(2, nc, 64, SM, p_rule="wave_fill1_p1tie", nc_frags=frags).P == 2 + ), frags + assert kp.p1_tie_steps(frags, 64, SM, 1) > kp.p1_tie_steps(frags, 64, SM, 2) + assert _p1tie_ref(2, frags, 64, SM) == 2 + assert ( + kp.p1_tie_steps([60, 62], 64, SM, 2) == 61 + ) # LPT: 31 + 30 on 80 of the 152 SMs + assert kp.p1_tie_steps([170, 59], 64, SM, 2) == 115 # 85 + 30 + assert ( + kp.kcp_plan(2, 100, 64, SM, p_rule="wave_fill1_p1tie", nc_frags=[100, 100]).P + == 1 + ) + # an empty second fragment: one live chain group of 64 CTAs at P = 1 (111 steps) vs 128 CTAs x 56 at P = 2 -> stays 2 + assert ( + kp.kcp_plan(2, 111, 64, SM, p_rule="wave_fill1_p1tie", nc_frags=[111, 0]).P == 2 + ) + # three-fragment H = 64 ranks (P_occ = 1 already), ranks above the cap (the wave argmin), one-fragment ranks (the + # fill), many-fragment ranks: unchanged + for n_frag, nc in ( + (3, 122), + (3, 80), + (2, 257), + (2, 488), + (3, 300), + (1, 244), + (1, 100), + (1, 488), + (4, 100), + (8, 251), + ): + assert ( + kp.kcp_plan(n_frag, nc, 64, SM, p_rule="wave_fill1_p1tie").P + == kp.kcp_plan(n_frag, nc, 64, SM, p_rule="wave_fill1").P + ), (n_frag, nc) + # H <= 4 ranks keep their program: the base P fills the SMs with chains of ceil(NC / P) < NC steps + for n_frag, nc, H in ( + (2, 14, 4), + (2, 100, 4), + (2, 256, 4), + (3, 20, 4), + (3, 4, 4), + (3, 128, 4), + (2, 2, 4), + (2, 1, 4), + (2, 79, 1), + (3, 16, 1), + (2, 1, 1), + (3, 256, 1), + ): + assert ( + kp.kcp_plan(n_frag, nc, H, SM, p_rule="wave_fill1_p1tie").P + == kp.kcp_plan(n_frag, nc, H, SM, p_rule="wave_fill1").P + ), (n_frag, nc, H) + # 148 SMs (GB200): the same move; SM = 1: no occupancy pressure, base P = 1 already + assert kp.kcp_plan(2, 122, 64, 148, p_rule="wave_fill1_p1tie").P == 1 + assert ( + kp.kcp_plan(2, 122, 64, 1, p_rule="wave_fill1_p1tie").P + == kp.kcp_plan(2, 122, 64, 1, p_rule="wave_fill1").P + == 1 + ) + # the plan terms are reported alike; the rule name travels + a = kp.kcp_plan(2, 111, 64, SM, p_rule="wave_fill1_p1tie") + b = kp.kcp_plan(2, 111, 64, SM, p_rule="wave_fill1") + assert (a.P, a.NSEG, a.P_occ, a.P_chain, a.P_bal, a.n_frag, a.NC_max_frag) == ( + 1, + 2, + b.P_occ, + b.P_chain, + b.P_bal, + 2, + 111, + ) + assert (a.p_rule, b.P, b.NSEG, b.p_rule) == ("wave_fill1_p1tie", 2, 4, "wave_fill1") + # nc_frags is validated for every rule; the other rules ignore the counts + for rule in kp.P_RULES: + with_counts = kp.kcp_plan(2, 111, 64, SM, p_rule=rule, nc_frags=[111, 111]).P + plain = kp.kcp_plan(2, 111, 64, SM, p_rule=rule).P + assert with_counts == plain, rule + assert ( + plain == kp.kcp_plan(2, 111, 64, SM, p_rule=rule, nc_frags=[111, 50]).P + or rule == "wave_fill1_p1tie" + ) + for bad in ([111], [111, 112], [111, -1], [111, 50, 10]): + with pytest.raises(ValueError, match="nc_frags"): + kp.kcp_plan(2, 111, 64, SM, p_rule=rule, nc_frags=bad) + assert kp.kcp_plan(0, 0, 64, SM, p_rule="wave_fill1_p1tie").P == 1 + assert kp.kcp_plan(2, 0, 64, SM, p_rule="wave_fill1_p1tie").P == 1 + assert kp.kcp_plan(2, 111, 64, SM).P == 1 # the shipped default + assert kp.kcp_plan(2, 111, 64, SM, nc_frags=[111, 111]).P == 1 + assert kp.kcp_plan(2, 111, 64, SM, p_rule="wave_fill1").P == 2 + assert kp.kcp_plan(2, 111, 64, SM, p_rule="wave_fill1", nc_frags=[111, 111]).P == 2 + with pytest.raises(ValueError, match="p_rule"): + kp.kcp_plan(2, 111, 64, SM, p_rule="p1tie") + _evidence("plan_p1tie_examples", {"pass": True}) + + +def _two_doc_two_rank_partition(chunks: int = 100): + """Two documents of two equal fragments across two ranks: every rank holds a document-start (crossing-out) fragment + and a document-last (crossing-in) fragment of ``chunks`` chunks -- the two-fragment Z rank shape of the tie rule.""" + n = chunks * 64 + cu = torch.tensor([0, 2 * n, 4 * n], dtype=torch.int64) + return kp.kcp_partition( + cu, 2, fragments=[[(0, 0, n), (1, n, n)], [(0, n, n), (1, 0, n)]] + ) + + +def test_p1tie_knob_is_gated_and_plans(monkeypatch) -> None: + monkeypatch.delenv(kp.TEST_KNOBS_ENV, raising=False) + with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): + kp.validate_test_knobs({"p_rule": "wave_fill1_p1tie"}) + part = _two_doc_two_rank_partition(100) + ctx0 = kp.build_kcp_context(part, None, 0, "cpu", 4, sm_count=SM_COUNT) + assert (ctx0.p_rule, ctx0.heads(64).P, ctx0.heads(64).NSEG) == ( + "wave_fill1_p1tie", + 1, + 2, + ) + monkeypatch.setenv(kp.TEST_KNOBS_ENV, "1") + ctx_fill1 = kp.build_kcp_context( + part, + None, + 0, + "cpu", + 4, + sm_count=SM_COUNT, + _kcp_test_knobs={"p_rule": "wave_fill1"}, + ) + assert (ctx_fill1.p_rule, ctx_fill1.heads(64).P, ctx_fill1.heads(64).NSEG) == ( + "wave_fill1", + 2, + 4, + ) + assert kp.validate_test_knobs({"p_rule": "wave_fill1_p1tie"}) == { + "precision_mode": 0, + "p_rule": "wave_fill1_p1tie", + } + for r in range(2): + ctx = kp.build_kcp_context( + part, + None, + r, + "cpu", + 4, + sm_count=SM_COUNT, + _kcp_test_knobs={"p_rule": "wave_fill1_p1tie"}, + ) + assert (ctx.p_rule, ctx.heads(64).P, ctx.heads(64).NSEG) == ( + "wave_fill1_p1tie", + 1, + 2, + ) + assert ctx.describe(64)["p_rule"] == "wave_fill1_p1tie" + assert ctx.heads(4).P == ctx0.heads(4).P # H = 4: unchanged + # an unequal pair (170 + 59 chunks on one rank) keeps P = 2 under the knob + n1, n2 = 170 * 64, 59 * 64 + part_u = kp.kcp_partition( + torch.tensor([0, 2 * n1, 2 * n1 + 2 * n2], dtype=torch.int64), + 2, + fragments=[[(0, 0, n1), (1, n2, n2)], [(0, n1, n1), (1, 0, n2)]], + ) + ctx_u = kp.build_kcp_context( + part_u, + None, + 0, + "cpu", + 4, + sm_count=SM_COUNT, + _kcp_test_knobs={"p_rule": "wave_fill1_p1tie"}, + ) + assert ctx_u.heads(64).P == 2 + + +_INV_WALKS = ("offset", "all", "live", "in", "m_b1", "start", "end") + + +def _inventory(ctx, H: int) -> dict: + """Inventory of one rank's tables under ``ctx.p_rule`` (CPU): P / NSEG, the K4 chain families (chains, + elements), the walk lists (segments, CTAs = segments x H, LPT step-times at one CTA per SM, longest segment) and + whether the fused carry launches save a launch (both families non-empty).""" + tb = cake_kcp._Tables(ctx, H) + seg_table = [int(x) for x in tb.hp.seg_chunk_start] + + def walk(lst) -> dict: + steps = [seg_table[s + 1] - seg_table[s] for s in lst] + groups: dict[int, int] = {} + for st in steps: + if st > 0: + groups[st] = groups.get(st, 0) + H + return { + "segs": len(lst), + "ctas": len(lst) * H, + "lpt_steps": kp.lpt_makespan(list(groups.items()), ctx.sm_count) + if groups + else 0, + "max_steps": max(steps, default=0), + } + + chains = {k: int(c.n_chains) for k, c in tb.chains.items()} + return { + "P": int(tb.P), + "NSEG": int(tb.NSEG), + "n_in": int(tb.n_in), + "n_out": int(tb.n_out), + "n_frag": int(tb.n_frag), + "chains": chains, + "elems": {k: int(c.n_elems) for k, c in tb.chains.items()}, + "walks": { + "offset": walk(tb.walk_offset), + "all": walk(tb.walk_all), + "live": walk(tb.walk_live), + "in": walk(tb.walk_in), + "m_b1": walk(tb.walk_m_b1), + "start": walk(tb.walk_start), + "end": walk(tb.walk_end), + }, + "transpose_tiles": int(sum(len(row) for row in tb.mf_tiles)), + "fuse_saves_fwd": int(chains["carry_start"] > 0 and chains["carry_in"] > 0), + "fuse_saves_bwd": int(chains["dcarry_end"] > 0 and chains["dcarry_cross"] > 0), + } + + +# the rows of :func:`_kb_rows` on which the P = 1 tie rule changes at least one rank's P at 152 SMs (the enumeration +# of record of the rule's adoption): the two-fragment Z ranks of the perf packs at CP 8 / 16 and four coverage rows, all +# H = 64 +P1TIE_CHANGED_ROWS = frozenset( + [(f"Q{i}", W, "Z", 64) for i in range(1, 7) for W in (8, 16)] + + [("C2", 4, "D", 64), ("C2", 4, "Z", 64), ("C6", 8, "Z", 64), ("C8", 8, "Z", 64)] +) + + +@HARNESS +@pytest.mark.parametrize( + "case", KB_ROWS, ids=[f"{n}-cp{W}-{lay}-H{H}" for n, W, lay, H in KB_ROWS] +) +def test_kcp_plan_p1tie_lane_rows(monkeypatch, case) -> None: + """Plan enumeration of the P = 1 tie rule: per rank of every row of :func:`_kb_rows` the P under the previous rule + "wave_fill1" (named explicitly: the default is the tie rule itself) and under the shipped default "wave_fill1_p1tie" + (with the actual fragment chunk counts), the P = 1 / base-P walk makespans of the tie test, and -- for the H = 64 rows + at CP 8 / 16 -- the K4 / walk inventory of both programs (the pre-estimate reads it). The module equals the oracle, the + rule changes only two- / three-fragment ranks below the cap whose base P exceeds 1, never an H <= 4 rank, the rows it + changes are exactly :data:`P1TIE_CHANGED_ROWS`, and a context built without a knob plans with it. Evidence record + `plan_p1tie_rows`.""" + name, W, lay, H = case + row = kcp_shapes.ROWS[name] + _frags, part = harness_partition(row, W, lay) + monkeypatch.setenv(kp.TEST_KNOBS_ENV, "1") + assert ( + kp.P_RULE_DEFAULT == "wave_fill1_p1tie" + and kp.P1_TIE_BASE[kp.P_RULE_DEFAULT] == "wave_fill1" + ) + want_inv = H == 64 and W >= 8 + ranks = [] + for r in range(W): + mine = part.fragments[r] + actual = [-(-f.length // 64) for f in mine] + n_frag, nc = len(mine), max(actual, default=0) + base = kp.kcp_plan( + n_frag, nc, H, SM_COUNT, p_rule="wave_fill1", nc_frags=actual + ).P + cand = kp.kcp_plan(n_frag, nc, H, SM_COUNT, nc_frags=actual).P + assert ( + cand + == kp.kcp_plan( + n_frag, nc, H, SM_COUNT, p_rule="wave_fill1_p1tie", nc_frags=actual + ).P + ), (case, r) + rec = { + "rank": r, + "nc_frags": actual, + "n_frag": n_frag, + "nc_max": nc, + "frags": [ + (f.doc, f.off, f.length, bool(f.crossing_in), bool(f.crossing_out)) + for f in mine + ], + "wave_fill1": base, + "wave_fill1_p1tie": cand, + "steps_p1": kp.p1_tie_steps(actual, H, SM_COUNT, 1) if nc else 0, + "steps_base": kp.p1_tie_steps(actual, H, SM_COUNT, base) if nc else 0, + "sm148": kp.kcp_plan( + n_frag, nc, H, 148, p_rule="wave_fill1_p1tie", nc_frags=actual + ).P, + } + assert cand == _p1tie_ref(n_frag, actual, H, SM_COUNT), (case, rec) + if cand != base: + assert cand == 1 and base > 1 and 2 <= n_frag <= 3 and 0 < nc < 257, ( + case, + rec, + ) + assert rec["steps_p1"] <= rec["steps_base"], (case, rec) + if H <= 4: + assert cand == base, (case, rec) + if want_inv: + ctx_b = kp.build_kcp_context( + part, + None, + r, + "cpu", + 4, + sm_count=SM_COUNT, + _kcp_test_knobs={"p_rule": "wave_fill1"}, + ) + ctx_c = kp.build_kcp_context(part, None, r, "cpu", 4, sm_count=SM_COUNT) + inv_b, inv_c = _inventory(ctx_b, H), _inventory(ctx_c, H) + assert (inv_b["P"], inv_c["P"]) == (base, cand), (case, rec) + if cand == base: + assert inv_b == inv_c + rec["inv_wave_fill1"], rec["inv_p1tie"] = inv_b, inv_c + ranks.append(rec) + changed = [x["rank"] for x in ranks if x["wave_fill1_p1tie"] != x["wave_fill1"]] + assert bool(changed) == ((name, W, lay, H) in P1TIE_CHANGED_ROWS), (case, changed) + _evidence( + "plan_p1tie_rows", + { + "row": name, + "cp": W, + "layout": lay, + "H": H, + "sm": SM_COUNT, + "wave_fill1": [x["wave_fill1"] for x in ranks], + "wave_fill1_p1tie": [x["wave_fill1_p1tie"] for x in ranks], + "changed": changed, + "inventory": want_inv, + "ranks": ranks, + }, + ) + for r in changed[:1] or [0]: + ctx = kp.build_kcp_context(part, None, r, "cpu", 4, sm_count=SM_COUNT) + assert (ctx.p_rule, ctx.heads(H).P, ctx.describe(H)["p_rule"]) == ( + "wave_fill1_p1tie", + ranks[r]["wave_fill1_p1tie"], + "wave_fill1_p1tie", + ) + + +def test_merge_fuse_knob_and_fused_carry_chains(monkeypatch) -> None: + """``merge_fuse`` is a gated test knob ("none" shipped | "carry"); the fused chain families + carry_both / dcarry_both are the concatenation of carry_start + carry_in / dcarry_end + dcarry_cross (same pools; every + chain keeps its elements, initial state and final slot), non-empty at P > 1 on a rank holding both fragment kinds, empty + at P = 1; the default options and the other families are untouched.""" + monkeypatch.delenv(cake_kcp.TEST_KNOBS_ENV, raising=False) + assert cake_kcp.MERGE_FUSES == ("none", "carry") + assert cake_kcp.MERGE_FUSE_DEFAULT == "none" + assert cake_kcp.P2_KNOBS["merge_fuse"] == cake_kcp.MERGE_FUSES + assert cake_kcp._options({}).merge_fuse == "none" + assert "merge_fuse" not in cake_kcp._OPTION_KEYS + with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): + cake_kcp._options({"_kcp_test_knobs": {"merge_fuse": "carry"}}) + with pytest.raises(TypeError): + cake_kcp._options({"merge_fuse": "carry"}) + monkeypatch.setenv(cake_kcp.TEST_KNOBS_ENV, "1") + o = cake_kcp._options({"_kcp_test_knobs": {"merge_fuse": "carry"}}) + assert (o.merge_fuse, o.schedule, o.fused_walks, o.cols_rule) == ( + "carry", + "wait_first", + 1, + "fixed", + ) + o = cake_kcp._options( + { + "_kcp_test_knobs": { + "merge_fuse": "carry", + "cols_rule": "by_chains", + "p_rule": "wave_fill1_p1tie", + } + } + ) + assert (o.merge_fuse, o.cols_rule, o.p_rule) == ( + "carry", + "by_chains", + "wave_fill1_p1tie", + ) + assert ( + cake_kcp._options({"_kcp_test_knobs": {"merge_fuse": "none"}}).merge_fuse + == "none" + ) + with pytest.raises(ValueError, match="merge_fuse"): + cake_kcp._options({"_kcp_test_knobs": {"merge_fuse": "all"}}) + spec = cake_kcp._Tables.FAMILIES + assert ( + spec["carry_both"] == {**spec["carry_start"], "x0": "S_in"} == spec["carry_in"] + ) + assert ( + spec["dcarry_both"] + == {**spec["dcarry_end"], "x0": "b2_recv"} + == spec["dcarry_cross"] + ) + part = _two_doc_two_rank_partition(100) + for r in range(2): + ctx = kp.build_kcp_context( + part, + None, + r, + "cpu", + 4, + sm_count=SM_COUNT, + _kcp_test_knobs={"p_rule": "wave_fill1"}, + ) # P = 2 on the tie shape: the shipped tie rule plans P = 1 here + tb = cake_kcp._Tables(ctx, 64) + assert tb.P == 2 and (tb.n_in, tb.n_out) == (1, 1) + for both, a, b in ( + ("carry_both", "carry_start", "carry_in"), + ("dcarry_both", "dcarry_end", "dcarry_cross"), + ): + ca, cb, cf = tb.chains[a], tb.chains[b], tb.chains[both] + assert (ca.n_chains, cb.n_chains, cf.n_chains) == (64, 64, 128) + assert cf.ncols == ca.ncols == cb.ncols == cake_kcp.D + assert cf.elem_m == ca.elem_m + cb.elem_m + assert cf.elem_c == ca.elem_c + cb.elem_c + assert cf.elem_y == ca.elem_y + cb.elem_y + assert cf.x0_index == ca.x0_index + cb.x0_index + assert cf.y_final == ca.y_final + cb.y_final + assert cf.chain_start == ca.chain_start + [ + ca.n_elems + s for s in cb.chain_start[1:] + ] + assert all(x == -1 for x in ca.x0_index) and all( + x >= 0 for x in cb.x0_index + ) + cf.check_alignment() + tb._check_chain(both, cf) + d = tb.describe()["chains"] + assert d["carry_both"] == { + "n_chains": 128, + "n_elems": d["carry_start"]["n_elems"] + d["carry_in"]["n_elems"], + } + assert d["dcarry_both"] == { + "n_chains": 128, + "n_elems": d["dcarry_end"]["n_elems"] + d["dcarry_cross"]["n_elems"], + } + inv = _inventory(ctx, 64) + assert (inv["fuse_saves_fwd"], inv["fuse_saves_bwd"]) == (1, 1) + e = cake_kcp.Chains.concat(cake_kcp.Chains(), cake_kcp.Chains()) + assert (e.n_chains, e.n_elems, e.chain_start) == (0, 0, [0]) + with pytest.raises(ValueError, match="ncols"): + cake_kcp.Chains.concat(cake_kcp.Chains(128), cake_kcp.Chains(256)) + one = cake_kcp.Chains() + one.add([(0, -1, -1), (16384, 16384, -1)], x0=-1, y_final=32768) + two = cake_kcp.Chains.concat(one, one) + assert ( + two.chain_start == [0, 2, 4] + and two.x0_index == [-1, -1] + and two.y_final == [32768, 32768] + ) + # P = 1 (an eight-fragment D rank): the fused families are empty like their parts + tb1 = cake_kcp._Tables( + kp.build_kcp_context(_small_partition(4), None, 1, "cpu", 4, sm_count=SM_COUNT), + 64, + ) + assert tb1.P == 1 + assert all( + tb1.chains[k].n_chains == 0 + for k in ( + "carry_start", + "carry_in", + "carry_both", + "dcarry_end", + "dcarry_cross", + "dcarry_both", + ) + ) + inv1 = _inventory( + kp.build_kcp_context(_small_partition(4), None, 1, "cpu", 4, sm_count=SM_COUNT), + 64, + ) + assert (inv1["fuse_saves_fwd"], inv1["fuse_saves_bwd"]) == (0, 0) + + def test_limits_are_rejected_before_launch() -> None: """Gate 3: n_frag > 65536 and int32 overflows are refused by the metadata layer.""" W = 2 @@ -1592,8 +2147,11 @@ def test_partition_validation_and_defaults() -> None: def test_tables_follow_the_kernel_contract() -> None: - """chunk_bos / chunk_len / seq_chunk_start / seq_order / seg_chunk_start of the spec formulas; the table layout - equals the existing backend's with the segment table present.""" + """chunk_bos / chunk_len / seq_chunk_start / seq_order / seg_chunk_start of the spec formulas; the four chunk tables + are laid out as the existing backend lays them out, and the boundary kernels' ``[NSEG + 1]`` segment table follows + them in the same 256 B-aligned form. The existing backend's own fifth table is a different object since its segmented + recurrences took a block table (``seg_table``, ``NSEG + 2 * N + 1`` words): the two layouts agree on the first four + entries, and the KCP's fifth entry keeps the width the K1-K4 ``seq_chunk_start`` argument has.""" cu = (0, 1, 1, 66, 130, 130, 193) # 1 token, empty, 65, 64, empty, 63 t = kp.chunk_tables(cu) assert t["seq_chunk_start"] == [0, 1, 1, 3, 4, 4, 5] @@ -1620,7 +2178,16 @@ def test_tables_follow_the_kernel_contract() -> None: seg[9:12] == [3, 3, 3] and seg[12:15] == [4, 4, 4] and seg[15:18] == [4, 4, 4] ) layout = kp.table_layout(nc_pad, 6, 18) - assert layout == cake_backend._table_layout(nc_pad, 6, 18) + assert layout[:4] == cake_backend._table_layout(nc_pad, 6, 0) + assert layout[4] == ( + "seg_chunk_start", + 256, + 19, + ) # NSEG + 1 words in the slot after seq_order + assert cake_backend.TABLE_NAMES[:4] == kp.TABLE_NAMES[:4] + assert ( + cake_backend._table_layout(nc_pad, 6, 18)[4][0] != kp.TABLE_NAMES[4] + ) # the backend's block table: not this table assert ( kp.tables_bytes(nc_pad, 6, 18) == 4 * 64 * 5 ) # five 256 B-aligned tables of < 64 words each @@ -1644,11 +2211,14 @@ def test_byte_queries_and_knobs(monkeypatch) -> None: ) # literal form: H = 64, light, no transfer, conv, C = 8192 assert ( base - == 256 * ctx.T_r * 64 - + kp.tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) + == kp.tables_bytes(ctx.NC_pad, ctx.n_frag, hp.NSEG) + ctx.n_in * s + ctx.halo.n_hist_rows * 8192 * 2 - ) + ) # no token-row Aqk / Akk term: the saved set carries none + assert cake_kcp.KCP_SAVED_VERSION == 2 + fields = cake_kcp.KcpSaved.__dataclass_fields__ + assert "Aqk" not in fields and "Akk" not in fields + assert cake_kcp.KcpSaved._OPTIONAL == ("c1_recv", "M_seg") assert ( kp.kcp_saved_bytes(ctx, save_transfer=False, with_conv=False) == base - ctx.halo.n_hist_rows * 8192 * 2 @@ -1727,7 +2297,10 @@ def test_byte_queries_and_knobs(monkeypatch) -> None: assert ws_t <= ws # the kept c1 tensor leaves the arena (SPEC-4) # knobs monkeypatch.delenv(kp.TEST_KNOBS_ENV, raising=False) - assert kp.validate_test_knobs(None) == {"precision_mode": 0, "p_rule": "wave_fill1"} + assert kp.validate_test_knobs(None) == { + "precision_mode": 0, + "p_rule": "wave_fill1_p1tie", + } assert kp.validate_test_knobs(None)["p_rule"] == kp.P_RULE_DEFAULT with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): kp.validate_test_knobs({"precision_mode": 2}) @@ -1738,9 +2311,12 @@ def test_byte_queries_and_knobs(monkeypatch) -> None: monkeypatch.setenv(kp.TEST_KNOBS_ENV, "1") assert kp.validate_test_knobs({"precision_mode": 1}) == { "precision_mode": 1, - "p_rule": "wave_fill1", + "p_rule": "wave_fill1_p1tie", + } + assert kp.validate_test_knobs({}) == { + "precision_mode": 0, + "p_rule": "wave_fill1_p1tie", } - assert kp.validate_test_knobs({}) == {"precision_mode": 0, "p_rule": "wave_fill1"} assert kp.validate_test_knobs({"p_rule": "occ"}) == { "precision_mode": 0, "p_rule": "occ", @@ -1767,14 +2343,14 @@ def test_byte_queries_and_knobs(monkeypatch) -> None: assert set(ctx.describe(64)) >= {"rank", "T_r", "halo", "heads"} -# --------------------------------------------------------------------------- the schedule / fused-walk / column-rule levers +# --------------------------------------------------------------------------- the schedule / fused-walk / column-rule knobs cake_kcp = cake_package("cake_kda_train", "cake_kcp") LEVER_TENSORS = ("out", "dq", "dk", "dv", "dg", "dbeta") def test_round1_knobs_are_gated_and_validated(monkeypatch) -> None: - """``schedule`` (L-A), ``fused_walks`` (L-B) and ``cols_rule`` (D1) are test knobs: refused + """``schedule``, ``fused_walks`` and ``cols_rule`` are test knobs: refused without ``CAKE_KCP_TEST_KNOBS=1``, never public keywords (``graph`` is the exception, public since this revision), validated against their admitted values; the shipped defaults since an earlier revision are the pair wait_first + fused_walks=1 (one tape launch per @@ -1855,6 +2431,42 @@ def test_ring_exchange_contract() -> None: ) +def test_host_trim_knob_is_gated_and_validated(monkeypatch) -> None: + """``host_trim`` is a private test knob: refused without + ``CAKE_KCP_TEST_KNOBS=1``, never a public keyword, validated against :data:`HOST_TRIMS`; the shipped default is 1 + (the trimmed host path of the eager call), 0 selects the previous host path (the same launches + in the same order) for the A/B record.""" + monkeypatch.delenv(cake_kcp.TEST_KNOBS_ENV, raising=False) + assert cake_kcp.HOST_TRIM_DEFAULT == 1 and cake_kcp.HOST_TRIMS == (0, 1) + assert cake_kcp._options({}).host_trim == 1 + with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): + cake_kcp._options({"_kcp_test_knobs": {"host_trim": 0}}) + with pytest.raises(TypeError): + cake_kcp._options({"host_trim": 1}) + monkeypatch.setenv(cake_kcp.TEST_KNOBS_ENV, "1") + assert cake_kcp._options({"_kcp_test_knobs": {"host_trim": 1}}).host_trim == 1 + assert cake_kcp._options({"_kcp_test_knobs": {"host_trim": 0}}).host_trim == 0 + o = cake_kcp._options( + { + "_kcp_test_knobs": { + "host_trim": 1, + "fused_walks": 1, + "walk_geometry": "lpt_tape_one_cta_per_sm", + } + } + ) + assert (o.host_trim, o.fused_walks, o.walk_geometry) == ( + 1, + 1, + "lpt_tape_one_cta_per_sm", + ) + for bad in ({"host_trim": 2}, {"host_trim": -1}): + with pytest.raises(ValueError): + cake_kcp._options({"_kcp_test_knobs": bad}) + with pytest.raises(TypeError): + cake_kcp._options({"host_trim": 0}) + + def test_ring_exchange_resolves_save_transfer() -> None: """``exchange="ring"`` has no ``save_transfer=True`` form (the ring path receives no c1 tensor), so an unset ``save_transfer`` resolves to ``False`` for the ring and to ``SAVE_TRANSFER_DEFAULT`` (``True``) @@ -1938,9 +2550,18 @@ def test_graph_knob_is_gated_and_validated(monkeypatch) -> None: with pytest.raises(ValueError): cake_kcp._options({"_kcp_test_knobs": {"graph": "capture"}}) # graph_cap (the cache-config form of a captured call) is a gated knob, never a public keyword - assert cake_kcp.GRAPH_CAPS == ("none", "capture_pref") + # (plus the kernel-node forms "node_attr" / "node_attr_post"; "node_attr" is + # GRAPH_CAP_DEFAULT) + assert cake_kcp.GRAPH_CAPS == ( + "none", + "capture_pref", + "node_attr", + "node_attr_post", + ) + assert cake_kcp.GRAPH_NODE_CAPS == ("node_attr", "node_attr_post") assert cake_kcp.P2_KNOBS["graph_cap"] == cake_kcp.GRAPH_CAPS assert cake_kcp.GRAPH_CAP_DEFAULT in cake_kcp.GRAPH_CAPS + assert cake_kcp.GRAPH_CAP_DEFAULT == "node_attr" assert "graph_cap" not in cake_kcp._OPTION_KEYS assert cake_kcp._options({}).graph_cap == cake_kcp.GRAPH_CAP_DEFAULT for cap in cake_kcp.GRAPH_CAPS: @@ -2052,7 +2673,7 @@ def test_cols_rule_thresholds() -> None: def _check_fused_tables(ctx, H: int) -> dict: - """The L-B tables of one rank against the separate-walk tables (``fused_walks=0``) of ``cake_kcp``.""" + """The fused-walk tables of one rank against the separate-walk tables (``fused_walks=0``) of ``cake_kcp``.""" TILE = cake_kcp.TILE tb0 = cake_kcp._Tables(ctx, H, fused_walks=0) tb1 = cake_kcp._Tables(ctx, H, fused_walks=1) @@ -2123,6 +2744,18 @@ def _check_fused_tables(ctx, H: int) -> dict: outs = [f for f in live if f.crossing_out] for name, c0 in tb0.chains.items(): c1 = tb1.chains[name] + if name in ("carry_both", "dcarry_both"): + # the fused chain families: in each table the concatenation of their two parts (whose own fused / unfused + # relation is checked below); the carry_in / dcarry_cross first elements differ by design + a, b = ( + ("carry_start", "carry_in") + if name == "carry_both" + else ("dcarry_end", "dcarry_cross") + ) + for tb in (tb0, tb1): + both = cake_kcp.Chains.concat(tb.chains[a], tb.chains[b]) + assert tb.chains[name].as_dict() == both.as_dict(), name + continue d0, d1 = c0.as_dict(), c1.as_dict() for key in ( "ncols", @@ -2215,7 +2848,7 @@ def _fused_cases(): @pytest.mark.parametrize("case", FUSED_CASES, ids=[c[0] for c in FUSED_CASES]) @pytest.mark.parametrize("H", (4, 64)) def test_fused_walk_tables(case, H) -> None: - """L-B: with ``fused_walks=1`` every live segment of every rank has a valid initial-state + """With ``fused_walks=1`` every live segment of every rank has a valid initial-state (final-adjoint) slot in one pool, the walk lists have one launch per phase, the carry plane grows by one slot per crossing fragment, and the only chain-table change is the publication of the ``S_in`` / ``dS_out`` copies.""" name, part = case @@ -2531,15 +3164,17 @@ def device() -> torch.device: 1, 3, 64, -) # segment counts launched beside the plan's own P (empty segments, P > chunks, P = 1) +) # segment counts checked beside the plan's own P (empty segments, P > chunks, P = 1) @HARNESS @pytest.mark.parametrize("case", PARITY_CASES, ids=PARITY_IDS) def test_gpu_table_parity(device, case) -> None: - """Gate 5 (R15): the CPU-built chunk_bos / chunk_len / seq_chunk_start / seq_order / seg_chunk_start are bitwise - the kda_chunk_train_layout_table_v1 / seg_table_v1 outputs on every rank of every C row, adversarial row and - ragged table (one launch of each program per rank and segment count; pad entries pre-filled with a sentinel).""" + """Gate 5 (R15): the CPU-built chunk_bos / chunk_len / seq_chunk_start / seq_order are bitwise the + kda_chunk_train_layout_table_v1 output on every rank of every C row, adversarial row and ragged table (one launch + per rank; pad entries pre-filled with a sentinel), and the boundary kernels' segment table follows its formula from + that output for the plan's P and the extra segment counts (its own program left the package with the segmented + recurrences' block table; the kernels that read it are covered by the FP64-referenced rows).""" name, W, layout = case row = kcp_shapes.ROWS[name] _frags, part = harness_partition(row, W, layout) @@ -2615,23 +3250,12 @@ def test_gpu_table_parity(device, case) -> None: failures.append((r, table)) hp = ctx.heads(row.h) for P in sorted({hp.P, *PARITY_EXTRA_P}): - seg = torch.full((N * P + 1,), -7, dtype=i32, device=device) - cake_jit.run_module( - "seg_table", - { - "seq_chunk_start": scs, - "num_seqs": N, - "segments": P, - "seg_chunk_start": seg, - }, - grid=(1, 1, 1), - device_index=device.index, - ) - torch.cuda.synchronize(device) + # the segment table from the device-built seq_chunk_start (the table program's output) must equal the + # CPU formula from the CPU-built one; at the plan's P it is the head plan's table want_seg = kp.seg_chunk_start(list(ctx.seq_chunk_start), P) if P == hp.P: assert tuple(want_seg) == hp.seg_chunk_start - equal = seg.cpu().tolist() == want_seg + equal = kp.seg_chunk_start(got["seq_chunk_start"], P) == want_seg checked += 1 _evidence( "parity", @@ -2830,7 +3454,9 @@ def test_fake_backend_transfer_matches_lane_p(fake) -> None: ) -@pytest.mark.parametrize("graph_cap", ["none", "capture_pref"]) +@pytest.mark.parametrize( + "graph_cap", ["none", "capture_pref", "node_attr", "node_attr_post"] +) def test_gpu_graph_replay_is_bitwise(device, monkeypatch, graph_cap) -> None: """The graph-replay option, KDA module, one rank (``world == 1``: no collective): under the public keyword ``graph="replay"`` the first call of a validity key runs eagerly, the second is captured and replayed once, later @@ -2840,7 +3466,10 @@ def test_gpu_graph_replay_is_bitwise(device, monkeypatch, graph_cap) -> None: Both cache-config forms of a capture (knob ``graph_cap``: ``"none"`` / ``"capture_pref"``) capture and replay under the shipped capped geometry: the per-launch bracket is skipped inside the capture (``capture_skips`` grows), the eager calls record ``cache_config="bracket"`` and the captured / replayed calls - record the arm (with the set / reset evidence of the capture-wide preference under ``"capture_pref"``).""" + record the arm (with the set / reset evidence of the capture-wide preference under ``"capture_pref"``). The kernel-node + forms ``"node_attr"`` / ``"node_attr_post"`` capture and replay bitwise too: no + launch is skipped (``capture_skips`` unchanged), every captured tape launch gets the carveout attribute (``node_attrs`` + grows by the tape launches of the captured calls, recorded as ``cache_config_nodes``).""" monkeypatch.setenv(cake_kcp.TEST_KNOBS_ENV, "1") knobs = {"_kcp_test_knobs": {"graph_cap": graph_cap}} backend = cake_kcp.get_backend() @@ -2905,7 +3534,8 @@ def run(**kw): ), recs0 fallback0 = cake_kcp.GRAPH_STATS["fallback"] skips0 = programs.capture_skips if capped else None - events, ccs, brackets = [], [], [] + nodes0 = programs.node_attrs if capped else None + events, ccs, brackets, nodes, tapes = [], [], [], [], [] try: for _ in range(4): res, ev, recs = run(graph="replay", **knobs) @@ -2914,6 +3544,18 @@ def run(**kw): brackets.append( (recs[0]["cache_config_bracket"], recs[1]["cache_config_bracket"]) ) + nodes.append( + (recs[0].get("cache_config_nodes"), recs[1].get("cache_config_nodes")) + ) + tapes.append( + tuple( + sum( + r["sequence"].count(p) + for p in ("kcp_fwdh_tape", "kcp_dhu_tape") + ) + for r in recs + ) + ) assert all(torch.equal(a, b) for a, b in zip(res, ref, strict=True)), events assert [e[0] for e in events] == ["eager", "capture", "replay", "replay"], ( events @@ -2924,7 +3566,33 @@ def run(**kw): cap_cc = graph_cap if capped else "none" assert [c[0] for c in ccs] == [eager_cc, cap_cc, cap_cc, cap_cc], ccs assert [c[1] for c in ccs] == [eager_cc, eager_cc, cap_cc, cap_cc], ccs - if capped: + if capped and graph_cap in cake_kcp.GRAPH_NODE_CAPS: + # kernel-node cap forms: nothing launched bare; every captured tape launch carries the node attribute + assert programs.capture_skips == skips0, (programs.capture_skips, skips0) + assert programs.node_attrs == nodes0 + tapes[1][0] + tapes[2][1], ( + programs.node_attrs, + nodes0, + tapes, + ) + assert ( + nodes[1][0] == tapes[1][0] >= 1 and nodes[2][1] == tapes[2][1] >= 1 + ), ( + nodes, + tapes, + ) + assert [n[0] for n in nodes[1:]] == [tapes[1][0]] * 3, ( + nodes + ) # the replays carry the capture's record + assert brackets[1][0] is None and brackets[2][1] is None, brackets + evidence = programs.node_attr_log + assert evidence and all( + e["rc"] == 0 + and e["node_type"] == 0 + and e["readback"] == cake_kcp.GRAPH_NODE_CARVEOUT + for log in evidence.values() + for e in log + ), evidence + elif capped: assert ( programs.capture_skips > skips0 ) # the capped programs launched bare inside the captures @@ -2968,6 +3636,14 @@ def run(**kw): "capture_skips_delta": (programs.capture_skips - skips0) if capped else None, + "node_attrs_delta": (programs.node_attrs - nodes0) if capped else None, + "cache_config_nodes": nodes, + "tape_launches": tapes, + "node_attr_log": { + f"{k[0]}:{k[1]}": v for k, v in programs.node_attr_log.items() + } + if capped + else None, "bitwise": True, "fallbacks": cake_kcp.GRAPH_STATS["fallback"] - fallback0, }, @@ -3106,3 +3782,894 @@ def fake_set(device_index, config): walk_geometry=cake_kcp.WALK_GEOMETRY_DEFAULT, ) assert cake_kcp._graph_cap_bracket(none_) == ("none", None) + + +def test_host_trim_options_passthrough() -> None: + """Trimmed host path: ``chunk_kda_kcp`` validates the options once; the autograd node passes the validated ``KcpOptions`` + through ``_kcp_options`` and ``_options`` returns it unchanged (it travels alone).""" + opts = cake_kcp._options({"exchange": "a2a"}) + assert cake_kcp._options({"_kcp_options": opts}) is opts + with pytest.raises(TypeError): + cake_kcp._options({"_kcp_options": opts, "exchange": "a2a"}) + with pytest.raises(TypeError): + cake_kcp._options({"_kcp_options": {"exchange": "a2a"}}) + + +def test_host_trim_workspace_fast_paths() -> None: + """Trimmed host path: ``ensure(host_trim=True)`` returns the last tables for the same (ctx, H, device, backward, knobs) and + takes the full path otherwise; ``pool(cached=True)`` returns the same ``[:n]`` view while the pool is unchanged.""" + part = _small_partition(4, lengths=(3000,), layout="contiguous") + ctx = kp.build_kcp_context(part, None, 1, "cpu", 4, sm_count=SM_COUNT) + ws = cake_kcp.KcpWorkspace(local=False) + tb = ws.ensure(ctx, 64, device="cpu", backward=False, host_trim=True) + assert ws.ensure(ctx, 64, device="cpu", backward=False, host_trim=True) is tb + assert ws._last_ensure is not None and ws._last_ensure[0] is ctx + assert not ws.has("b1_send") + assert ( + ws.ensure(ctx, 64, device="cpu", backward=True, host_trim=True) is tb + ) # another signature: the full path + assert ws.has("b1_send") + assert ( + ws.ensure(ctx, 64, device="cpu", backward=False) is tb + ) # the untrimmed path never consults the fast path + n = tb.sizes["c1_send"] + v1 = ws.pool("c1_send", n, cached=True) + assert ws.pool("c1_send", n, cached=True) is v1 + v0 = ws.pool("c1_send", n) + assert v0 is not v1 and v0.data_ptr() == v1.data_ptr() and v0.shape == v1.shape + ws._pools["c1_send"] = torch.empty_like( + ws._pools["c1_send"] + ) # a re-sized pool invalidates the cached view + v2 = ws.pool("c1_send", n, cached=True) + assert ( + v2 is not v1 + and v2.data_ptr() == ws._pools["c1_send"].data_ptr() + and v2.shape == v0.shape + ) + + +def test_host_trim_backward_walk_sublists_are_table_owned() -> None: + """Trimmed host path: the backward's tape-replay walk lists ``s % P == 0`` / ``!= 0`` of ``walk_all`` are the table set's own + list objects (one device copy per table set on the trimmed path, not one per call), and an identity-keyed cache + fed with fresh objects stays bounded.""" + for lengths, layout in ( + ((3000,), "contiguous"), + ((700, 1, 1300, 999), "zigzag_per_document"), + ((40, 40), "zigzag"), + ): + part = _small_partition(4, lengths=lengths, layout=layout) + ctx = kp.build_kcp_context(part, None, 1, "cpu", 4, sm_count=SM_COUNT) + tb = cake_kcp._Tables(ctx, 64, fused_walks=1) + p0, rest = tb.walk_all_p0, tb.walk_all_rest + assert tb.walk_all_p0 is p0 and tb.walk_all_rest is rest + assert p0 == [s for s in tb.walk_all if s % tb.P == 0] + assert rest == [s for s in tb.walk_all if s % tb.P != 0] + assert sorted(p0 + rest) == sorted(tb.walk_all) + if p0: + d0 = tb.dev_list(p0, "cpu") + assert tb.dev_list(tb.walk_all_p0, "cpu") is d0 and d0.tolist() == p0 + if tb.P == 1: + assert rest == [] and p0 == list(tb.walk_all) + # the bound: fresh list objects miss every time and the cache is emptied at ID_CACHE_MAX entries + keep = [] + for i in range(cake_kcp._Tables.ID_CACHE_MAX + 5): + fresh = [i, i + 1] + keep.append(fresh) + tb.dev_list(fresh, "cpu") + assert 0 < len(tb._dev_lists) <= cake_kcp._Tables.ID_CACHE_MAX + assert tb.dev_list(tb.walk_offset, "cpu") is tb.dev_list(tb.walk_offset, "cpu") + + +def test_host_trim_table_identity_caches() -> None: + """Trimmed host path: the table set's identity-keyed device copies (walk lists, the segment table, tile tables, Indexed + metadata, chain tables) hold the same integers as the content-keyed copies of today's path.""" + from flashinfer.experimental.cake_kda_train import cake_kcp_jit + + part = _small_partition(4, lengths=(3000,), layout="contiguous") + ctx = kp.build_kcp_context(part, None, 1, "cpu", 4, sm_count=SM_COUNT) + tb = cake_kcp._Tables(ctx, 64, fused_walks=1) + assert tb.P > 1 + w1 = tb.dev_list(tb.walk_offset, "cpu") + assert tb.dev_list(tb.walk_offset, "cpu") is w1 + assert w1.dtype == torch.int32 and w1.tolist() == [int(s) for s in tb.walk_offset] + copy = list(tb.walk_offset) + w2 = tb.dev_list(copy, "cpu") + assert w2 is not w1 and torch.equal( + w2, w1 + ) # another list object: another entry, the same integers + seg = tb.dev_list(tb.hp.seg_chunk_start, "cpu") + assert seg.tolist() == [int(x) for x in tb.hp.seg_chunk_start] + cpu = torch.device("cpu") + t_id = tb.device_tiles(tb.mf_tiles, cpu, by_identity=True) + assert tb.device_tiles(tb.mf_tiles, cpu, by_identity=True) is t_id + assert torch.equal(t_id, tb.device_tiles(tb.mf_tiles, cpu)) + meta = tb.indexed_meta(tb.h0_index, tb.s_in_tiles, 64, "cpu") + assert tb.indexed_meta(tb.h0_index, tb.s_in_tiles, 64, "cpu") is meta + plain = cake_kcp.Indexed( + torch.empty(0), + tuple(int(i) for i in tb.h0_index), + tuple(tuple(int(t) for t in row) for row in tb.s_in_tiles), + ) + assert (meta.index, meta.tiles, meta.addr) == ( + plain.index, + plain.tiles, + plain.addressing(64), + ) + assert meta.dev_index.dtype == torch.int32 and meta.dev_index.tolist() == list( + plain.index + ) + ix = cake_kcp.Indexed(torch.empty(0), meta.index, meta.tiles, meta) + assert ix.addressing(64) == plain.addressing(64) and ix.meta is meta + assert cake_kcp.Indexed(torch.empty(0), meta.index, meta.tiles).meta is None + for name, ch in tb.chains.items(): + mine = ch.device_tables("cpu") + assert ch.device_tables("cpu") is mine + ref = cake_kcp_jit._chain_tables(ch, "cpu") + assert set(mine) == set(ref), name + for k in ref: + assert mine[k].dtype == ref[k].dtype and torch.equal(mine[k], ref[k]), ( + name, + k, + ) + empty = cake_kcp.Chains(8) + e_mine, e_ref = empty.device_tables("cpu"), cake_kcp_jit._chain_tables(empty, "cpu") + for k in e_ref: + assert e_mine[k].dtype == e_ref[k].dtype and torch.equal(e_mine[k], e_ref[k]), k + + +def test_host_trim_kernel_backend_run_skips_the_discarded_abi() -> None: + """Trimmed host path: under host_trim a launch goes to the provider's ``run_program`` without the positional ABI (which that + form discards), keeping the D1 bookkeeping of the merge (alignment check, ``cols_per_cta`` None, ``merge_log``); + an explicit cols rule takes today's path (ABI + the ``_c`` record).""" + calls: list = [] + + class FakeJit: + def run_program(self, name, args, *, device_index): + calls.append(("run_program", name, args.get("chains"), device_index)) + + def run_module(self, record, values, *, grid, device_index): + calls.append(("run_module", record, grid, device_index)) + + ch = cake_kcp.Chains(cake_kcp.D) + ch.add([(0, -1, -1)], x0=-1, y_final=0) + big = cake_kcp.Chains( + cake_kcp.D + ) # > COLS_RULE_CHAINS state chains: by_chains asks for the explicit c16 record + for i in range(cake_kcp.COLS_RULE_CHAINS + 1): + big.add([(4 * i, -1, -1)], x0=-1, y_final=4 * i) + y, m = torch.zeros(2 * cake_kcp.TILE), torch.zeros(2 * cake_kcp.TILE) + + def call( + host_trim: int, + cols_rule: str = "fixed", + geometry: str = cake_kcp.WALK_GEOMETRY_DEFAULT, + ): + return types.SimpleNamespace( + host_trim=host_trim, + cols_rule=cols_rule, + merge_log=[], + device=torch.device("cpu"), + walk_geometry=geometry, + scratch={}, + precision_mode=0, + transfer_walk="tf32_hilo", + ) + + backend = cake_kcp.KernelBackend(programs=cake_kcp._FiPrograms(FakeJit())) + c1 = call(1) + args = { + "family": "fwd", + "m": m, + "c": None, + "x0": None, + "y": y, + "chains": ch, + "transpose_m": 0, + "write_states": 1, + "_call": c1, + } + backend.run("kcp_affine_merge", args) + assert calls == [("run_program", "kcp_affine_merge", ch, None)] + assert ( + args["cols_per_cta"] is None + and c1.merge_log == [("fwd", cake_kcp.D, 1, None)] + and id(ch) in backend._checked_chains + ) + calls.clear() + c0 = call(0) + backend.run("kcp_affine_merge", {**args, "_call": c0}) + assert calls == [ + ("run_program", "kcp_affine_merge", ch, None) + ] and c0.merge_log == [("fwd", cake_kcp.D, 1, None)] + calls.clear() + c2 = call(1, cols_rule="by_chains") + backend.run("kcp_affine_merge", {**args, "chains": big, "_call": c2}) + n_cta = (cake_kcp.COLS_RULE_CHAINS + 1) * (cake_kcp.D // 16) + assert calls == [("run_module", "kcp_affine_merge_c16", (n_cta, 1, 1), None)] + assert ( + c2.merge_log == [("fwd", cake_kcp.D, cake_kcp.COLS_RULE_CHAINS + 1, 16)] + and id(big) in backend._checked_chains + ) + calls.clear() + c3 = call(1, geometry="shipped") + backend.run("kcp_fwdh_tape", {"walk": [0, 1], "h0": None, "out": None, "_call": c3}) + assert calls == [("run_program", "kcp_fwdh_tape", None, None)] + + +class _FakeGraphRt: + """A CUDA-runtime double of the capture-graph entry points of the kernel-node cap forms: one + active capture whose next operation depends on one kernel node; records every call.""" + + def __init__( + self, *, status: int = 1, n_deps: int = 1, node_type: int = 0, rc_set: int = 0 + ) -> None: + self.status, self.n_deps, self.node_type, self.rc_set = ( + status, + n_deps, + node_type, + rc_set, + ) + self.calls: list[tuple] = [] + self.values: dict[int, int] = {} + self.next_node = 0x1000 + self._keep: list = [] + + def cudaRuntimeGetVersion(self, ver): + ver._obj.value = 13000 # the double speaks CUDA 13: the plain export has the 7-argument (v3) form + return 0 + + def cudaStreamGetCaptureInfo(self, stream, status, cid, graph, deps, edges, n_deps): + node = self.next_node + self.next_node += 0x10 + status._obj.value = self.status + cid._obj.value = 7 + graph._obj.value = 0xBEEF + handle = ctypes.c_void_p(node) + self._keep.append( + handle + ) # the array the runtime would own: valid until the next capture call + deps._obj.contents = handle + n_deps._obj.value = self.n_deps + self.calls.append(("info", stream.value)) + return 0 + + def cudaGraphNodeGetType(self, node, node_type): + node_type._obj.value = self.node_type + self.calls.append(("type", node.value)) + return 0 + + def cudaGraphKernelNodeSetAttribute(self, node, attr, value): + self.calls.append(("set", node.value, attr, value._obj[0])) + if self.rc_set: + return self.rc_set + self.values[node.value] = value._obj[0] + return 0 + + def cudaGraphKernelNodeGetAttribute(self, node, attr, value): + self.calls.append(("get", node.value, attr)) + value._obj[0] = self.values.get(node.value, 0xFFFFFFFF) + return 0 + + def cudaGetLastError(self): + self.calls.append(("lasterror",)) + return 0 + + +def _fake_node_call(form: str) -> types.SimpleNamespace: + return types.SimpleNamespace( + walk_geometry=cake_kcp.WALK_GEOMETRY_DEFAULT, + scratch={}, + opts=types.SimpleNamespace(graph_cap=form), + ) + + +def _fake_stream(monkeypatch) -> None: + monkeypatch.setattr( + torch.cuda, + "current_stream", + lambda device_index=None: types.SimpleNamespace(cuda_stream=0x5151), + ) + + +def test_capped_launch_sets_the_node_attribute_inside_a_capture(monkeypatch) -> None: + """The kernel-node cap ``graph_cap="node_attr"``: inside a stream capture a capped program + launches bare and the kernel node it just added (``cudaStreamGetCaptureInfo``, ``_v2`` on CUDA 12: one dependency node of an active + capture, type kernel) gets the preferred shared-memory carveout launch attribute ``GRAPH_NODE_CARVEOUT`` with a + read-back; nothing is counted as a capture skip; the call records the number of capped launches; the evidence log keeps + the first records per (program, device). ``"node_attr_post"`` keeps the node handle for :meth:`apply_node_caps` instead, + which sets the attributes after the capture. Without the graph entry points both forms report ``unavailable`` and the + launch falls back to the bare capture-skip form. A capture-info reply without exactly one kernel-node dependency degrades + the same way: uncapped launch, one warning per process, the reason on the call.""" + rt = _FakeGraphRt() + monkeypatch.setattr(cake_kcp, "_CUDART_CACHE", {"rt": rt, "graph": rt}) + monkeypatch.setattr(cake_kcp, "_capturing", lambda: True) + _fake_stream(monkeypatch) + stats0 = dict(cake_kcp.CACHE_CONFIG_STATS) + jit = _FakeJit() + programs = cake_kcp._FiPrograms(jit) + program = cake_kcp.WALK_CAP_PROGRAMS[cake_kcp.WALK_GEOMETRY_DEFAULT][0] + call = _fake_node_call("node_attr") + programs.launch(program, [], grid=4, device_index=0, args={"_call": call}) + programs.launch(program, [], grid=4, device_index=0, args={"_call": call}) + assert jit.launched == [(program, 4, 0, 0, "tf32_hilo")] * 2 + assert programs.capture_skips == 0 and programs.node_attrs == 2 + assert ( + call.scratch["node_attr_launches"] == 2 + and "node_attr_pending" not in call.scratch + ) + kinds = [c[0] for c in rt.calls] + assert kinds == ["info", "type", "set", "get"] * 2, kinds + sets = [c for c in rt.calls if c[0] == "set"] + assert all( + c[2] == cake_kcp._LAUNCH_ATTR_SHMEM_CARVEOUT + and c[3] == cake_kcp.GRAPH_NODE_CARVEOUT + for c in sets + ), sets + assert [c[1] for c in sets] == [ + 0x1000, + 0x1010, + ] # one node per launch, the just-captured one + log = programs.node_attr_log[(program, 0)] + assert [e["node"] for e in log] == [0x1000, 0x1010] + assert all( + e["rc"] == 0 + and e["readback"] == cake_kcp.GRAPH_NODE_CARVEOUT + and e["node_type"] == 0 + and e["status"] == 1 + and e["n_deps"] == 1 + and e["stream"] == 0x5151 + and e["immediate"] is True + for e in log + ), log + assert ( + cake_kcp.CACHE_CONFIG_STATS["node_attr_launches"] + == stats0["node_attr_launches"] + 2 + ) + assert cake_kcp.CACHE_CONFIG_STATS["capture_skips"] == stats0["capture_skips"] + # the evidence log keeps four records per (program, device); the counters keep counting + for _ in range(4): + programs.launch(program, [], grid=4, device_index=0, args={"_call": call}) + assert len(programs.node_attr_log[(program, 0)]) == 4 and programs.node_attrs == 6 + # the bracket form of a captured call: node_attr / node_attr_post, or unavailable without the entry points + cap = types.SimpleNamespace( + opts=types.SimpleNamespace(graph_cap="node_attr"), + walk_geometry=cake_kcp.WALK_GEOMETRY_DEFAULT, + ) + assert cake_kcp._graph_cap_bracket(cap) == ("node_attr", None) + cap.opts.graph_cap = "node_attr_post" + assert cake_kcp._graph_cap_bracket(cap) == ("node_attr_post", None) + cap.walk_geometry = "shipped" + assert cake_kcp._graph_cap_bracket(cap) == ("none", None) + # node_attr_post: the handles are kept on the call and set afterwards + post = _fake_node_call("node_attr_post") + programs.launch(program, [], grid=4, device_index=0, args={"_call": post}) + programs.launch(program, [], grid=4, device_index=0, args={"_call": post}) + assert ( + post.scratch["node_attr_launches"] == 0 + and len(post.scratch["node_attr_pending"]) == 2 + ) + n_set = len([c for c in rt.calls if c[0] == "set"]) + assert programs.apply_node_caps(post) == 2 + assert len([c for c in rt.calls if c[0] == "set"]) == n_set + 2 + assert ( + post.scratch["node_attr_launches"] == 2 + and "node_attr_pending" not in post.scratch + ) + assert programs.node_attrs == 8 + # a failing set inside the capture raises (the capture then falls back to the eager path, recorded by _graph_capture) + bad = _FakeGraphRt(rc_set=1) + monkeypatch.setattr(cake_kcp, "_CUDART_CACHE", {"rt": bad, "graph": bad}) + programs2 = cake_kcp._FiPrograms(_FakeJit()) + with pytest.raises(RuntimeError, match="cudaGraphKernelNodeSetAttribute"): + programs2.launch( + program, + [], + grid=4, + device_index=0, + args={"_call": _fake_node_call("node_attr")}, + ) + assert ("lasterror",) in bad.calls + # a capture whose dependency set is not one kernel node DEGRADES: the launch is + # issued uncapped, one RuntimeWarning per process, the call records the reason -- never an exception / eager fallback + odd = _FakeGraphRt(n_deps=2) + monkeypatch.setattr(cake_kcp, "_CUDART_CACHE", {"rt": odd, "graph": odd}) + odd_call = _fake_node_call("node_attr") + skips0 = programs2.capture_skips + degraded0 = cake_kcp.CACHE_CONFIG_STATS["node_attr_degraded"] + with pytest.warns(RuntimeWarning, match="one dependency node"): + programs2.launch( + program, + [], + grid=4, + device_index=0, + args={"_call": odd_call}, + ) + assert programs2.capture_skips == skips0 + 1 + assert cake_kcp.CACHE_CONFIG_STATS["node_attr_degraded"] == degraded0 + 1 + assert "one dependency node" in odd_call.scratch["node_attr_degraded"] + assert "node_attr_launches" not in odd_call.scratch + assert not any(c[0] == "set" for c in odd.calls) + with warnings.catch_warnings(): + warnings.simplefilter("error") # remembered: no second warning in the process + programs2.launch( + program, + [], + grid=4, + device_index=0, + args={"_call": odd_call}, + ) + assert programs2.capture_skips == skips0 + 2 and programs2.node_attrs == 0 + # without the graph entry points: unavailable, and the launch takes the capture-skip form + monkeypatch.setattr(cake_kcp, "_CUDART_CACHE", {"rt": object(), "graph": None}) + programs3 = cake_kcp._FiPrograms(_FakeJit()) + plain = _fake_node_call("node_attr") + programs3.launch(program, [], grid=4, device_index=0, args={"_call": plain}) + assert ( + programs3.capture_skips == 1 + and programs3.node_attrs == 0 + and plain.scratch == {} + ) + cap.walk_geometry = cake_kcp.WALK_GEOMETRY_DEFAULT + cap.opts.graph_cap = "node_attr" + assert cake_kcp._graph_cap_bracket(cap) == ("unavailable", None) + # the shipped forms are untouched by the kernel-node forms (the capture skip) + none_ = _fake_node_call("none") + programs3.launch(program, [], grid=4, device_index=0, args={"_call": none_}) + assert programs3.capture_skips == 2 and none_.scratch == {} + + +def test_graph_entry_points_missing_degrades(monkeypatch) -> None: + """A CUDA runtime without the capture-graph entry points: one ``RuntimeWarning`` naming them, ``_cudart_graph_available()`` + is ``None`` (remembered), the shipped forms are unaffected.""" + monkeypatch.setattr(cake_kcp, "_CUDART_CACHE", {"rt": object()}) + monkeypatch.setattr(cake_kcp, "_cudart", lambda: types.SimpleNamespace()) + with pytest.warns(RuntimeWarning, match="cudaStreamGetCaptureInfo") as rec: + assert cake_kcp._cudart_graph_available() is None + assert len([w for w in rec if issubclass(w.category, RuntimeWarning)]) == 1 + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert ( + cake_kcp._cudart_graph_available() is None + ) # remembered: no second warning + assert cake_kcp._CUDART_CACHE["graph"] is None + + +def test_capture_info_export_name_resolution() -> None: + """CUDA 12.3+ exports the 7-argument stream-capture-info call as ``cudaStreamGetCaptureInfo_v3``, CUDA 11.3-12.x the + 6-argument form as ``..._v2``, and CUDA 13 (13000+) the 7-argument form under the plain name (observed: libcudart.so.13 + exports neither suffixed name); the module resolves the first name present and its arity, and a plain name on a pre-13 runtime + (the pre-11.3 form without dependency nodes) is unusable.""" + assert cake_kcp._CUDART_CAPTURE_INFO_NAMES == ( + "cudaStreamGetCaptureInfo_v3", + "cudaStreamGetCaptureInfo_v2", + "cudaStreamGetCaptureInfo", + ) + ns = types.SimpleNamespace + assert ( + cake_kcp._capture_info_name( + ns(cudaStreamGetCaptureInfo_v2=1, cudaStreamGetCaptureInfo=2) + ) + == "cudaStreamGetCaptureInfo_v2" + ) + assert ( + cake_kcp._capture_info_name( + ns(cudaStreamGetCaptureInfo_v3=1, cudaStreamGetCaptureInfo_v2=1) + ) + == "cudaStreamGetCaptureInfo_v3" + ) + assert ( + cake_kcp._capture_info_name(ns(cudaStreamGetCaptureInfo=2)) + == "cudaStreamGetCaptureInfo" + ) + assert cake_kcp._capture_info_name(ns()) is None + assert cake_kcp._capture_info_arity(ns(), "cudaStreamGetCaptureInfo_v3") == 7 + assert cake_kcp._capture_info_arity(ns(), "cudaStreamGetCaptureInfo_v2") == 6 + + def version(v): + def fn(ver): + ver._obj.value = v + return 0 + + return fn + + assert ( + cake_kcp._capture_info_arity( + ns(cudaRuntimeGetVersion=version(13000)), "cudaStreamGetCaptureInfo" + ) + == 7 + ) + assert ( + cake_kcp._capture_info_arity( + ns(cudaRuntimeGetVersion=version(12090)), "cudaStreamGetCaptureInfo" + ) + == 0 + ) + assert ( + cake_kcp._capture_info_arity(ns(), "cudaStreamGetCaptureInfo") == 0 + ) # no version call: unusable + assert cake_kcp._CUDART_GRAPH_ENTRY_POINTS == ( + "cudaGraphNodeGetType", + "cudaGraphKernelNodeSetAttribute", + "cudaGraphKernelNodeGetAttribute", + ) + rt = _FakeGraphRt() # the double speaks CUDA 13 + assert cake_kcp._runtime_version(rt) == 13000 + assert cake_kcp._capture_info_arity(rt, cake_kcp._capture_info_name(rt)) == 7 + + +def test_host_trim_templates_live_on_the_arena_and_weakly_key_the_table_set() -> None: + """Template-cache ownership: the trimmed path's ``begin_call`` template (the call-independent values of + one (arena, table set, direction)) is cached on the local arena whose views it holds, keyed by the table set's + identity with a WEAK reference to it (a hit is valid only while that table set object is alive and identical); + the backend keeps no template cache, so a table set or an arena dies when its owner drops it -- a training loop + with a new context / workspace per step retains nothing. Dead entries are purged on the next miss.""" + + class Arena: # the local arena's surface used by begin_call: plan geometry and the layout views + def __init__(self, ctx, H: int) -> None: + self.plan = types.SimpleNamespace( + S=int(ctx.T_r), + N=int(ctx.n_frag), + H=int(H), + NC_pad=int(ctx.NC_pad), + R=int(ctx.R), + NSEG=4, + P=2, + seg_arg=2, + proportional=0, + ) + n = max(int(ctx.NC_pad), 1) + self._views = { + mode: types.SimpleNamespace( + h=torch.zeros(n, 2, cake_kcp.D), + dh=torch.zeros(n, 2, cake_kcp.D), + w=torch.zeros(4), + ) + for mode in (cake_kcp.LOCAL_LAYOUT_FWD, cake_kcp.LOCAL_LAYOUT_BWD) + } + + def views(self, mode: str): + return self._views[mode] + + def call_for(arena, ctx, tb, H: int, backward: bool, host_trim: int = 1): + return types.SimpleNamespace( + local=arena, + ctx=ctx, + tb=tb, + T_r=int(ctx.T_r), + H=int(H), + backward=backward, + host_trim=host_trim, + ops={}, + scale=1.0, + scratch={}, + local_values=None, + ) + + part = _small_partition(4, lengths=(3000,), layout="contiguous") + ctx = kp.build_kcp_context(part, None, 1, "cpu", 4, sm_count=SM_COUNT) + H = 64 + tb = cake_kcp._Tables(ctx, H, fused_walks=1) + arena = Arena(ctx, H) + backend = cake_kcp.KernelBackend() + assert not hasattr( + backend, "_local_templates" + ) # no template cache on the long-lived backend + c1 = call_for(arena, ctx, tb, H, False) + backend.begin_call(c1) + c2 = call_for(arena, ctx, tb, H, False) + backend.begin_call(c2) + assert ( + c1.local_values["h2d"] is c2.local_values["h2d"] + ) # the cached template: the same view objects per call + assert c1.local_values is not c2.local_values and c1.scratch[ + "views" + ] is arena.views("fwd") + cache = getattr(arena, cake_kcp.TEMPLATE_CACHE_ATTR) + assert set(cache) == {(id(tb), False)} + assert ( + isinstance(cache[(id(tb), False)][0], weakref.ref) + and cache[(id(tb), False)][0]() is tb + ) + c3 = call_for(arena, ctx, tb, H, True) + backend.begin_call(c3) + assert set(cache) == {(id(tb), False), (id(tb), True)} and "dh2d" in c3.local_values + # the table set dies although the arena and the backend live on: only a weak reference points at it + wr_tb = weakref.ref(tb) + del tb, c1, c2, c3 + gc.collect() + assert wr_tb() is None + # a new table set on the same arena misses, purges the dead entries and leaves one live entry + ctx2 = kp.build_kcp_context(part, None, 1, "cpu", 4, sm_count=SM_COUNT) + tb2 = cake_kcp._Tables(ctx2, H, fused_walks=1) + c4 = call_for(arena, ctx2, tb2, H, False) + backend.begin_call(c4) + assert set(cache) == {(id(tb2), False)} and cache[(id(tb2), False)][0]() is tb2 + # the arena dies with its templates; the backend holds nothing of it + wr_arena = weakref.ref(arena) + del arena, c4, cache + gc.collect() + assert wr_arena() is None + # the previous host path (host_trim = 0) builds the values per call and caches nothing on the arena + arena_b = Arena(ctx2, H) + c0 = call_for(arena_b, ctx2, tb2, H, False, host_trim=0) + backend.begin_call(c0) + assert not hasattr(arena_b, cake_kcp.TEMPLATE_CACHE_ATTR) + assert c0.local_values["h2d"].data_ptr() == arena_b.views("fwd").h.data_ptr() + + +LOCAL_RECORD_CASES = [ + (4, "sm_103a", 152), + (64, "sm_103a", 152), + (64, "sm_107a", 216), + (4, "sm_100a", 148), +] + + +@pytest.mark.parametrize( + "H,arch,sm", + LOCAL_RECORD_CASES, + ids=[f"h{h}-{a}-{s}" for h, a, s in LOCAL_RECORD_CASES], +) +def test_local_stage_values_cover_the_records(H, arch, sm, monkeypatch) -> None: + """The local stages against the existing backend's registry on the CPU: the kernel backend's value dict of a + forward (no token-row copies: the saved set carries none), a saved-mode backward and a replay backward binds + every logical name of every + record it launches -- the plan's records on the forward path without a saved set and on the replay path, the + direct WY record where the plan fuses the WY representation into a recurrence the KCP does not run, and the + KCP-owned ``kcp_fwd_o`` / ``kcp_dav`` on their grids; the two backward modes launch the same records on the same + grids; the intra record that computes the layout itself gets the arena's scratch table set (and the workspace + bytes count it) while every other program reads the context's tables.""" + monkeypatch.setenv(cake_kcp.SPEC3_ENV, "0") + part = _small_partition(2) + ctx = kp.build_kcp_context(part, None, 0, "cpu", 4, sm_count=sm) + assert ctx.T_r > 0 + plan = cake_backend.plan(ctx.T_r, ctx.n_frag, H, sm, arch) + tb = cake_kcp._Tables(ctx, H, fused_walks=1) + launches: list[tuple[str, str, tuple[int, int, int], dict]] = [] + + class Arena: # the arena's surface used by begin_call / stage: the plan, the layout views, the launcher + def __init__(self) -> None: + self.plan = plan + self.device = torch.device("cpu") + self.device_index = 0 + self.arch = arch + self.launcher = self._launch + self._views: dict = {} + + def views(self, mode: str): + ns = self._views.get(mode) + if ns is None: + ns = types.SimpleNamespace() + for r in plan.layout(mode): + small = ( + torch.zeros(2, cake_kcp.D) + if r.name in ("h", "dh") + else torch.zeros(1) + ) + setattr(ns, r.name, small) + self._views[mode] = ns + return ns + + def _launch(self, record, values, *, grid, device_index) -> None: + launches.append(("backend", record, grid, dict(values))) + + class Provider: + def run_record(self, record, values, *, grid, device_index) -> None: + launches.append(("kcp", record, grid, dict(values))) + + arena = Arena() + backend = cake_kcp.KernelBackend(programs=Provider()) + + def call_for(backward: bool, *, ops, out=None, grads=None): + return types.SimpleNamespace( + local=arena, + ctx=ctx, + tb=tb, + T_r=int(ctx.T_r), + H=H, + backward=backward, + host_trim=1, + ops=ops, + scale=0.5, + scratch={}, + local_values=None, + out=out, + grads=grads, + device=torch.device("cpu"), + ) + + def run(call, stages) -> list: + backend.begin_call(call) + for name, extra in stages: + backend.stage(name, call, **extra) + done = list(launches) + launches.clear() + return done + + ops = {n: torch.zeros(1) for n in ("q", "k", "v", "g", "beta")} + fwd_call = call_for(False, ops=ops, out=torch.zeros(1)) + fwd = run(fwd_call, (("fwd_intra", {}), ("wy_fwd", {}), ("fwd_o", {}))) + bwd_ops = {**ops, "dout": torch.zeros(1)} + grads = {n: torch.zeros(1) for n in ("dq", "dk", "dv", "dg", "dbeta")} + saved = types.SimpleNamespace(S_in=torch.zeros(1), c1_recv=None, M_seg=None) + bwd_stages = ( + ("fwd_intra", {}), + ("wy", {}), + ("dav", {}), + ("dqkg", {}), + ("glm_intra", {}), + ("gate_epilogue", {}), + ) + bwd = run( + call_for(True, ops=bwd_ops, grads=grads), + (("bwd_prep", {"saved": saved}), *bwd_stages), + ) + replay = run( + call_for(True, ops=bwd_ops, grads=grads), + (("bwd_prep", {"saved": None}), *bwd_stages), + ) + # the two backward modes: the same records on the same grids + assert [(k, r, g) for k, r, g, _ in replay] == [(k, r, g) for k, r, g, _ in bwd] + # every launched record of the existing backend's registry finds its logical names in the value dict + for kind, record, _grid, values in fwd + bwd: + if kind == "backend": + missing = set(cake_jit.MODULES[record]["logical"].values()) - set(values) + assert not missing, (record, sorted(missing)) + # the records of record and their paths + f, b = cake_kcp.LOCAL_PATH_FWD, cake_kcp.LOCAL_PATH_BWD + + def expect(stage: str, path: str) -> str: + if stage in cake_kcp.WY_DIRECT_RECORDS and not plan.dispatched(stage, path): + return cake_kcp.WY_DIRECT_RECORDS[stage] + return plan.record(stage, path) + + assert [r for _, r, _, _ in fwd] == [ + expect("fwd_intra", f), + expect("wy_fwd", f), + "kcp_fwd_o", + ] + assert [r for _, r, _, _ in bwd] == [ + expect("bwd_prep", b), + expect("fwd_intra", b), + expect("wy", b), + "kcp_dav", + expect("dqkg", b), + expect("glm_intra", b), + expect("gate_epilogue", b), + ] + if ( + H == 64 and arch != "sm_107a" + ): # the plan fuses the WY representation into its recurrences there + assert ( + not plan.dispatched("wy_fwd", f) + and [r for _, r, _, _ in fwd][1] == "wy_fwd" + ) + # the write switches (no token-row copies on this route: the saved set carries none), the scratch tables and + # the grids + intra, bintra = fwd[0][3], bwd[1][3] + assert (intra["write_ext"], intra["write_int"], intra["write_prep"]) == (0, 1, 1) + assert (bintra["write_ext"], bintra["write_int"], bintra["write_prep"]) == (0, 1, 0) + assert ( + intra["aqk_ext"] + is intra["akk_ext"] + is arena.views(cake_kcp.LOCAL_LAYOUT_FWD).akk_int + ) + assert ( + bintra["aqk_ext"] + is bintra["akk_ext"] + is arena.views(cake_kcp.LOCAL_LAYOUT_BWD).akk_int + ) + assert "Aqk" not in fwd_call.scratch and "Akk" not in fwd_call.scratch + if kp.intra_writes_tables(plan): + scratch = getattr(arena, cake_kcp.SCRATCH_TABLES_ATTR) + assert intra["seg_table"] is scratch["seg_table"] + assert intra["chunk_bos"] is scratch["chunk_bos"] + assert intra["cu_seqlens"] is ctx.cu_seqlens_local + assert kp.local_scratch_table_bytes(plan) == plan.table_words() * 4 > 0 + else: + assert "seg_table" not in intra and "cu_seqlens" not in intra + assert kp.local_scratch_table_bytes(plan) == 0 + assert ( + fwd[1][3]["chunk_bos"] is ctx.tables["chunk_bos"] + ) # the WY program reads the context's tables + assert fwd[2][2] == (max(1, min(plan.NC_pad * H, sm)), 1, 1) + assert bwd[3][2] == (plan.NC_pad // 2, H, 1) + for kind, record, grid, _values in fwd + bwd: + if kind == "backend": + stage = cake_jit.MODULES[record]["stage"] + path = f if record in [r for _, r, _, _ in fwd] else b + if plan.dispatched(stage, path): + assert grid == plan.grid(stage, path) + else: + assert grid == cake_jit.launch_grid( + record, + num_chunks_pad=plan.NC_pad, + num_seqs=plan.N, + num_heads=H, + sm_count=sm, + ) + # the workspace prediction counts the scratch set + scratch_bytes = ( + cake_backend.tables_bytes(ctx.T_r, ctx.n_frag, H, sm_count=sm) + if kp.intra_writes_tables(plan) + else 0 + ) + assert kp.kcp_workspace_bytes(ctx, H=H) == ( + cake_backend.workspace_bytes( + ctx.T_r, ctx.n_frag, H, backward=True, save_mode="light", sm_count=sm + ) + + scratch_bytes + + kp.kcp_arena_layout(ctx, H=H)["arena"] + ) + # the KCP-owned records (the registry's kcp_fwd_o / kcp_dav) find their logical names too + for kind, record, _grid, values in fwd + bwd: + if kind == "kcp": + missing = set(cake_kcp_jit.MODULES[record]["logical"].values()) - set( + values + ) + assert not missing, (record, sorted(missing)) + + +def test_gpu_host_trim_fresh_context_per_call_releases_everything(device) -> None: + """Template-cache ownership (KDA module, one rank): a training loop that builds a NEW context and + workspace per step must not grow device memory -- after each step's objects are dropped the workspace, its + local arena and the context are collected (weak references die) and ``torch.cuda.memory_allocated`` returns to + the level after the first step; the fresh-context results are bitwise the first step's (same kernels, same + argument order).""" + backend = cake_kcp.get_backend() + dt = backend.dtypes + H = 2 + cu = torch.tensor([0, 300, 1024, 1500], dtype=torch.int64) + part = kp.kcp_partition(cu, 1) + T = int(kp.build_kcp_context(part, None, 0, device).T_r) + gen = torch.Generator(device=device).manual_seed(0x3837) + + def randn(*shape): + return torch.randn(*shape, device=device, generator=gen, dtype=torch.float32) + + def unit( + x, + ): # the delta rule's operand contract (harness kcp_inputs): q / k L2-normalized per head + return x * torch.rsqrt((x * x).sum(-1, keepdim=True) + 1e-6) + + q = unit(randn(T, H, cake_kcp.D)).to(dt["q"]) + k = unit(randn(T, H, cake_kcp.D)).to(dt["k"]) + v = (0.5 * randn(T, H, cake_kcp.D)).to(dt["v"]) + g = (-5.0 * torch.sigmoid(randn(T, H, cake_kcp.D))).to(dt["g"]) + beta = torch.sigmoid(randn(T, H)).to(dt["beta"]) + dout = (0.5 * randn(T, H, cake_kcp.D)).to(dt["dout"]) + + def step(): + ctx = kp.build_kcp_context(part, None, 0, device) + ws = cake_kcp.KcpWorkspace(state_dtype=dt["state"], local=backend.local_stages) + out, saved = cake_kcp.chunk_kda_kcp_fwd(q, k, v, g, beta, ctx, workspace=ws) + grads = cake_kcp.chunk_kda_kcp_bwd( + q, k, v, g, beta, ctx, dout, saved=saved, workspace=ws + ) + assert cake_kcp.LAST_CALL["bwd"]["host_trim"] == cake_kcp.HOST_TRIM_DEFAULT + refs = [weakref.ref(ws), weakref.ref(ctx)] + ( + [weakref.ref(ws.local)] if ws.local is not None else [] + ) + return tuple(x.detach().clone() for x in (out, *grads)), refs + + cake_kcp.graph_cache_clear() + first, refs = step() + assert all(torch.isfinite(x.float()).all().item() for x in first) + gc.collect() + torch.cuda.synchronize(device) + assert all(r() is None for r in refs), ( + "the first step's workspace / arena / context are retained" + ) + base = torch.cuda.memory_allocated(device) + for i in range(3): + res, refs = step() + assert all(torch.equal(a, b) for a, b in zip(first, res, strict=True)), i + del res + gc.collect() + torch.cuda.synchronize(device) + assert all(r() is None for r in refs), ( + i, + "a fresh context's workspace / arena / context is retained", + ) + now = torch.cuda.memory_allocated(device) + assert now == base, (i, now - base) diff --git a/tests/experimental/test_cake_kda_train_kcp_build.py b/tests/experimental/test_cake_kda_train_kcp_build.py index 6efdd348d67..f11ccf83059 100644 --- a/tests/experimental/test_cake_kda_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_train_kcp_build.py @@ -37,7 +37,7 @@ ARCHES = ("sm_100a", "sm_103a", "sm_107a") GRID = [["grid", "grid_x"], ["grid", "grid_y"], ["grid", "grid_z"]] -# The frozen kernel argument lists (name, kind) of the four KCP kernel packages. +# The frozen kernel argument lists (name, kind) of the four KCP kernel packages and the two chunk stages of the route. FWDH_PLAN = [ ["buffer", "w_tma"], ["buffer", "kg_tma"], ["buffer", "u"], ["buffer", "gk"], ["parameter", "num_heads"], ["buffer", "seq_chunk_start"], ["parameter", "num_seqs"], ["buffer", "walk_seg"], ["buffer", "h0"], @@ -64,6 +64,15 @@ ["parameter", "ncols"], ["parameter", "transpose_m"], ["parameter", "write_states"], ["parameter", "cols_per_cta"], *GRID, ] # fmt: skip +FWD_O_PLAN = [ + ["buffer", "qg_tma"], ["buffer", "h_tma"], ["buffer", "aqk_tma"], ["buffer", "vn_tma"], ["buffer", "o_out"], + ["buffer", "chunk_bos"], ["buffer", "chunk_len"], ["buffer", "seq_chunk_start"], ["parameter", "num_seqs"], + ["parameter", "num_heads"], ["parameter", "scale"], *GRID, +] # fmt: skip +DAV_PLAN = [ + ["buffer", "do_tma"], ["buffer", "vnew_tma"], ["buffer", "aqk_tma"], ["buffer", "dAqk"], ["buffer", "dv1"], + ["parameter", "num_heads"], ["parameter", "scale"], *GRID, +] # fmt: skip PLANS = { "kcp_fwdh_offset": FWDH_PLAN, "kcp_fwdh_tape": FWDH_PLAN, @@ -72,6 +81,8 @@ "kcp_transfer_walk": TRANSFER_PLAN, "kcp_affine_merge_c8": MERGE_PLAN, "kcp_affine_merge_c16": MERGE_PLAN, + "kcp_fwd_o": FWD_O_PLAN, + "kcp_dav": DAV_PLAN, "kcp_transfer_walk_three_mma": TRANSFER_PLAN, "kcp_transfer_walk_plain_tf32": TRANSFER_PLAN, "kcp_transfer_walk_bf16x3": TRANSFER_PLAN, @@ -108,6 +119,8 @@ f"tart.kda_chunk_kcp_affine_merge.{_W}.kda_chunk_kcp_affine_merge_fp32_v1", {"cols_per_cta": 16}, ), + "kcp_fwd_o": (f"tart.kda_chunk_fwd_o.{_W}.kda_chunk_train_fwd_o_v1", {}), + "kcp_dav": (f"tart.kda_chunk_bwd_dav.{_W}.kda_chunk_train_bwd_dav_v1", {}), "kcp_transfer_walk_three_mma": ( f"tart.kda_chunk_kcp_transfer_walk.{_W}.kda_chunk_kcp_transfer_walk_v1", {"precision_mode": 1}, @@ -147,6 +160,8 @@ def test_the_tables_are_disjoint_and_complete() -> None: "kcp_dhu_tape", "kcp_transfer_walk", "kcp_affine_merge", + "kcp_fwd_o", + "kcp_dav", ) assert kcp_jit.PRECISION_MODES == (0, 1, 2) assert kcp_jit.TRANSFER_WALKS == ("tf32_hilo", "bf16x3", "bf16x3_hw64") @@ -227,6 +242,16 @@ def test_selection_and_grids() -> None: ) with pytest.raises(ValueError): kcp_jit.launch_grid("kcp_fwdh_offset", n_walk=0, num_heads=4) + # the chunk stages: persistent CTAs bounded by the SMs, one CTA per chunk pair and head + assert kcp_jit.launch_grid( + "kcp_fwd_o", num_chunks_pad=10, num_heads=3, sm_count=152 + ) == (30, 1, 1) + assert kcp_jit.launch_grid( + "kcp_fwd_o", num_chunks_pad=400, num_heads=4, sm_count=152 + ) == (152, 1, 1) + assert kcp_jit.launch_grid("kcp_dav", num_chunks_pad=10, num_heads=3) == (5, 3, 1) + with pytest.raises(ValueError): + kcp_jit.launch_grid("kcp_dav", num_chunks_pad=9, num_heads=3) def test_test_record_selection_by_knobs() -> None: @@ -296,7 +321,13 @@ def test_test_manifest_is_the_test_only_programs_under_the_opt_in( == shipped["package"] == "flashinfer.experimental.cake_kda_train" ) - assert document["generator"] == shipped["generator"] + # the shipped manifest is re-rendered by every kernel update of the package, the test-only manifest only when its + # programs change: the two generator records share their form, not their stamp + assert ( + set(document["generator"]) + == set(shipped["generator"]) + == {"cake_revision", "cuda", "rendered_at"} + ) names = {entry["name"] for entry in document["libraries"]} assert names == { f"{record['library']}_{arch}" diff --git a/tests/experimental/test_cake_kda_train_kcp_dist.py b/tests/experimental/test_cake_kda_train_kcp_dist.py index 4da79a00c98..39d2c45f6ce 100644 --- a/tests/experimental/test_cake_kda_train_kcp_dist.py +++ b/tests/experimental/test_cake_kda_train_kcp_dist.py @@ -223,12 +223,14 @@ def fn(*a): "ckpt_eq_direct": all(torch.equal(res_auto[n], res_ckpt[n]) for n in TENSORS), "save_transfer_eq": all(torch.equal(res[n], res_recompute[n]) for n in TENSORS), "ring_eq_a2a": ring_ok, - # a backend with local stages must have run the SPEC-3 assertion (unless disabled) and passed it; a rank - # without tokens launches no local stage, so it has nothing to check + # a backend with local stages must have run the SPEC-3 probe (unless disabled); a rank without tokens + # launches no local stage, so it has nothing to probe. The probe's row comparison (the forward's intra + # prep against the backward's prep program) is reported in spec3_by_rank, not gated: the requirement of + # record is saved_eq_replay (both backward modes run the same local programs) "spec3": (not backend.local_stages) or (spec3 is None and os.environ.get(cake_kcp.SPEC3_ENV, "1") == "0") or (spec3 is None and q.shape[0] == 0) - or bool(spec3 and spec3.get("equal")), + or bool(spec3 and spec3.get("checked")), } gathered = {} for n in TENSORS: @@ -421,7 +423,15 @@ def test_public_api_cannot_select_a_transfer_walk(monkeypatch) -> None: def test_module_imports_only_torch_and_its_package() -> None: """The design note §11.3: the host module imports torch, torch.distributed, tvm_ffi and its own package (plus ``ctypes`` for the CUDA runtime's device cache-config preference around the tape launches, - ``cake_kcp.walk_device_cache_config``).""" + ``cake_kcp.walk_device_cache_config``; the same ``ctypes`` import also reaches the capture-graph entry points of the + replay cap of record + ``graph_cap="node_attr"`` -- ``cudaStreamGetCaptureInfo`` (``_v2`` / ``_v3``), ``cudaGraphNodeGetType``, + ``cudaGraphKernelNodeSetAttribute``, ``cudaGraphKernelNodeGetAttribute``, ``cudaRuntimeGetVersion``, + ``cudaGetLastError``; no new module; the trimmed host path: + ``functools.cached_property`` on the host table set (``walk_all_p0`` / ``walk_all_rest``, host ints only) -- stdlib, + no new module either; ``weakref`` for the trimmed + path's begin_call template cache, owned by the local arena and weakly keyed by the table set -- stdlib, no new + module).""" import ast src = Path(cake_kcp.__file__).read_text() @@ -440,7 +450,9 @@ def test_module_imports_only_torch_and_its_package() -> None: "tvm_ffi", "ctypes", "warnings", - } # ctypes: the cache-config bracket of the walk-launch geometry; warnings: its one RuntimeWarning without a loadable CUDA runtime; no other use + "functools", + "weakref", + } # weakref: the arena-owned template cache of the trimmed host path; functools: cached_property of the trimmed host path; ctypes: the cache-config bracket of the walk-launch geometry and the capture-graph entry points of the replay cap; warnings: the one RuntimeWarning without a loadable CUDA runtime or for a degraded node cap; no other use for node in ast.walk(tree): if isinstance(node, ast.Import): for a in node.names: diff --git a/tests/test_cake_causal_conv1d_train_kcp_knobs.py b/tests/test_cake_causal_conv1d_train_kcp_knobs.py index cbf6c15683b..692a9f6bf2a 100644 --- a/tests/test_cake_causal_conv1d_train_kcp_knobs.py +++ b/tests/test_cake_causal_conv1d_train_kcp_knobs.py @@ -32,10 +32,12 @@ def test_graph_knob_is_gated_and_validated(monkeypatch) -> None: monkeypatch.delenv(conv_kcp.TEST_KNOBS_ENV, raising=False) assert conv_kcp.GRAPH_MODES == ("off", "replay") - assert conv_kcp._set_test_knobs(None) == {"graph": "off"} + assert conv_kcp._set_test_knobs(None) == {"graph": "off", "host_trim": 1} with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): conv_kcp._set_test_knobs({"graph": "replay"}) - assert conv_kcp._TEST_KNOBS == {"graph": "off"} + with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): + conv_kcp._set_test_knobs({"host_trim": 1}) + assert conv_kcp._TEST_KNOBS == {"graph": "off", "host_trim": 1} monkeypatch.setenv(conv_kcp.TEST_KNOBS_ENV, "1") try: with pytest.raises(ValueError, match="unknown"): @@ -44,10 +46,58 @@ def test_graph_knob_is_gated_and_validated(monkeypatch) -> None: conv_kcp._set_test_knobs({"graph": "capture"}) with pytest.raises(TypeError): conv_kcp._set_test_knobs(["replay"]) - assert conv_kcp._set_test_knobs({"graph": "replay"}) == {"graph": "replay"} - assert conv_kcp._set_test_knobs({"graph": "off"}) == {"graph": "off"} + assert conv_kcp._set_test_knobs({"graph": "replay"}) == { + "graph": "replay", + "host_trim": 1, + } + assert conv_kcp._set_test_knobs({"graph": "off"}) == { + "graph": "off", + "host_trim": 1, + } + finally: + conv_kcp._set_test_knobs(None) + + +def test_host_trim_knob_is_gated_and_validated(monkeypatch) -> None: + """``host_trim`` (the trimmed host path): refused without ``CAKE_KCP_TEST_KNOBS=1``, + validated against ``HOST_TRIMS`` (1 shipped = the trimmed host path of the eager halo call, + 0 = the previous host path), independent of ``graph``; a dict naming one knob keeps the other at its shipped default.""" + monkeypatch.delenv(conv_kcp.TEST_KNOBS_ENV, raising=False) + assert conv_kcp.HOST_TRIMS == (0, 1) and conv_kcp.HOST_TRIM_DEFAULT == 1 + with pytest.raises(ValueError, match="CAKE_KCP_TEST_KNOBS"): + conv_kcp._set_test_knobs({"host_trim": 0}) + assert conv_kcp._TEST_KNOBS["host_trim"] == 1 + monkeypatch.setenv(conv_kcp.TEST_KNOBS_ENV, "1") + try: + assert conv_kcp._set_test_knobs({"host_trim": 1}) == { + "graph": "off", + "host_trim": 1, + } + assert conv_kcp._graph_mode(None) == "off" + assert conv_kcp._set_test_knobs({"host_trim": 1, "graph": "replay"}) == { + "graph": "replay", + "host_trim": 1, + } + assert conv_kcp._set_test_knobs({"graph": "replay"}) == { + "graph": "replay", + "host_trim": 1, + } + assert conv_kcp._set_test_knobs({}) == {"graph": "off", "host_trim": 1} + for bad in ( + {"host_trim": 2}, + {"host_trim": "1"}, + {"host_trim": True}, + {"host_tri": 1}, + ): + with pytest.raises(ValueError): + conv_kcp._set_test_knobs(bad) + assert conv_kcp._set_test_knobs({"host_trim": 0}) == { + "graph": "off", + "host_trim": 0, + } finally: conv_kcp._set_test_knobs(None) + assert conv_kcp._TEST_KNOBS == {"graph": "off", "host_trim": 1} def test_graph_cache_key_and_eager_first_call(monkeypatch) -> None: @@ -159,3 +209,51 @@ def test_graph_public_keyword_mode_and_key(monkeypatch) -> None: conv_kcp._set_test_knobs(None) conv_kcp._GRAPHS.clear() conv_kcp.LAST_CALL.clear() + + +def test_host_trim_cache_is_per_context_and_bounded(monkeypatch) -> None: + """The trimmed host path keeps one entry set per context object (by identity, the context pinned), bounded + by ``HOST_CACHE_MAX`` (the oldest context dropped) and cleared by ``host_cache_clear()``.""" + conv_kcp.host_cache_clear() + a, b = types.SimpleNamespace(), types.SimpleNamespace() + key = ("fwd", 8, 16, "cpu", (0, 8)) + ea = conv_kcp._host_cache_entry(a, key) + assert ea == {} and conv_kcp._host_cache_entry(a, key) is ea + ea["x"] = 1 + eb = conv_kcp._host_cache_entry(b, key) + assert eb == {} and eb is not ea + assert conv_kcp._host_cache_entry(a, ("bwd", 8, 16, "cpu", (0, 8))) is not ea + assert conv_kcp._host_cache_entry(a, key) is ea + monkeypatch.setattr(conv_kcp, "HOST_CACHE_MAX", 2) + conv_kcp._host_cache_entry( + types.SimpleNamespace(), key + ) # a third context evicts the oldest (a) + assert id(a) not in conv_kcp._HOST_CACHE and id(b) in conv_kcp._HOST_CACHE + assert conv_kcp.host_cache_clear() == 2 and conv_kcp._HOST_CACHE == {} + + +def test_host_trim_ffi_device_and_cpu_refusal(monkeypatch) -> None: + """Trimmed host path: the FFI device object is built once per index; the trimmed forward / backward are selected by the knob + and refuse CPU tensors like today's path (the operand checks first, then the architecture check).""" + dev = conv_kcp._ffi_device(0) + assert ( + conv_kcp._ffi_device(0) is dev + and int(dev.index) == 0 + and int(dev.dlpack_device_type()) == 2 + ) + x = torch.zeros(4, 16, dtype=torch.bfloat16) + w = torch.zeros(16, 4, dtype=torch.float32) + hist = torch.zeros(1, 16, dtype=torch.bfloat16) + ctx = types.SimpleNamespace(cu_seqlens_local_cpu=(0, 4)) + monkeypatch.setenv(conv_kcp.TEST_KNOBS_ENV, "1") + try: + conv_kcp._set_test_knobs({"host_trim": 1}) + with pytest.raises(ValueError, match="CUDA"): + conv_kcp._kcp_forward(x, w, ctx) + with pytest.raises(ValueError, match="CUDA"): + conv_kcp._kcp_backward(x, w, ctx, x, hist) + finally: + conv_kcp._set_test_knobs(None) + with pytest.raises(ValueError, match="CUDA"): + conv_kcp._kcp_forward(x, w, ctx) + conv_kcp.host_cache_clear()