diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 26b34ab3ef3..6960dc59073 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -335,7 +335,7 @@ jobs: strategy: fail-fast: false matrix: - suite: [fast, adapters, deepseek, lanes, deferred] + suite: [fast, adapters, native, deepseek, lanes, deferred] steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 with: @@ -344,11 +344,29 @@ jobs: with: version: 0.17.0 use-cache: false + - name: Metal toolchain for native engine host checks + if: matrix.suite == 'native' + run: xcrun -sdk macosx metal -v || { xcodebuild -downloadComponent MetalToolchain && xcrun -sdk macosx metal -v; } - name: Host checks and aggregate wall time shell: bash env: SUITE: ${{ matrix.suite }} run: | printf 'Runner class: %s %s\n' "$RUNNER_OS" "$RUNNER_ARCH" + if [ "$SUITE" = native ]; then + printf 'Native source: %s\n' "$GITHUB_SHA" + status=0 + for phase in initial warm; do + TIMEFORMAT="test-native $phase wall seconds: %R" + if { time zig build test-host-ci-macos-native --test-timeout 30m --summary all; }; then + printf 'test-native %s exit: 0\n' "$phase" + else + code=$? + printf 'test-native %s exit: %s\n' "$phase" "$code" + status=$code + fi + done + exit "$status" + fi TIMEFORMAT='Host checks wall seconds: %R' time zig build "test-host-ci-macos-$SUITE" --test-timeout 30m diff --git a/CHANGELOG.md b/CHANGELOG.md index 716bd8c6a93..73e51edec80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,42 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- OCR: the native HIP engine reads page images through `POST /v1/ocr` in Mistral's OCR contract (its Python and TypeScript SDKs), with dots.mocr's reader (a box for each layout block). Mistral crops share one image budget per request (256 images, 4 times the page pixels, 16 MiB plus 2 MiB a page). An option that changes content and that TensorFold cannot apply, such as annotations, is refused, and so are PDF documents and `pages`. A page cut at the context limit or out of its model's format fails the request with HTTP 422; `docs/api.md` lists the contract's fields, statuses and limits (#2570, #2440). Thanks to @laudney. + +- Video on Linux: `libtfvideo.so` exports only its `tf_video_*` functions, so FFmpeg's x86_64 assembly links into it, and a CI check refuses a build that exports more (#2861). Thanks to @swbiggart. + +- Docs: the Intel Arc (XPU) page notes that its B70 prefill row predates 1.1.0's more precise prompt DeltaNet, which costs about 5% of a 32,768-token prompt with tokens unchanged (#2906, measured in #2859). Thanks to @dinadur. + +- Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. + +- Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. + +- Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. + +- Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. + +- GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. + +- HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. + +- Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. + +- CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). + +- Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). + +- Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. + +- GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. + +- Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). + +- GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). + +- CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. + +- DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. + - Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). @@ -79,9 +115,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. - - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. +- Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. @@ -127,6 +163,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - HIP on gfx12 under `--policy matrix=on`: Qwen3 dense checkpoints load each layer's gate and up weights interleaved as one weight, and one WMMA product writes silu(gate) x up in its epilogue instead of two products and `tf_silu_mul`, with the same output bytes: Qwen3-Reranker-8B ranks 16 documents in 0.657 to 0.661 s instead of 0.687 to 0.697 s and Qwen3-Embedding-8B embeds 100 short inputs in 0.901 to 0.911 s instead of 0.940 to 0.948 s on a Radeon AI PRO R9700. The kernel contract checks both fused blocks byte for byte against the unfused launches (part of #2245). Thanks to @laudney. - Nemotron on Macs: the masked prompt shader calls the shared `tf_paged_row` page helper again, so its prompt sources compile; #2322 had used the name #2283 retired the same day (#2548). + ## 1.2.7 - GLM-5.3-Flash on CUDA: with `TF_GLM_LAYERS=N`, serve's start-up memory check counts the weights by tensor from the shard headers, only those the load reads (the first N layers, the MTP layer when drafts are on, the embedding, norm and head), where it counted every checkpoint file. A 20-layer thin load of the 163.6 GiB checkpoint now counts 66.6 GiB without drafts. Unset, the count is the files' as before (#2788, fixes #2745). Thanks to @wojo. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index c894ea02bd1..7a72107791f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -292,8 +292,7 @@ below); `tools/zig/bolt_ggml.py` reads them from that file. The IQ2_XXS and Q8_0 Bolt's GGML checkpoints (`zig/kernels/cuda/ggml.cu`, `zig/src/cuda/ggml.zig`) are TensorFold's own. TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card states Apache-2.0. -The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use -without derivatives; only the Python engine reads that drafter. Each checkpoint keeps its own license. +The optional `incoai/GLM-5.3-Flash-DFlash2` model card states Apache-2.0 too. Each checkpoint keeps its own license. ## Structured output: libtfgrammar diff --git a/docs/api.md b/docs/api.md index 98f10a003f4..2d9ac207bc5 100644 --- a/docs/api.md +++ b/docs/api.md @@ -45,7 +45,7 @@ different scopes, or scopes a key the flags or environment already give, is refu | Scope | Routes | | --- | --- | -| `inference` | POST `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/messages`, `/v1/messages/count_tokens`, `/v1/tokenize`, `/v1/detokenize`, `/v1/audio/speech`, `/v1/images/generations`, `/v1/embeddings`, `/v1/decisions`, `/v1/systemone`, `/v1/rerank` (also by PUT), and the same paths without `/v1` where the server has them; GET `/v1/models`, `/models`, `/v1/voices`; a WebSocket upgrade on GET `/v1/responses` or `/responses` | +| `inference` | POST `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/messages`, `/v1/messages/count_tokens`, `/v1/tokenize`, `/v1/detokenize`, `/v1/audio/speech`, `/v1/images/generations`, `/v1/embeddings`, `/v1/decisions`, `/v1/systemone`, `/v1/rerank` (also by PUT), `/v1/ocr`, and the same paths without `/v1` where the server has them; GET `/v1/models`, `/models`, `/v1/voices`; a WebSocket upgrade on GET `/v1/responses` or `/responses` | | `read` | GET `/metrics`, `/v1/metrics`, `/v1/load`, `/memory`, `/v1/memory`, `/stats`, `/dashboard`, `/v1/benchmark/manifest`, `/v1/models`, `/models`, `/v1/voices` | A scoped key reaches nothing else: not the stored Responses (GET or DELETE `/v1/responses/ID`), Recall, Sliding Weights, @@ -193,6 +193,95 @@ context window before the first is scored, so a refused request costs no GPU wor `max_length` above it) returns HTTP 400 with `code` `context_length_exceeded`. An image or video part returns HTTP 400. Errors use the server's envelope, as on every route: `{"error": {"message", "type"}}`. +## OCR + +The native HIP engine reads page images through Mistral's OCR contract when it serves a model with an OCR reader +and `--vision`. + +| Route | Contract | Request | Reply | Written for | +| --- | --- | --- | --- | --- | +| `POST /v1/ocr` | Mistral OCR | JSON `{"model", "document", ...}` | `{"pages", "model", "document_annotation", "usage_info"}` | the `mistralai` Python and `@mistralai/mistralai` TypeScript SDKs, with `server_url` set to the server root | + +One family has a reader. A reader gives Markdown blocks in pixels of the page the client sent (after EXIF +orientation), never in the model's resized frame. + +| Model | What the reader returns | Geometry | +| --- | --- | --- | +| dots.mocr | one block a layout region (title, text, table, formula, picture, header, footer, ...), tables as pipe tables or, when cells span, HTML | a box a block | + +Any other model answers the route with HTTP 400 "this model has no OCR reader; use POST /v1/chat/completions", +and a server started without `--vision` answers "OCR needs image input; start the server with --vision". +Embedding, reranker and readout engines refuse the OCR route as they refuse chat. The route takes the `inference` +key scope. + +### Reader Markdown + +The reader writes block text in one Markdown flavour, the one Mistral's page Markdown uses: + +- Headings with `#` markers: a title is `#`, a section header `##` unless the model gives a level. +- GFM pipe tables, or one HTML `` when a table has row or column spans a pipe table cannot hold; a table + block holds exactly one table. +- Display math as `$$…$$` and inline math as `$…$`, the LaTeX kept as the model wrote it. +- Lists with `-` or `1.` markers, and code in fenced blocks. +- No images: a picture block has empty text, and its caption, if any, is its own block. The contract writes its + own image references. + +### Pages and budgets + +- Images are PNG, JPEG or WebP, one page each, up to the 32 MiB body limit and the reader's decoded-pixel limit: + 16 megapixels and 8192 pixels a side for dots.mocr; a larger image is bad input. An image whose header and + decoded pixels disagree is refused as bad input. +- PDF documents and `pages` are not read: a PDF document is HTTP 400 "PDF documents are not read on POST /v1/ocr; + send each page as a PNG, JPEG or WebP image", and any non-null `pages` is HTTP 400 "pages selects pages of a PDF + document; POST /v1/ocr reads one page image". +- `https://` URLs (`document_url` and `image_url`) are fetched only with `--vision-urls`: images, at most 32 MiB. + `http://` URLs are refused. +- Image budget, for Mistral crops (`include_image_base64`): at most 256 encoded images, 4 times the request's page + pixels, and 16 MiB plus 2 MiB a page of encoded bytes. Crossing one fails the whole request as an unreadable + page, "too many image regions". The encode needs libjpeg-turbo 3.1.4.1 (`TENSORFOLD_LIBTURBOJPEG` names its + file); without it `include_image_base64` is HTTP 400. + +The page is read as one user turn (the page image, then the reader's prompt), greedy, with thinking off and +`max_tokens` equal to all the context the prompt leaves. It succeeds only when the reply stops on its end token +and the reader can parse it. + +### Failures + +| Failure | Mistral | +| --- | --- | +| Bad input: an undecodable or unsupported file, a page or option the server cannot serve | 400 | +| Unreadable page: the reply was cut at the context limit, repeated without end, broke the model's format or gave a box that is not finite; the image budget was crossed | 422 | +| The engine stopped the page (it closed or drained) | 503, `Retry-After: 1` | +| Engine fault or out of memory | 500 | +| Envelope | `{"detail": [{"loc", "msg", "type"}]}`, `type` `bad_input`, `unreadable_page`, `capacity` or `internal` | + +An unreadable page's message names its index, and the model's raw text goes to the server log, never to the +client. 422 means the same request fails the same way again under greedy decoding; a 503 carries `Retry-After`. +A client that disconnects stops the request before its page is read. + +### Mistral: `POST /v1/ocr` + +| Field | Meaning | +| --- | --- | +| `model` | Any string or null. The reply names the requested id when it is the served name or an `--alias`, else the served name, so `--alias mistral-ocr-latest` lets copied examples echo it | +| `document` | `{"type": "document_url", "document_url"}` or `{"type": "image_url", "image_url"}` (a string or `{"url"}`): a `data:` URL, or an `https://` URL as above. The bytes must be a PNG, JPEG or WebP image; a PDF is HTTP 400 (above). `{"type": "file"}` is HTTP 400 | +| `pages` | Any non-null value is HTTP 400 (above): an image has one page | +| `table_format` | Null: tables stay in the page Markdown as pipe tables (a spanned HTML table becomes a pipe table with the spanned text in its first cell). `"markdown"` or `"html"`: each table becomes a `tables[]` entry (`tbl-N.md` or `tbl-N.html`) and the Markdown links it; under `"markdown"` a table with spanned cells stays HTML, the only form that keeps it | +| `extract_header`, `extract_footer` | Move page headers or footers out of `markdown` into `header` and `footer` | +| `include_blocks` | Default `true`: one block for each region with a box. `false` leaves `blocks` out | +| `image_limit`, `image_min_size` | Choose which pictures to keep, in reading order (`image_min_size` against the box's width and height, then the running count against `image_limit`). A dropped picture has no image entry, Markdown link or block | +| `include_image_base64` | Add `image_base64`, a `data:image/jpeg;base64,...` crop of the picture's box, to each kept image | +| `confidence_scores_granularity` | Accepted; no scores are returned | +| `bbox_annotation_format`, `document_annotation_format`, `document_annotation_prompt` | Any non-null value is HTTP 400 "annotations are not supported" | + +`null` means absent for every optional field, and unknown keys are ignored. Each page carries `index`, `markdown` +(blocks joined by a blank line, pictures as `![img-N.jpeg](img-N.jpeg)`), `images`, `tables`, `header`, `footer`, +`dimensions` (`dpi`, `height`, `width`) and `blocks`; `hyperlinks` is left out. A block has `type` (`text`, +`title`, `list`, `table`, `equation`, `code`, `caption`, `header`, `footer` or `image`), an integer box and +`content`, its own Markdown with its tables in the `table_format` form; image blocks add `image_id` and table +blocks `table_id`. `usage_info` gives `pages_processed` and `doc_size_bytes`, the decoded document's size. An +image's `dimensions.dpi` is 72. + ## System One on a Clef engine `POST /v1/systemone` on a Clef checkpoint (Cloudflare's clef-flash: BF16 Qwen3.5 weights with diff --git a/docs/environment.md b/docs/environment.md index 0af40a6f12c..b1818515681 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -43,6 +43,7 @@ Diagnostics quote and escape the supplied text so it cannot create another termi | `TF_GLM_SHARED` | Exactly `0` or `1`; default on (GLM windows share a forward); read on rank 0 | | `TF_GLM_DFLASH2_LOG` | Whole unsigned 64-bit count of DFlash2 rounds whose drafts are logged; default 0 | | `TF_GLM_DRAFT_CONFIDENCE` | Finite number 0 through 1 where a DFlash2 chain's chance product stops it; default 0.7 | +| `TF_GLM_CONTEXT_CEILING` | GLM-5.3-Flash on CUDA: whole token count 1 through 2147483616, the window `--context` may ask for and `--context 0` serves, in place of config.json's 1048576; read at GLM's open. Past that length the model's quality is the operator's to check. Unset keeps config.json's window | | `TF_PROMPT_ROWS` | Native serving: a multiple of 16 from 128 through 2048; unset keeps the chip/family default | | `TF_ANE_COLUMNS` | Whole u32 count; zero disables the experimental share; default zero | | `TF_FLASHNEXT_SCALE_INDEX` | Flash Next on Metal before tensor units: exactly `0` or `1`; `1` has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs, the same values with fewer weight bytes; default 0 | diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 6b4d03873ff..520a89bc5d5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -401,6 +401,10 @@ CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 down, each bit-equal to the decode instance it replaces; the scaled-code prompt tile is the planted control). +CUDA also checks the bf16 experts (`experts_bf16.cu`, a checkpoint's own bf16 weights: Nemotron 3.5 Transcribe's routed +and shared experts): relu^2 up, fp32 down and bf16 down, in the decode and prompt forms, against float64 (two controls: +the fp32 down launched one 32-input group short, and a down summing in bf16). nvcc 13.2's prompt forms and 13.3's +decode forms on sm_120 and sm_121 take other registers than 13.0's and have their own records. CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 8d5efda9499..e8ed3350863 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -162,7 +162,7 @@ concurrent one its run alone. | DFlash2 | Prose, 16 streams | 526.0 | 583.6 | One-stream runs are `tools/bench_openai.py` at 512 tokens and stream runs `tools/bench_concurrent.py` at 256. DFlash2's -weights are CC BY-NC-ND 4.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts +weights are Apache-2.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts are unchanged at about 4,430 tok/s from 8K to 131K tokens, and a repeated code reply reaches about 850 tok/s with windows of up to 32 rows. Other GPUs and rank counts are not covered here. @@ -250,6 +250,15 @@ Long prompts and long runs of replies heat a GB10. The runbook's [heat bands](.. 1,035,295-token prompt at TP=3), [CPU pinning](../../RUNBOOK.md#gb10-cpu-pinning-for-long-cuda-runs) and [GPU clock cap](../../RUNBOOK.md#gb10-gpu-clock-cap-for-long-cuda-runs) each carry a measured GLM-5.3-Flash run. +## Windows past 1,048,576 tokens + +GLM-5.3-Flash's attention keeps no rotary positions (`qk_rope_head_dim` is 0, and KDA keeps none), so only its trained +length, config.json's 1,048,576 tokens, bounds the window. `TF_GLM_CONTEXT_CEILING` sets the window `--context` may ask +for and `--context 0` serves. GLM reads it at its open, and the startup log says when the window passes the trained +length, past which the quality is the operator's to check. On three GB10s with `Bizuayeu/GLM-5.3-Flash-XSFP4`, +`TF_GLM_CONTEXT_CEILING=2097152`, `--context 2097152 --parallel 1` and no prompt cache, rank 0 took 4,097 KV pages, and +a 1,996,430-token prompt answered all three of its passphrases. + ## Shared KDA operator `cuda.kda` owns the fused and three-kernel chains and single-layer and batched saved-row replay. diff --git a/docs/recipes/glm-5.3-flash.md b/docs/recipes/glm-5.3-flash.md index 538adf74bdc..f2e7a68c365 100644 --- a/docs/recipes/glm-5.3-flash.md +++ b/docs/recipes/glm-5.3-flash.md @@ -42,8 +42,8 @@ tensorfold serve TensorFold/GLM-5.3-Flash-MLX-4bit-MTP --tp 2 --rank 0 --master ``` Rank 1 starts first and rank 0 serves HTTP. Use the same context and drafting settings on both ranks. -The optional `incoai/GLM-5.3-Flash-DFlash2` model has CC BY-NC-ND 4.0 terms; pull it on both ranks only -when those terms fit the intended use. The CLI uses it automatically once it has been pulled. +The optional `incoai/GLM-5.3-Flash-DFlash2` model is Apache-2.0; pull it on both ranks to use it. +The CLI uses it automatically once it has been pulled. Without it, the engine uses MTP drafts; `--drafter none` explicitly selects MTP-only drafting. Give both ranks the same drafter setting. `--no-drafts` disables all drafting for the serial reference. A checkpoint with neither an MTP head nor a supplied DFlash2 model is refused unless drafts are disabled. diff --git a/docs/xpu/README.md b/docs/xpu/README.md index bd649f64412..edddef11648 100644 --- a/docs/xpu/README.md +++ b/docs/xpu/README.md @@ -90,6 +90,11 @@ Options through environment variables: | 2 / 4 / 8 simultaneous streams, chat, greedy (aggregate) | 151 / 239 / 300 | 148 / 215 / 208 | 162 / 230 / 224 | 162 / 249 / 227 | | `bench_openai` decode, 1 / 2 / 4 / 8 requests (aggregate) | 173 / 243 / 343 / 350 | 165 / 300 / 337 / 439 | 165 / 336 / 367 / 456 | 179 / 325 / 364 / 463 | +The prefill row predates 1.1.0, which feeds the prompt DeltaNet's fp32 intermediates to the matrix units at close to fp32 +precision (`489baa2673`, `17facba3f8`). On a B70 with GPTQ g128 that costs about 5% of a 32,768-token prompt in the +raw runner (`xpu-qwen27 --bench=32768`: 1,647 to 1,571 tok/s, tokens unchanged; #2859). The table's sizes have not been +re-measured since. + Before the per-stream DFlash2 drafts (copies only in the shared rounds), 2 / 4 / 8 code streams ran at 64 / 120 / 182 tok/s in GPTQ. Validation: `zig build test`, oracle, drafted == serial, multi-stream == serial (chains and trees), and `bench_concurrent --alone --serial` on 1, 2, 4 and 8 streams, greedy and sampled, for the four diff --git a/tools/zig/check_tfvideo_exports.py b/tools/zig/check_tfvideo_exports.py new file mode 100644 index 00000000000..c1a9d57db62 --- /dev/null +++ b/tools/zig/check_tfvideo_exports.py @@ -0,0 +1,70 @@ +"""An ELF libtfvideo.so exports exactly the four tf_video_* functions, keeps a planted table local, has no TEXTREL.""" +import struct +import sys +from pathlib import Path + +WANT = {'tf_video_open', 'tf_video_frames', 'tf_video_close', 'tf_video_version'} +SHT_SYMTAB, SHT_DYNAMIC, SHT_DYNSYM = 2, 6, 11 +DT_TEXTREL, DT_FLAGS, DF_TEXTREL = 22, 30, 0x4 + + +def sections(data: bytes) -> list[tuple]: + """(type, offset, size, link, entsize) for each section of a little-endian ELF64 file.""" + if data[:4] != b'\x7fELF' or data[4] != 2 or data[5] != 1: + raise SystemExit('not a little-endian ELF64 file') + shoff, = struct.unpack_from(' dict[str, tuple[int, int]]: + """Defined symbols of the given table: name -> (binding, section index).""" + out = {} + for t, off, size, link, ent in secs: + if t != kind: + continue + strtab = secs[link][1] + for at in range(off + ent, off + size, ent): + name, info, _, shndx = struct.unpack_from('> 4, shndx) + return out + + +def textrel(data: bytes, secs: list[tuple]) -> bool: + """True when .dynamic has DT_TEXTREL or DF_TEXTREL in DT_FLAGS.""" + for t, off, size, _, _ in secs: + if t == SHT_DYNAMIC: + for at in range(off, off + size, 16): + tag, val = struct.unpack_from(' int: + failed = 0 + for arg in sys.argv[1:]: + data = Path(arg).read_bytes() + secs = sections(data) + exported = {n for n, (bind, _) in symbols(data, secs, SHT_DYNSYM).items() if bind in (1, 2)} + planted = symbols(data, secs, SHT_SYMTAB).get('tfvideo_planted_table') + problems = [f'exports {sorted(exported ^ WANT)} beyond or short of the four'] if exported != WANT else [] + if planted is None or planted[0] != 0: + problems.append(f'the planted table is not a local symbol: {planted}') + if textrel(data, secs): + problems.append('text relocations') + machine = {62: 'x86_64', 183: 'aarch64'}.get(struct.unpack_from(' @@ -31,6 +32,7 @@ template __device__ T __shfl_xor_sync(unsigned,T,int,int=32); template __device__ T __shfl_up_sync(unsigned,T,unsigned,int=32); template __device__ T __shfl_down_sync(unsigned,T,unsigned,int=32); __device__ void __syncthreads(); __device__ int __syncthreads_count(int); __device__ void __syncwarp(unsigned=~0u); +__device__ int __syncthreads_or(int); __device__ void __threadfence(); __device__ void __threadfence_system(); __device__ void __trap(); __device__ void __nanosleep(unsigned); __device__ unsigned long long clock64(); __device__ unsigned __ballot_sync(unsigned,int); __device__ int __any_sync(unsigned,int); @@ -49,6 +51,7 @@ __device__ int __float_as_int(float); __device__ unsigned __float_as_uint(float); __device__ float __int_as_float(int); __device__ float __uint_as_float(unsigned); __device__ double __longlong_as_double(long long); __device__ long long __double2ll_rz(double); __device__ float __double2float_rn(double); +__device__ float __double2float_rz(double); __device__ int __float2int_rn(float); __device__ float __fadd_rn(float,float); __device__ float __fsub_rn(float,float); __device__ float __fmul_rn(float,float); __device__ float __fdiv_rn(float,float); __device__ float __fsqrt_rn(float); __device__ float __fmaf_rn(float,float,float); __device__ float __expf(float); diff --git a/tools/zig/cuda_host/math.h b/tools/zig/cuda_host/math.h index 4aade403d68..65cf5c4b1c8 100644 --- a/tools/zig/cuda_host/math.h +++ b/tools/zig/cuda_host/math.h @@ -19,5 +19,6 @@ template __host__ __device__ bool isfinite(T); #define M_2_SQRTPI 1.12837916709551257390 #define M_SQRT1_2 0.70710678118654752440 __host__ __device__ float erff(float); __host__ __device__ float rintf(float); +__host__ __device__ float ldexpf(float,int); __host__ __device__ double pow(double,double); __host__ __device__ double sin(double); __host__ __device__ double cos(double); template __host__ __device__ bool signbit(T); diff --git a/tools/zig/dots_mocr_reference.py b/tools/zig/dots_mocr_reference.py index 4c0038cfdb9..a29921aba2e 100644 --- a/tools/zig/dots_mocr_reference.py +++ b/tools/zig/dots_mocr_reference.py @@ -1,4 +1,4 @@ -"""dots.mocr references: the model card's Hugging Face example oracle, processor fixtures, native trace compare.""" +"""dots.mocr references: the Hugging Face oracle, processor and OCR fixtures, our test pages, native trace compare.""" import argparse import ast import ctypes @@ -130,7 +130,7 @@ def oracle(args): if pci != args.pci.lower(): raise RuntimeError(f"selected GPU {pci} differs from requested GPU {args.pci}") args.output.mkdir(parents=True, exist_ok=False) - prompt = example_prompt(args.checkpoint) + prompt = args.prompt.read_text(encoding="utf-8") if args.prompt else example_prompt(args.checkpoint) image_paths = [args.image, *args.extra_image] messages = [{"role": "user", "content": [*({"type": "image", "image": str(p)} for p in image_paths), {"type": "text", "text": prompt}]}] @@ -227,6 +227,69 @@ def hook(_module, _args, output): print(json.dumps({"output_tokens": len(ids), "seconds": seconds, "text": output_text}, ensure_ascii=False), flush=True) +def pages(args): + """Draw the OCR test pages, our own text in Liberation fonts: page.png and one-line.png.""" + from PIL import Image, ImageDraw, ImageFont + + def font(name, size): + return ImageFont.truetype(str(args.fonts / f"Liberation{name}.ttf"), size) + + args.output.mkdir(parents=True, exist_ok=True) + page = Image.new("RGB", (1240, 1754), "white") + d = ImageDraw.Draw(page) + d.text((80, 60), "TENSORFOLD READER NOTES", font=font("Sans-Regular", 22), fill="black") + d.text((1000, 60), "March 2026", font=font("Sans-Regular", 22), fill="black") + d.line((80, 95, 1160, 95), fill="black", width=2) + d.text((80, 150), "Reading a Page Twice", font=font("Sans-Bold", 48), fill="black") + d.text((80, 260), "1. Why the second read matters", font=font("Sans-Bold", 32), fill="black") + for i, line in enumerate(["A reader that sees the same page twice must give the same blocks twice.", + "We feed one scanned page to two readers and count the tokens each", + "one writes. The table below lists the counts for a page and a line."]): + d.text((80, 330 + 44 * i), line, font=font("Serif-Regular", 30), fill="black") + d.text((80, 520), "Table 1: Tokens written per input", font=font("Sans-Regular", 26), fill="black") + x, y = [80, 420, 760, 1100], [570, 630, 690, 750, 810] + for row in y: + d.line((x[0], row, x[3], row), fill="black", width=2) + for column, top in [(x[0], y[0]), (x[1], y[0]), (x[3], y[0]), (x[2], y[1])]: + d.line((column, top, column, y[4]), fill="black", width=2) + d.rectangle((x[0] + 2, y[1] - 1, x[1] - 1, y[1] + 1), fill="white") # "Reader" spans both header rows + for at, cell in [((100, 615), "Reader"), ((520, 585), "Tokens written"), ((560, 645), "Page"), ((900, 645), "Line")]: + d.text(at, cell, font=font("Sans-Bold", 28), fill="black") + for r, row in enumerate([("First reader", "612", "41"), ("Second reader", "598", "39")]): + for c, cell in enumerate(row): + d.text((x[c] + 20, y[2 + r] + 15), cell, font=font("Sans-Regular", 28), fill="black") + d.text((80, 880), "2. What we check", font=font("Sans-Bold", 32), fill="black") + for i, item in enumerate(["The blocks come back in reading order.", "Each box stays inside the page.", + "A table keeps its spanned header cell."]): + d.text((100, 950 + 50 * i), "\u2022 " + item, font=font("Serif-Regular", 30), fill="black") + d.text((612, 1680), "3", font=font("Sans-Regular", 24), fill="black") + page.save(args.output / "page.png") + line = Image.new("RGB", (960, 96), "white") + ImageDraw.Draw(line).text((24, 24), "TensorFold reads one line of text.", font=font("Sans-Regular", 36), fill="black") + line.save(args.output / "one-line.png") + print(json.dumps({name: sha(args.output / name) for name in ["page.png", "one-line.png"]}, indent=2)) + + +def record(args): + """Write an OCR fixture from an oracle run: its input, frame and hashes, without the model's text.""" + result = json.loads((args.reference_dir / "result.json").read_text()) + if len(result["prepared_images"]) != 1: + raise ValueError("an OCR fixture records a run on one image") + from PIL import Image + with Image.open(args.image) as image: + original = list(image.size) + if sha(args.image) != result["image_sha256"]: + raise ValueError(f"{args.image} is not the image the oracle read") + ids = np.asarray(result["output_ids"], dtype=" tuple[int, Info]: + """tf_video_open's return code and the stream's facts.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc == 0: + self.lib.tf_video_close(handle) + return rc, info + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """The frames delivered (hashed, or index only), the seconds tf_video_frames took and its return code.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc != 0: + return [], 0.0, rc + out: list[tuple[int, str]] = [] + + def take(_ctx, _k, index, rgb, width, height, stride): + if not hashed: # timing: the bytes are left alone, so the decoder's threads never wait on Python + out.append((index, "")) + return 0 + raw = ctypes.string_at(rgb, stride * (height - 1) + width * 3) + out.append((index, hashlib.sha256(b"".join(raw[r * stride:r * stride + width * 3] + for r in range(height))).hexdigest())) + return 0 + + callback = FRAME_FN(take) + start = time.perf_counter() + rc = self.lib.tf_video_frames(handle, (ctypes.c_int * len(indices))(*indices), len(indices), callback, None) + seconds = time.perf_counter() - start + self.lib.tf_video_close(handle) + return out, seconds, rc + + +class PyAV: + """transformers' read_video_pyav: decode from the start, keep the wanted indices, to_ndarray(format="rgb24").""" + + def __init__(self): + import av + self.av, self.name, self.version = av, f"PyAV {av.__version__}", str(av.ffmpeg_version_info) + + def open(self, data: bytes) -> tuple[int, None]: + """0 when PyAV opens the container with a video stream, -2 (TF_NO_STREAM) without one, else -1 (TF_INVALID).""" + import io + try: + with self.av.open(io.BytesIO(data)) as c: + return (0 if c.streams.video else -2), None + except self.av.error.FFmpegError: + return -1, None + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """As Library.frames; an FFmpeg error counts as TF_DECODE and an early end as TF_SHORT.""" + import io + want, out, rc = set(indices), [], 0 + start = time.perf_counter() + try: + with self.av.open(io.BytesIO(data)) as c: + c.seek(0) + for i, frame in enumerate(c.decode(video=0)): + if i > indices[-1]: + break + if i in want: + rgb = frame.to_ndarray(format="rgb24") + out.append((i, hashlib.sha256(rgb.tobytes()).hexdigest() if hashed else "")) + except self.av.error.FFmpegError: + rc = TF_DECODE + seconds = time.perf_counter() - start + return out, seconds, rc or (len(out) if len(out) == len(indices) else TF_SHORT) + + +def index_sets(info: Info) -> dict[str, list[int]]: + """Every frame, GLM's 2 frames a second, the first three (an early stop) and the last (the end-of-stream drain).""" + n, rate = int(info.frames), info.rate_num / info.rate_den if info.rate_den else 0.0 + two = sorted({min(n - 1, round(i * rate / 2)) for i in range(int(n / rate * 2) + 1)}) if rate else [0] + return {"all": list(range(n)), "2fps": two, "first3": list(range(min(3, n))), "last": [n - 1]} + + +def best(source, data: bytes, indices: list[int], repeat: int, hashed: bool) -> tuple[list, float, int]: + frames, seconds, rc = source.frames(data, indices, hashed) + for _ in range(repeat - 1): + seconds = min(seconds, source.frames(data, indices, hashed)[1]) + return frames, seconds, rc + + +def compare(lib, ref, clip: Path, repeat: int, rows: dict) -> int: + """One clip's sets against the reference (or timed alone when ref is None); returns the sets that differ.""" + data = clip.read_bytes() + rc, info = lib.open(data) + ref_rc, ref_info = ref.open(data) if ref else (rc, info) + if rc != 0 or ref_rc != 0 or info.frames <= 0: # a damaged clip may not open, or may count no frames + rows["open"] = {"test_rc": rc, "reference_rc": ref_rc, "frames": info.frames} + theirs = (ref_info or info).frames # PyAV counts nothing here: libtfvideo's count stands for both + ok = ref is not None and rc == ref_rc and (info.frames > 0) == (theirs > 0) + print(f"{'ok ' if ok else 'FAIL'} {clip} open: test rc {rc}, reference rc {ref_rc}, {info.frames} frames") + return int(not ok) + failures = 0 + for name, indices in index_sets(info).items(): + mine, mine_s, mine_rc = best(lib, data, indices, repeat, ref is not None) + theirs, theirs_s, theirs_rc = best(ref, data, indices, repeat, True) if ref else (mine, 0.0, mine_rc) + same = sum(a == b for a, b in zip(mine, theirs)) + ok = mine == theirs and mine_rc == theirs_rc and (ref is not None or mine_rc == len(indices)) + failures += not ok + rows[name] = {"indices": indices, "test": mine, "reference": theirs if ref else None, "test_rc": mine_rc, + "reference_rc": theirs_rc, "test_s": mine_s, "reference_s": theirs_s} + codes = f", rc {mine_rc}/{theirs_rc}" if (mine_rc, theirs_rc) != (len(indices), len(indices)) else "" + timing = f"test {mine_s:.3f} s" + (f", reference {theirs_s:.3f} s" if ref else "") + what = "frames equal" if ref else "frames delivered" + print(f"{'ok ' if ok else 'FAIL'} {clip} {name}: {same}/{len(indices)} {what}{codes}, {timing}") + return failures + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--lib", required=True, type=Path, help="the libtfvideo build under test") + ap.add_argument("--against", type=Path, help="another libtfvideo build; PyAV when omitted") + ap.add_argument("--time-only", action="store_true", + help="no reference and no hashing: time tf_video_frames alone (its callback leaves the bytes)") + ap.add_argument("--repeat", type=int, default=1, help="decodes per set; the fastest time is reported") + ap.add_argument("--json", type=Path, help="write every frame's hash, return code and each set's time here") + ap.add_argument("clips", nargs="+", type=Path) + args = ap.parse_args() + lib = Library(args.lib) + ref = None if args.time_only else Library(args.against) if args.against else PyAV() + print(f"test {lib.name} (FFmpeg {lib.version}); reference " + + (f"{ref.name} (FFmpeg {ref.version})" if ref else "none (--time-only)")) + failures, report = 0, {"lib": lib.name, "reference": ref.name if ref else None, "clips": {}} + for clip in args.clips: + rows = report["clips"][str(clip)] = {} + try: + failures += compare(lib, ref, clip, args.repeat, rows) + except Exception as e: # one unreadable clip is a failed clip, not the end of the run + rows["error"] = repr(e) + failures += 1 + print(f"FAIL {clip}: {e!r}") + if args.json: + args.json.write_text(json.dumps(report, indent=1)) + print(f"{failures} set(s) differ" if failures else "every frame equal") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/zig/build/checkpoint_io_host.zig b/zig/build/checkpoint_io_host.zig index 77767ffa710..583d42f9b1e 100644 --- a/zig/build/checkpoint_io_host.zig +++ b/zig/build/checkpoint_io_host.zig @@ -11,6 +11,8 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget) void { step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"checkpoint admission:"} })).step); if (mode == .debug) b.step("test-modelopt-part-host", "The shared ModelOpt part reader (core/modelopt_part.zig) without a GPU") .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"part:"} })).step); + if (mode == .debug) b.step("test-fp8-block-host", "The shared block-FP8 reader and exact scale widening (core/quant/fp8.zig)") + .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"fp8 block:"} })).step); const qwen = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/loader.zig"), .target = target, .optimize = mode, .link_libc = true, .imports = &.{.{ .name = "core", .module = core }} }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = qwen, .filters = &.{"Qwen shard integrity:"} })).step); } diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index c0c8b503336..950d2d5586f 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -97,6 +97,7 @@ pub const kernels = [_]Kernel{ .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order .{ .name = "nvfp4_experts_tile", .flags = &.{"-O3"}, .min_sm = 80 }, + .{ .name = "experts_bf16", .flags = &.{"-O3"} }, // ours: grouped bf16 experts on the plan, mma fragment blocks .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "nvfp4_experts32", .flags = &.{"-O3"} }, // ours: GLM's routed experts at block 32 with csf0 scales .{ .name = "rank_sum", .flags = &.{"-O3"} }, // ours: the rank-order sum of a reduce-scatter (rank_scatter.zig) diff --git a/zig/build/cuda_server_host.zig b/zig/build/cuda_server_host.zig index c77531a9b00..31e7b5314f5 100644 --- a/zig/build/cuda_server_host.zig +++ b/zig/build/cuda_server_host.zig @@ -29,6 +29,7 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = mods.tokenizer, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), .filters = &.{"EOS policy:"} })).step); + b.step("test-serve-flags-host", "Serve flag parsing a later flag replaces, without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--compact-at:"} })).step); b.step("test-dory-cli-host", "Dory CUDA loop option forwarding and invalid-value refusals").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--device and --segments:"} })).step); b.step("test-dory-chat-host", "Dory chat template, EOS, reasoning split and qwen3_coder call on the CPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Dory chat:"} })).step); b.step("test-embeddings-protocol-host", "Embedding request and response contracts without an admitted model or GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"embeddings protocol:"} })).step); @@ -66,6 +67,7 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod b.step("test-http-limits-host", "Total HTTP read deadlines and connection admission without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"HTTP limits:"} })).step); b.step("test-server-recovery-host", "Incomplete replies, token allocation and bounded SSE storage without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Server recovery:"} })).step); b.step("test-prompt-interior-host", "Last-user metadata excludes later messages without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"interior prompt:"} })).step); + b.step("test-ocr-host", "OCR route readers, shared page steps and dialects without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"ocr "} })).step); const server_test = b.addRunArtifact(b.addTest(.{ .root_module = server_tests })); @import("request_starts_host.zig").steps(b, native_modules.engines, server_tests); const json_test = b.addRunArtifact(b.addTest(.{ .root_module = server_json })); diff --git a/zig/build/dots_mocr.zig b/zig/build/dots_mocr.zig index aa60eb0c418..90765b85fa1 100644 --- a/zig/build/dots_mocr.zig +++ b/zig/build/dots_mocr.zig @@ -5,13 +5,15 @@ pub fn hostTests(b: *std.Build, test_step: *std.Build.Step) void { const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = b.graph.host, .link_libc = true }); core.addImport("gpu_memory", @import("gpu_memory.zig").module(b, b.graph.host)); core.addImport("tokenizer", b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = b.graph.host, .link_libc = true })); + const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = b.graph.host, .link_libc = true }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| { const family = b.createModule(.{ .root_source_file = b.path("zig/src/families/dots_mocr/dots_mocr.zig"), .target = b.graph.host, .optimize = mode, .link_libc = true, - .imports = &.{.{ .name = "core", .module = core }}, + .imports = &.{ .{ .name = "core", .module = core }, .{ .name = "engine_api", .module = api } }, }); const run = b.addRunArtifact(b.addTest(.{ .root_module = family, .filters = &.{"dots:"} })); run.has_side_effects = true; diff --git a/zig/build/dsv41_assets.zig b/zig/build/dsv41_assets.zig index 06c612faa61..aa950c2ffa5 100644 --- a/zig/build/dsv41_assets.zig +++ b/zig/build/dsv41_assets.zig @@ -2,7 +2,8 @@ const std = @import("std"); pub fn steps(b: *std.Build) void { const step = b.step("dsv41-assets", "Install CPU-only DeepSeek token-map and AOT-plan generators"); - const asset_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false); + // ReleaseSafe whatever -Doptimize says: a 512K aot-needs pass took ~45 min in Debug on a GB10, under 7 here (#2270) + const asset_module = @import("dsv41_cuda.zig").moduleAt(b, b.graph.host, .ReleaseSafe, "asset_main.zig", false); const assets = b.addExecutable(.{ .name = "tf-dsv41-assets", .root_module = asset_module }); step.dependOn(&b.addInstallArtifact(assets, .{}).step); const token_map = b.addExecutable(.{ .name = "tf-dsv41-token-map", .root_module = b.createModule(.{ @@ -12,6 +13,6 @@ pub fn steps(b: *std.Build) void { .link_libc = true, }) }); step.dependOn(&b.addInstallArtifact(token_map, .{}).step); - const tests = b.addRunArtifact(b.addTest(.{ .root_module = asset_module })); + const tests = b.addRunArtifact(b.addTest(.{ .root_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false) })); b.step("test-dsv41-assets-host", "DeepSeek asset commands refuse GPU operations and malformed arguments").dependOn(&tests.step); } diff --git a/zig/build/dsv41_cuda.zig b/zig/build/dsv41_cuda.zig index 6b64f7d380e..b3e19f0cc64 100644 --- a/zig/build/dsv41_cuda.zig +++ b/zig/build/dsv41_cuda.zig @@ -2,11 +2,14 @@ const std = @import("std"); const shared = @import("dsv41_shared.zig"); const units = @import("dsv41_units.zig"); pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *std.Build.Module { + return runtimeAt(b, target, .Debug, images); +} +pub fn runtimeAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, images: bool) *std.Build.Module { const options = @import("cuda_runtime.zig").kernelOptions(b, images, null); - const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); cuda.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); cuda.addOptions("kernel_options", options); - cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = .Debug })); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); if (images) { const empty = b.addWriteFiles().add("compile-only.fatbin", "SYNTHETIC-COMPILE-ONLY"); const all = @import("cuda.zig").kernels; @@ -15,16 +18,20 @@ pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *s return cuda; } pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8, images: bool) *std.Build.Module { - const cuda = runtime(b, target, images); + return moduleAt(b, target, .Debug, file, images); +} +/// `module` with the runtime, kernel and family modules at `optimize`. +pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, images: bool) *std.Build.Module { + const cuda = runtimeAt(b, target, optimize, images); const opts = b.addOptions(); opts.addOption(bool, "with_kernels", images); opts.addOption(u32, "family_scope", 41); - const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = optimize, .link_libc = true }); dk.addImport("cuda", cuda); dk.addOptions("dsv41_kernel_options", opts); const empty = b.addWriteFiles().add("compile-only-dsv41.fatbin", "SYNTHETIC-COMPILE-ONLY"); for (units.units ++ units.optional) |k| dk.addAnonymousImport(b.fmt("dsv41_fatbin_{s}", .{k.name}), .{ .root_source_file = empty }); - return assemble(b, target, .Debug, file, cuda, dk); + return assemble(b, target, optimize, file, cuda, dk); } /// The family's root `file` over the runtime `cuda` and kernel module `dk`, every family module at `optimize`. pub fn assemble(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, cuda: *std.Build.Module, dk: *std.Build.Module) *std.Build.Module { diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 70c3f26e360..88c0f6fb363 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -518,15 +518,15 @@ const groups = [_]Source{ }; /// What the group sources include, so an edit to one rebuilds them. const headers = [_][]const u8{ - "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", - "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", - "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", - "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", - "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", - "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", "flashnext/streams.hpp", - "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", "quant/bf16_decoder.hpp", - "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", "columns_controls.hpp", - "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", + "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", + "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", + "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", + "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", + "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", + "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/matrix_item.hpp", "tiles/plan.hpp", "tiles/stream.hpp", + "flashnext/streams.hpp", "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", + "quant/bf16_decoder.hpp", "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", + "columns_controls.hpp", "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. const flags = [_][]const u8{ "-D__HIP_PLATFORM_AMD__=1", "-DUSE_ROCM=1", "-DHIPBLAS_V2", "-fPIC", "-DCUDA_HAS_FP16=1", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", "-std=c++20", "-fno-gpu-rdc", "-mno-wavefrontsize64", "-ffp-contract=off" }; diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index 85206a8dd76..fed6dc007a9 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -14,6 +14,7 @@ pub fn steps(b: *std.Build) void { .macos => { wire(b, "test-host-ci-macos-fast", "Nightly host checks without a GPU or shader compiler", &macos_fast); wire(b, "test-host-ci-macos-adapters", "Nightly host checks without a GPU or shader compiler", &macos_adapters); + wire(b, "test-host-ci-macos-native", "Nightly Metal engine host checks with the Metal toolchain", &macos_native); wire(b, "test-host-ci-macos-deepseek", "Nightly host checks without a GPU or shader compiler", &macos_deepseek); wire(b, "test-host-ci-macos-lanes", "Nightly host checks without a GPU or shader compiler", &macos_lanes); wire(b, "test-host-ci-macos-deferred", "Nightly host checks without a GPU or shader compiler", &macos_deferred); @@ -250,6 +251,7 @@ const macos_adapters = [_][]const u8{ "check-nemotron-prefill-widths-host", "check-qwen35-cache-host", "check-sessions-linux", + "check-video-exports-linux", "check-video-layout-linux", "compile-disk-identity-metal-host", "compile-dsv41-engram-alloc-aarch64", @@ -289,6 +291,8 @@ const macos_adapters = [_][]const u8{ "test-verbs", }; +const macos_native = [_][]const u8{"test-native"}; + const macos_deepseek = [_][]const u8{ "test-dsv41-device-host", "test-dsv41-draft", diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 88f939e7984..2faa0416cf6 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -16,7 +16,7 @@ pub const Shared = struct { }; /// The 27B's Metal sources, embedded by name (zig/kernels/metal/qwen27). -const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg", "prompt_glue", "prompt_state" }; +const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg" }; /// The shared modules, imported into `engine` (whose `core` is itself) for the server and the release binary. pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module) Shared { diff --git a/zig/build/video.zig b/zig/build/video.zig index 3e33b3fae41..e6752a24734 100644 --- a/zig/build/video.zig +++ b/zig/build/video.zig @@ -15,6 +15,7 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, host_tests: *std.B b.step("test-video-library-host", "Video candidate order with a recording opener, without real libraries").dependOn(&run.step); host_tests.dependOn(&run.step); const cross = b.step("check-video-layout-linux", "Compile the Linux video probe and stub without execution"); + const exports = b.step("check-video-exports-linux", "Link a planted Linux stub as `video` links; read its exports"); inline for (.{ .aarch64, .x86_64 }) |arch| { const linux = b.resolveTargetQuery(.{ .cpu_arch = arch, .os_tag = .linux, .abi = .gnu }); const probe = layoutProbe(b, linux); @@ -25,6 +26,15 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, host_tests: *std.B const lib = b.addLibrary(.{ .name = "tfvideo", .linkage = .dynamic, .root_module = stub }); _ = lib.getEmittedBin(); cross.dependOn(&lib.step); + const planted = b.createModule(.{ .target = linux, .link_libc = true }); + for ([_][]const u8{ "zig/tests/tfvideo_stub.c", "zig/tests/tfvideo_planted.c" }) |c| + planted.addCSourceFile(.{ .file = b.path(c), .flags = &.{ "-std=c11", "-Wall", "-Wextra", "-Werror" } }); + const scripted = b.addLibrary(.{ .name = "tfvideo", .linkage = .dynamic, .root_module = planted }); + scripted.link_gc_sections = false; + exportOnlyTfVideo(b, scripted); + const check = b.addSystemCommand(&.{ "python3", "tools/zig/check_tfvideo_exports.py" }); + check.addArtifactArg(scripted); + exports.dependOn(&check.step); } if (b.graph.host.result.os.tag == .linux) { const layout = b.addSystemCommand(&.{ "python3", "tools/zig/test_tfvideo_layout.py", "--probe" }); @@ -56,6 +66,7 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, host_tests: *std.B }) |a| mod.addObjectFile(root.path(b, b.pathJoin(&.{ "lib", a }))); mod.linkSystemLibrary("m", .{}); const lib = b.addLibrary(.{ .name = "tfvideo", .linkage = .dynamic, .root_module = mod }); + exportOnlyTfVideo(b, lib); step.dependOn(&b.addInstallArtifact(lib, .{ .dest_dir = .{ .override = .{ .custom = "native/lib" } } }).step); } @@ -64,3 +75,9 @@ fn layoutProbe(b: *std.Build, target: std.Build.ResolvedTarget) *std.Build.Step. root.addImport("tfvideo", b.createModule(.{ .root_source_file = b.path("zig/src/core/tfvideo.zig"), .target = target, .link_libc = true })); return b.addExecutable(.{ .name = "tfvideo-layout", .root_module = root }); } + +/// ELF exports only tf_video_*: FFmpeg's x86_64 assembly then reaches its C tables PC-relative inside the library. +fn exportOnlyTfVideo(b: *std.Build, lib: *std.Build.Step.Compile) void { + if (lib.rootModuleTarget().ofmt != .elf) return; + lib.setVersionScript(b.addWriteFiles().add("tfvideo.map", "{ global: tf_video_*; local: *; };\n")); +} diff --git a/zig/gemma4_shared_affine_test.zig b/zig/gemma4_shared_affine_test.zig index 33e99b67ad7..4ab73b90968 100644 --- a/zig/gemma4_shared_affine_test.zig +++ b/zig/gemma4_shared_affine_test.zig @@ -109,8 +109,8 @@ test "landed Metal source bytes are pinned without shader compilation" { { var hash: [32]u8 = undefined; std.crypto.hash.sha2.Sha256.hash(ks.core_lane_projection, &hash, .{}); - // 66155d5ee's bytes (8-bit lane codes and producer sums, Refs #2502, under the metal_lane_q8 contract) - try std.testing.expectEqualStrings("159f39224b5bb60fd2ced301496d91662d2958056bd291d13d8a5db5186dd41a", &std.fmt.bytesToHex(hash, .lower)); + // Packed affine lanes with producer sums and independent prompt row blocks. + try std.testing.expectEqualStrings("7befe2028dae5036541f311f230a02d5652e634450ef16dc66dce6e15e935277", &std.fmt.bytesToHex(hash, .lower)); } } test "released lane parameters and conservative scratch size" { diff --git a/zig/kernels/cuda/contract_controls_set.cu b/zig/kernels/cuda/contract_controls_set.cu index e4608a5b9f7..3624ac10b03 100644 --- a/zig/kernels/cuda/contract_controls_set.cu +++ b/zig/kernels/cuda/contract_controls_set.cu @@ -3,3 +3,4 @@ #include "contract_controls.cu" #include "contract_js_controls.cuh" #include "contract_dense_controls.cuh" +#include "contract_experts_bf16_controls.cuh" diff --git a/zig/kernels/cuda/contract_experts_bf16_controls.cuh b/zig/kernels/cuda/contract_experts_bf16_controls.cuh new file mode 100644 index 00000000000..1af8126ec29 --- /dev/null +++ b/zig/kernels/cuda/contract_experts_bf16_controls.cuh @@ -0,0 +1,24 @@ +// Gate 4's control for experts_bf16.cu: its fp32 down with each pair's sum rounded to bf16 at every input. +#pragma once + +extern "C" __global__ void tf_ctl_bf16_experts_bf16_sums(const __nv_bfloat16* X, int x_stride, const uint4* W, int KG, int NB, + const int* items, const int* counts, const int* members, float* out, + int N) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= N) return; + const int cb = col / 32, j = (col % 32) / 8; + for (int it = blockIdx.y; it < counts[0]; it += gridDim.y) { + const int e = items[3 * it], first = items[3 * it + 1], cnt = items[3 * it + 2]; + const __nv_bfloat16* blk = reinterpret_cast(W + (static_cast(e) * NB + cb) * KG * 128); + for (int i = 0; i < cnt; ++i) { + const int pr = members[first + i]; + float s = 0.0f; + for (int k = 0; k < KG * 32; ++k) { + const int q = k % 16, lane = 4 * (col % 8) + q / 4; + const float w = __bfloat162float(blk[(static_cast(k / 32) * 128 + 32 * j + lane) * 8 + 4 * (k % 32 / 16) + q % 4]); + s = __bfloat162float(__float2bfloat16_rn(s + tf_ctl_x(X, static_cast(pr) * x_stride + k) * w)); + } + out[static_cast(pr) * N + col] = s; + } + } +} diff --git a/zig/kernels/cuda/experts_bf16.cu b/zig/kernels/cuda/experts_bf16.cu new file mode 100644 index 00000000000..df6d756ec18 --- /dev/null +++ b/zig/kernels/cuda/experts_bf16.cu @@ -0,0 +1,114 @@ +// Ours: grouped bf16 experts on mma-fragment blocks; a row's sums chain in input order, whatever its tile or form. + +#include +#include + +#include "experts.cuh" + +namespace tf_bf16_experts { + +constexpr int WV = 4; // uint4 a lane a block + +__device__ __forceinline__ uint4 ld_w(const uint4* p) { return tf_common::ld_nc(p); } + +// RT row tiles of 16 pairs share each block's fragments; a unit is one item's 32-column block. +template +__device__ __forceinline__ void body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint4* __restrict__ W, int KG, int NB, const int* __restrict__ items, + const int* __restrict__ counts, const int* __restrict__ members, + void* __restrict__ out, int N, int skip) { + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x * WARPS + (threadIdx.x >> 5); unit < units; unit += gridDim.x * WARPS) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + if (e == skip) continue; + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (32 * WV) + lane; + for (int r0 = 0; r0 < cnt; r0 += 16 * RT) { + bool v[RT][2]; + int pr[RT][2]; + const __nv_bfloat16* x[RT][2]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = r0 + 16 * r + gq + 8 * h; + v[r][h] = i < cnt; + pr[r][h] = v[r][h] ? __ldg(members + first + i) : 0; + x[r][h] = X + (size_t)(slots ? pr[r][h] / slots : pr[r][h]) * x_stride + 4 * t; + } + float acc[1][RT][NTW][4]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[0][r][j][0] = acc[0][r][j][1] = acc[0][r][j][2] = acc[0][r][j][3] = 0.f; + const uint2 zero = make_uint2(0u, 0u); + for (int g = 0; g < KG; ++g) { + uint4 w[NTW]; +#pragma unroll + for (int j = 0; j < NTW; ++j) w[j] = ld_w(blk + (size_t)g * (32 * WV) + 32 * j); +#pragma unroll + for (int r = 0; r < RT; ++r) { + uint2 xa[2], xb[2]; +#pragma unroll + for (int q = 0; q < 2; ++q) { + xa[q] = v[r][0] ? __ldg(reinterpret_cast(x[r][0] + 32 * g + 16 * q)) : zero; + xb[q] = v[r][1] ? __ldg(reinterpret_cast(x[r][1] + 32 * g + 16 * q)) : zero; + } +#pragma unroll + for (int q = 0; q < 2; ++q) +#pragma unroll + for (int j = 0; j < NTW; ++j) + mma(acc[0][r][j], xa[q].x, xb[q].x, xa[q].y, xb[q].y, q ? w[j].z : w[j].x, q ? w[j].w : w[j].y); + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr[r][0], pr[r][1], v[r][0], v[r][1], 0.f); + } + } +} + +} // namespace tf_bf16_experts + +// Decode (items of up to 16 pairs) and prompt (up to 64) entries per epilogue: the NVFP4 entries' arguments, no scales. +#define TF_BF16_EXPERTS(NAME, EPI, RT, WARPS) \ + extern "C" __global__ void __launch_bounds__(WARPS * 32) \ + NAME(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, int NB, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int skip) { \ + tf_bf16_experts::body(X, x_stride, slots, W, KG, NB, items, counts, members, out, N, skip); \ + } +TF_BF16_EXPERTS(tf_bf16_experts_up, 1, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_f32, 0, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_bf16, 3, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_up, 1, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_f32, 0, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_bf16, 3, 4, 4) + +// The repack, one thread a uint4: matrix e's [k][n] (in_out, x @ W) or [n][k] (nn.Linear) rows into its blocks. +extern "C" __global__ void tf_bf16_experts_pack(const __nv_bfloat16* __restrict__ src, long long e_stride, int ld, + int n0, int k0, int in_out, int n, int k, uint4* __restrict__ dst) { + const int e = blockIdx.y; + const long long per = (long long)(n / 32) * (k / 32) * 32 * tf_bf16_experts::WV; + const long long q = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (q >= per) return; + const int lane = (int)(q % 32), j = (int)(q / 32 % tf_bf16_experts::WV); + const long long b = q / (32 * tf_bf16_experts::WV); + const int kg = k / 32, g = (int)(b % kg), cb = (int)(b / kg); + const int row = n0 + cb * 32 + 8 * j + (lane >> 2), col = k0 + g * 32 + 4 * (lane & 3); + const __nv_bfloat16* s = src + e * e_stride; + uint16_t h[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int c = col + (i & 3) + 16 * (i >> 2); + const __nv_bfloat16 v = in_out ? s[(size_t)c * ld + row] : s[(size_t)row * ld + c]; + h[i] = __bfloat16_as_ushort(v); + } + uint4 o; + o.x = h[0] | (uint32_t)h[1] << 16; + o.y = h[2] | (uint32_t)h[3] << 16; + o.z = h[4] | (uint32_t)h[5] << 16; + o.w = h[6] | (uint32_t)h[7] << 16; + dst[e * per + q] = o; +} diff --git a/zig/kernels/hip/contract_controls.hip b/zig/kernels/hip/contract_controls.hip index 231f769c61c..327da668856 100644 --- a/zig/kernels/hip/contract_controls.hip +++ b/zig/kernels/hip/contract_controls.hip @@ -7,6 +7,7 @@ #include "attention/prefill.hpp" #include "quant/dense_tiles.hpp" +#include "quant/mlx_tiles.hpp" #include "tiles/wmma12_gemm_k32_ref.hpp" __device__ __forceinline__ float tf_ctl_x(const uint16_t* x, int64_t i) { @@ -196,6 +197,18 @@ tf_ctl_fa_paged_unmasked_bf16(const float* q, const hip_bfloat16* k, const hip_b 0, 0, 0, 0, table, count); } +// Gate 3: the routed item tile's few-row passes stop after a wave's first, leaving the block's other columns unwritten. +struct TfCtlItemFewOnePass { + using Dot = tf::rocm::DotBF16; + using Epi = tf::rocm::GemmOut; + static constexpr bool every_pass = false; +}; + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_item_few_one_pass(tf::rocm::Affine a) { + using namespace tf::rocm; + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, 4, TfCtlItemFewOnePass>(a); +} + // Gate 3: the big block loading its operands without the stores' swizzle, and without its K loop's barriers. struct TfCtlUnswizzled { static constexpr bool swizzled_loads = false, ordered = true; diff --git a/zig/kernels/hip/quant/mlx_tiles.hpp b/zig/kernels/hip/quant/mlx_tiles.hpp index a3f6ca68940..70bdd72d1e5 100644 --- a/zig/kernels/hip/quant/mlx_tiles.hpp +++ b/zig/kernels/hip/quant/mlx_tiles.hpp @@ -45,11 +45,18 @@ __global__ void __launch_bounds__(256) affine_wmma_skew(Affine a) { wmma_gemm_tile, BF16Act, WmmaBF16, GemmOut, true>(a); } -// ---- the same on RM x 16 rows a block, for a routed item of few rows ---- +// ---- the same on RM x 16 rows a block, for a routed item of few rows; K-parallel passes for one of very few ---- + +// The item tile's few-row passes as served: BF16 dot2 into the fp32 sums, every pass of a wave's columns. +struct ItemFew { + using Dot = DotBF16; + using Epi = GemmOut; + static constexpr bool every_pass = true; +}; template __global__ void __launch_bounds__(256) affine_wmma_item(Affine a) { - wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM>(a); + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM, ItemFew>(a); } // ---- the decode stream tile ---- diff --git a/zig/kernels/hip/tiles/gemm_kp.hpp b/zig/kernels/hip/tiles/gemm_kp.hpp index d759c7da59f..f4b7da96a48 100644 --- a/zig/kernels/hip/tiles/gemm_kp.hpp +++ b/zig/kernels/hip/tiles/gemm_kp.hpp @@ -12,26 +12,23 @@ namespace rocm { // The lanes a column takes, a group each: 32, or 16 or 8 when it has few groups (a wave then takes more columns). __host__ __device__ inline int kp_lanes(int groups) { return groups > 16 ? 32 : groups > 8 ? 16 : 8; } -// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. -template -__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { - const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows - if (!Dec::take_item(a, blockIdx.z)) return; +// A lane's row of each of a wave's LDS tables, what `lds(lane)` names: R * CB dots, R sums of x, CB scales, CB biases. +struct KpRows { + float* part; + float* spart; + float* scl; + float* bil; +}; + +// A wave's passes over its columns from `colw`, CB a lane and R rows a pass: rows from `first`, every `step` when LOOP. +template +__device__ __forceinline__ void gemm_kp_wave(typename Dec::Args& a, int colw, int first, int step, Lds lds) { constexpr int NO = R * CB; // outputs of a lane set - constexpr int PS = NO + 1; constexpr int OPL = (4 * NO + 31) / 32; // outputs a lane folds, at most const int groups = a.k / a.group; const int gp = kp_lanes(groups); // the lanes of a set; a wave has 32 / gp sets const int sets = 32 / gp; - const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; - if (colw >= a.n) return; - __shared__ float part[WAVES][32][PS]; - __shared__ float spart[WAVES][32][R + 1]; - __shared__ float scl[WAVES][32][CB + 1]; - __shared__ float bil[WAVES][32][CB + 1]; - const int lane = threadIdx.x & 31; - const int wave = threadIdx.x >> 5; const int per = a.group / 32; const int rounds = (groups + gp - 1) / gp; const int gi = lane & (gp - 1); @@ -128,14 +125,14 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { // this group's results to LDS, then the round's groups in order #pragma unroll for (int r = 0; r < R; ++r) { - spart[wave][lane][r] = sx[r]; + lds(lane).spart[r] = sx[r]; #pragma unroll - for (int c = 0; c < CB; ++c) part[wave][lane][r * CB + c] = d[r][c]; + for (int c = 0; c < CB; ++c) lds(lane).part[r * CB + c] = d[r][c]; } #pragma unroll for (int c = 0; c < CB; ++c) { - scl[wave][lane][c] = Dec::scale_value(a, term[c]); - bil[wave][lane][c] = Dec::bias_value(a, term[c]); + lds(lane).scl[c] = Dec::scale_value(a, term[c]); + lds(lane).bil[c] = Dec::bias_value(a, term[c]); } if (round + 1 < rounds) tables(g + gp); __syncwarp(); @@ -150,8 +147,8 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { const int c = o % CB; float v = acc[t]; for (int gg = 0; gg < count; ++gg) { - v = Dec::fold_scale(v, part[wave][base + gg][o], scl[wave][base + gg][c]); - v = Dec::fold_bias(v, spart[wave][base + gg][r], bil[wave][base + gg][c]); + v = Dec::fold_scale(v, lds(base + gg).part[o], lds(base + gg).scl[c]); + v = Dec::fold_bias(v, lds(base + gg).spart[r], lds(base + gg).bil[c]); } acc[t] = v; } @@ -168,11 +165,28 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { } }; if constexpr (LOOP) { - for (int row0 = (blockIdx.x % blocks) * R; row0 < a.m; row0 += blocks * R) pass(row0); + for (int row0 = first; row0 < a.m; row0 += step) pass(row0); } else { - pass((blockIdx.x % blocks) * R); + pass(first); } } +// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. +template +__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { + const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows + if (!Dec::take_item(a, blockIdx.z)) return; + const int sets = 32 / kp_lanes(a.k / a.group); + const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; + if (colw >= a.n) return; + __shared__ float part[WAVES][32][R * CB + 1]; + __shared__ float spart[WAVES][32][R + 1]; + __shared__ float scl[WAVES][32][CB + 1]; + __shared__ float bil[WAVES][32][CB + 1]; + const int wave = threadIdx.x >> 5; + const auto lds = [wave](int lane) { return KpRows{part[wave][lane], spart[wave][lane], scl[wave][lane], bil[wave][lane]}; }; + gemm_kp_wave(a, colw, (blockIdx.x % blocks) * R, blocks * R, lds); +} + } // namespace rocm } // namespace tf diff --git a/zig/kernels/hip/tiles/matrix_item.hpp b/zig/kernels/hip/tiles/matrix_item.hpp index 06195bb373e..734ccf5a056 100644 --- a/zig/kernels/hip/tiles/matrix_item.hpp +++ b/zig/kernels/hip/tiles/matrix_item.hpp @@ -4,13 +4,19 @@ #include "common/wmma.hpp" #include "tiles/gemm.hpp" +#include "tiles/gemm_kp.hpp" #include "tiles/plan.hpp" namespace tf { namespace rocm { +// The most rows a launch's cut sends to K-parallel passes, the rows a pass takes and the columns a lane takes in it. +constexpr int kItemFewRows = 8; +constexpr int kItemFewPass = 2; +constexpr int kItemFewCols = 2; + // Same staging, stage order and per-element sums as wmma_gemm_tile, so each output equals that tile's byte for byte. -template +template __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { #if TF_DEVICE_WMMA_GFX11 if (!Dec::take_item(a, blockIdx.z)) return; @@ -27,6 +33,28 @@ __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { const int tid = threadIdx.x; const int n0 = bx * kGemmN; + + // An item within the launch's cut takes K-parallel passes with Few's Dot and epilogue: same bytes, no matrix core. + if (a.m <= (a.route.few < kItemFewRows ? a.route.few : kItemFewRows)) { + constexpr int CB = kItemFewCols; + constexpr int R = kItemFewPass; + // a pass's partials, a row a lane: each wave's own slice of the two code stages, which this path leaves unused + constexpr int kLane = (R * CB + 1) + (R + 1) + 2 * (CB + 1); + static_assert(8 * 32 * kLane <= 2 * kGemmN * kLdw, "the few-row passes' partials fit the tile's code stages"); + const int wave = __builtin_amdgcn_readfirstlane(tid >> 5); + const auto lds = [wave](int lane) { + float* const row = reinterpret_cast(&ws[0][0]) + (wave * 32 + lane) * kLane; + return KpRows{row, row + R * CB + 1, row + R * CB + 1 + R + 1, row + R * CB + 1 + R + 1 + CB + 1}; + }; + // The 8 waves walk the block's columns side by side, CB a lane set a pass, R rows of the item a pass. + const int span = CB * (32 / kp_lanes(a.k / a.group)); + const int end = n0 + kGemmN < a.n ? n0 + kGemmN : a.n; + for (int colw = n0 + wave * span; colw < end; colw += 8 * span) { + gemm_kp_wave(a, colw, 0, R, lds); + if constexpr (!Few::every_pass) break; // the contract's control stops after a wave's first + } + return; + } const int words_row = Dec::template row_words(a); const int groups = a.k / a.group; const int stages = a.k / kGemmK; diff --git a/zig/kernels/hip/tiles/plan.hpp b/zig/kernels/hip/tiles/plan.hpp index 6cd355ba754..260d14b3b1e 100644 --- a/zig/kernels/hip/tiles/plan.hpp +++ b/zig/kernels/hip/tiles/plan.hpp @@ -10,6 +10,7 @@ struct Routing { const int* members = nullptr; int x_div = 1; int first = 0; // set by the kernel from its item + int few = 0; // an item of at most this many rows takes its tile's K-parallel passes, where the tile has them }; // Row r of this block's x and out: the plain matrix's r, or the routed item's r-th pair. diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 2fedb0c13c2..7e65ddb08b6 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -12,6 +12,9 @@ using namespace mpp::tensor_ops; #ifndef TF_PRECOMPUTE_SUMS #define TF_PRECOMPUTE_SUMS 0 #endif +#ifndef TF_ROW_BLOCKS +#define TF_ROW_BLOCKS 0 +#endif #ifndef TF_YS #define TF_YS TF_N #endif @@ -141,12 +144,8 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #endif #if TF_PRECOMPUTE_SUMS -#if TF_BITS == 8 #define TF_SUM_VALUE(row, g) ((row) < M ? XS[(g) * mdims[1] + (row)] : 0.0f) #else -#define TF_SUM_VALUE(row, g) XS[(g) * 16 + (row)] -#endif -#else #define TF_SUM_VALUE(row, g) tf_lane_xsum(X, row, g, M) #endif @@ -178,7 +177,21 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in const device float* XS [[buffer(5)]], #endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], - uint tgx [[threadgroup_position_in_grid]]) { +#if TF_ROW_BLOCKS + uint2 tg [[threadgroup_position_in_grid]] +#else + uint tgx [[threadgroup_position_in_grid]] +#endif +) { +#if TF_ROW_BLOCKS + const int row0 = int(tg.y) * 16, rows = mdims[0] - row0; + if (rows <= 0) return; + X += (size_t)row0 * TF_K; + Y += (size_t)row0 * TF_YS; +#if TF_PRECOMPUTE_SUMS + XS += row0; +#endif +#endif const int lane = int(lanei), sg = int(sgi); const short qid = short(lane) >> 2; const short fm = (qid & 4) | ((short(lane) >> 1) & 3); @@ -187,8 +200,13 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #else const short fn = ((qid & 2) | (short(lane) & 1)) * 4; #endif +#if TF_ROW_BLOCKS + const int M = min(rows, 16); + const int tile = tf_tile(int(tg.x)); +#else const int M = mdims[0]; const int tile = tf_tile(int(tgx)); +#endif const int n0 = tile * NT; const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); diff --git a/zig/kernels/metal/core/shared_attention.metal b/zig/kernels/metal/core/shared_attention.metal index 126aef9685c..11ae5ce8910 100644 --- a/zig/kernels/metal/core/shared_attention.metal +++ b/zig/kernels/metal/core/shared_attention.metal @@ -149,34 +149,6 @@ inline void tf_context_row(const device bfloat* Q, const device bfloat* K, const tf_context_row(Q, K, V, (const device uchar*)nullptr, O, 0, head, args.keys, 0, as_type(args.scale_bits), lane); } -// Position is supplied per row. Caller-supplied inverse frequencies preserve the model's RoPE convention. -[[kernel]] void tf_lane_rope(const device bfloat* X [[buffer(0)]], const device uint* positions [[buffer(1)]], - const device float* frequencies [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], - device bfloat* Y [[buffer(4)]], uint index [[thread_position_in_grid]]) { - if ((size_t)index >= (size_t)args.rows * args.heads * TF_D) return; - const uint d = index % TF_D; -#if TF_RD > 0 - if (d >= TF_RD) { Y[index] = X[index]; return; } -#if TF_INTERLEAVED - const uint pair = d / 2; - const uint first = pair * 2, second = first + 1; - const bool high = (d & 1) != 0; -#else - const uint pair = d % (TF_RD / 2); - const uint first = pair, second = first + TF_RD / 2; - const bool high = d >= TF_RD / 2; -#endif - const uint row = index / (args.heads * TF_D); - const size_t base = index - d; - const float angle = float(positions[row]) * frequencies[pair]; - const float c = precise::cos(angle), s = precise::sin(angle); - const float x = float(X[base + first]), y = float(X[base + second]); - Y[index] = bfloat(high ? fma(x, s, y*c) : fma(x, c, -y*s)); -#else - Y[index] = X[index]; -#endif -} - // Only newly produced rows are written. Prompt and previous-round cache slots are never copied or rewritten. [[kernel]] void tf_lane_append_kv(const device bfloat* K [[buffer(0)]], const device bfloat* V [[buffer(1)]], constant uint* slots [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], diff --git a/zig/kernels/metal/qwen27/prompt_glue.metal b/zig/kernels/metal/qwen27/prompt_glue.metal deleted file mode 100644 index 3ff271a61ab..00000000000 --- a/zig/kernels/metal/qwen27/prompt_glue.metal +++ /dev/null @@ -1,92 +0,0 @@ -// Prompt intermediates round where the declared BF16 graph rounds, rather than at the lane fusion's final store. -#include -using namespace metal; - -struct Q27PromptProduct { uint width, rows, mode, pad; }; -struct Q27PromptQk { uint rows, heads, width, stride, offset; float eps, scale; uint pad; }; -struct Q27PromptDecay { uint rows, heads, weight_flags, pad; }; - -inline float q27_prompt_probability(float value) { - const float magnitude = metal::abs(value); - const float denominator = precise::exp(magnitude) + 1.0f; - const float tail = 1.0f / denominator; - if (value < 0.0f) return tail; - return 1.0f - tail; -} - -inline float q27_prompt_softplus(float value) { - const float small = metal::exp(-metal::abs(value)); - const float shifted = 1.0f + small; - float logarithm = small; - if (shifted != 1.0f) { - const float difference = shifted - 1.0f; - const float ratio = metal::log(shifted) / difference; - logarithm = ratio * small; - } - return metal::max(value, 0.0f) + logarithm; -} - -kernel void q27_prompt_product(device const bfloat* gate [[buffer(0)]], - device const bfloat* factor [[buffer(1)]], - constant Q27PromptProduct& p [[buffer(2)]], - device bfloat* result [[buffer(3)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float x = float(gate[at]); - const float probability = q27_prompt_probability(x); - float activated; - if (p.mode == 0) activated = float(bfloat(x * float(bfloat(probability)))); - else if (p.mode == 1) activated = x * probability; - else activated = x / (1.0f + metal::exp(-x)); - result[at] = bfloat(activated * float(factor[at])); -} - -kernel void q27_prompt_silu(device const bfloat* input [[buffer(0)]], - constant Q27PromptProduct& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float value = float(input[at]); - const bfloat probability = bfloat(q27_prompt_probability(value)); - result[at] = bfloat(value * float(probability)); -} - -kernel void q27_prompt_decay(device const bfloat* a [[buffer(0)]], - device const bfloat* b [[buffer(1)]], - device const void* a_log [[buffer(2)]], - device const void* time_bias [[buffer(3)]], - constant Q27PromptDecay& p [[buffer(4)]], - device float* decay [[buffer(5)]], - device bfloat* beta [[buffer(6)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.heads) return; - const uint head = at % p.heads; - const bool wide_bias = (p.weight_flags & 4u) != 0; - const float bias = wide_bias ? ((device const float*)time_bias)[head] : float(((device const bfloat*)time_bias)[head]); - const float logarithm = (p.weight_flags & 2u) != 0 ? ((device const float*)a_log)[head] : float(((device const bfloat*)a_log)[head]); - const float sum = float(a[at]) + bias; - const float shifted = wide_bias ? sum : float(bfloat(sum)); - const float raw = q27_prompt_softplus(shifted); - const float smooth = wide_bias ? raw : float(bfloat(raw)); - const float speed = precise::exp(logarithm); - decay[at] = precise::exp(-speed * smooth); - beta[at] = bfloat(q27_prompt_probability(float(b[at]))); -} - -kernel void q27_prompt_qk(device const bfloat* activation [[buffer(0)]], - constant Q27PromptQk& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint2 group [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint begin = group.y * p.stride + p.offset + group.x * p.width + lane * 4; - float values[4]; - float square = 0.0f; - for (uint part = 0; part < 4; ++part) { - values[part] = float(activation[begin + part]); - square += values[part] * values[part]; - } - const float inverse = precise::rsqrt(simd_sum(square) / float(p.width) + p.eps); - const float scale = float(bfloat(p.scale)); - const uint output = (group.y * p.heads + group.x) * p.width + lane * 4; - for (uint part = 0; part < 4; ++part) result[output + part] = bfloat(scale * float(bfloat(values[part] * inverse))); -} diff --git a/zig/kernels/metal/qwen27/prompt_state.metal b/zig/kernels/metal/qwen27/prompt_state.metal deleted file mode 100644 index 1f8808136fc..00000000000 --- a/zig/kernels/metal/qwen27/prompt_state.metal +++ /dev/null @@ -1,82 +0,0 @@ -// Prompt state advances through the core tap order and one fixed per-position reduction topology. -#include -using namespace metal; - -struct Q27PromptConv { uint rows, channels, taps, streams; }; -struct Q27PromptRecur { uint rows, nk, nv, dv, streams, pad0, pad1, pad2; }; - -kernel void q27_prompt_conv(device const bfloat* projected [[buffer(0)]], - device const bfloat* history [[buffer(1)]], - device const bfloat* weight [[buffer(2)]], - constant Q27PromptConv& p [[buffer(3)]], - device bfloat* output [[buffer(4)]], - device bfloat* next_history [[buffer(5)]], - uint3 at [[thread_position_in_grid]]) { - if (at.x >= p.channels || at.y >= p.rows || at.z >= p.streams) return; - const uint keep = p.taps - 1; - const uint old_base = at.z * keep * p.channels; - const uint new_base = at.z * p.rows * p.channels; - float sum = 0.0f; - for (uint tap = 0; tap < p.taps; ++tap) { - const uint row = at.y + tap; - const float value = row < keep ? float(history[old_base + row * p.channels + at.x]) : float(projected[new_base + (row - keep) * p.channels + at.x]); - sum += value * float(weight[at.x * p.taps + tap]); - } - output[new_base + at.y * p.channels + at.x] = bfloat(sum); - if (at.y == 0) { - for (uint tail = 0; tail < keep; ++tail) { - const uint row = p.rows + tail; - next_history[old_base + tail * p.channels + at.x] = row < keep ? history[old_base + row * p.channels + at.x] : projected[new_base + (row - keep) * p.channels + at.x]; - } - } -} - -inline float q27_prompt_fold(thread float* partial) { - for (uint stride = 1; stride < 8; stride *= 2) { - for (uint at = 0; at < 8; at += 2 * stride) partial[at] = partial[at] + partial[at + stride]; - } - float total = partial[0]; - for (uint neighbor = 1; neighbor < 4; neighbor *= 2) total = total + simd_shuffle_xor(total, neighbor); - return total; -} - -kernel void q27_prompt_recur(device const bfloat* query [[buffer(0)]], - device const bfloat* key [[buffer(1)]], - device const bfloat* value [[buffer(2)]], - device const float* decay [[buffer(3)]], - device const bfloat* beta [[buffer(4)]], - device const float* committed [[buffer(5)]], - constant Q27PromptRecur& p [[buffer(6)]], - device bfloat* result [[buffer(7)]], - device float* next_state [[buffer(8)]], - uint3 at [[thread_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint stream = at.z / p.nv; - const uint head = at.z % p.nv; - const uint dimension = at.y * 8 + lane / 4; - const uint first = (lane % 4) * 32; - const uint key_head = head / (p.nv / p.nk); - const uint state_base = ((stream * p.nv + head) * p.dv + dimension) * 128 + first; - float state[32]; - for (uint channel = 0; channel < 32; ++channel) state[channel] = committed[state_base + channel]; - for (uint row = 0; row < p.rows; ++row) { - const uint key_base = ((stream * p.rows + row) * p.nk + key_head) * 128 + first; - const uint head_at = (stream * p.rows + row) * p.nv + head; - float partial[8]; - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] * decay[head_at]; - partial[channel / 4] += state[channel] * float(key[key_base + channel]); - } - const float remembered = q27_prompt_fold(partial); - const float correction = (float(value[head_at * p.dv + dimension]) - remembered) * float(beta[head_at]); - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] + float(key[key_base + channel]) * correction; - partial[channel / 4] += state[channel] * float(query[key_base + channel]); - } - const float output = q27_prompt_fold(partial); - if (lane % 4 == 0) result[head_at * p.dv + dimension] = bfloat(output); - } - for (uint channel = 0; channel < 32; ++channel) next_state[state_base + channel] = state[channel]; -} diff --git a/zig/qwen27_quant_metal.zig b/zig/qwen27_quant_metal.zig index 02a97ba26d5..9684fb4bba7 100644 --- a/zig/qwen27_quant_metal.zig +++ b/zig/qwen27_quant_metal.zig @@ -32,9 +32,11 @@ pub const Record = struct { released: usize = 0, live_bytes: usize = 0, fail_pipeline: bool = false, + tensor_units: bool = true, current: usize = 0, rows: u32 = 0, - launches: [256]struct { pipeline: usize, rows: u32 } = undefined, + launches: [256]struct { pipeline: usize, rows: u32, grid: Size, stride: i32 } = undefined, + sums_stride: i32 = 0, count: usize = 0, pub fn device(r: *Record) Device { return .{ .r = r }; @@ -65,8 +67,8 @@ pub const Buffer = struct { }; pub const Device = struct { r: *Record, - pub fn tensorUnits(_: Device) bool { - return true; + pub fn tensorUnits(d: Device) bool { + return d.r.tensor_units; } pub fn generation(_: Device) usize { return 17; @@ -111,10 +113,13 @@ pub const ComputeEncoder = struct { } pub fn setBuffer(_: ComputeEncoder, _: Buffer, _: usize, _: usize) void {} pub fn setValue(e: ComputeEncoder, v: anytype, at: usize) void { - if (@TypeOf(v) == [2]i32 and at == 3) e.r.rows = @intCast(v[0]); + if (@TypeOf(v) == [2]i32 and at == 3) { + e.r.rows = @intCast(v[0]); + e.r.sums_stride = v[1]; + } } - pub fn dispatchGroups(e: ComputeEncoder, _: Size, _: Size) void { - e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows }; + pub fn dispatchGroups(e: ComputeEncoder, grid: Size, _: Size) void { + e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows, .grid = grid, .stride = e.r.sums_stride }; e.r.count += 1; } pub fn dispatchThreads(_: ComputeEncoder, _: Size, _: Size) void {} diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 73d432c8705..8fcfa9a5bcf 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -132,3 +132,127 @@ test "qwen27 quant: preparation skips projections that cannot take a decode lane try t.expectEqual(@as(u32, 0), k.lanes.count()); try t.expectEqual(@as(usize, 0), r.libraries); } + +const wide_formats = [_]lane_projection.Format{ + .{ .bits = 8, .group = 64 }, .{ .bits = 6, .group = 32 }, .{ .bits = 5, .group = 64 }, + .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, +}; + +test "qwen27 quant: non-reg layouts admit whole prompt chunks with row-block source variants" { + for (wide_formats) |format| { + const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048, .row_blocks = true }; + try long.validate(); + var short = long; + short.rows = 16; + const text = try lane_projection.source(a, long); + defer a.free(text); + const old = try lane_projection.source(a, short); + defer a.free(old); + try t.expect(std.mem.indexOf(u8, old, "#define TF_ROW_BLOCKS 0\n") != null); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 1\n") != null); + for ([_]usize{ 0, 15, 17, 2064 }) |rows| { + short.rows = rows; + try t.expectError(error.InvalidSchedule, short.validate()); + } + } +} + +test "qwen27 quant: each non-reg whole prompt is one projection dispatch" { + for (wide_formats) |format| for ([_]u32{ 17, 1065, 2048 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const producer: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 2080 } else null; + try k.quant(r.encoder(), p, ref, producer, ref, ref, rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(rows, r.launches[0].rows); + try t.expectEqual(@as(usize, (rows + 15) / 16), r.launches[0].grid.height); + try t.expectEqual(@as(i32, if (producer != null) 2080 else @intCast((rows + 15) / 16 * 16)), r.launches[0].stride); + var entries = k.lanes.valueIterator(); + const entry = entries.next().?; + try t.expectEqual(@as(u32, 2048), entry.max_rows); + if (entry.sums_buffer) |b| try t.expectEqual(@as(usize, 2048 * (5120 / format.group) * 4), b.length()); + // A decode or verify row still uses the short key and its original 16-row steps. + k.prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 17); + try t.expectEqual(@as(usize, 2), r.count); + try t.expectEqual(@as(u32, 16), r.launches[0].rows); + try t.expectEqual(@as(u32, 1), r.launches[1].rows); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectEqual(@as(u32, 2), k.lanes.count()); + k.prompt = true; + k.long_prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 2048); + try t.expectEqual(@as(usize, 128), r.count); + for (r.launches[0..r.count]) |call| { + try t.expectEqual(@as(u32, 16), call.rows); + try t.expectEqual(@as(usize, 1), call.grid.height); + } + }; +} + +test "qwen27 quant: cooperative and older-chip projections keep their 16-row limits" { + var r: mtl.Record = .{}; + var layout = lane_projection.Layout{ .n = 64, .k = 512, .format = .{ .bits = 4, .group = 64 }, .cooperative = true }; + const coop = try lane_projection.Projection.init(a, r.device(), layout); + defer coop.deinit(); + try t.expectEqual(@as(u32, 16), coop.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try coop.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, coop.encode(r.encoder(), ref, ref, ref, ref, 17)); + layout.rows = 2048; + try t.expectError(error.InvalidSchedule, layout.validate()); + layout.cooperative = false; + layout.format = .{ .bits = 6, .group = 32 }; + r.tensor_units = false; + const old = try lane_projection.Projection.init(a, r.device(), layout); + defer old.deinit(); + try t.expectEqual(@as(u32, 16), old.max_rows); + try t.expectError(error.InvalidRows, old.encode(r.encoder(), ref, ref, ref, ref, 17)); + const over_budget = lane_projection.Layout{ .n = 2048, .k = 5120, .format = .{ .bits = 8, .group = 128 }, .sk = 8 }; + try t.expectError(error.ThreadgroupMemoryLimit, over_budget.validate()); +} + +test "qwen27 quant: default non-reg callers retain scalar source and 16-row capacity" { + for (wide_formats) |format| { + var r: mtl.Record = .{}; + const layout = lane_projection.Layout{ .n = 64, .k = 5120, .format = format }; + const text = try lane_projection.source(a, layout); + defer a.free(text); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 0\n") != null); + const p = try lane_projection.Projection.init(a, r.device(), layout); + defer p.deinit(); + try t.expectEqual(@as(u32, 16), p.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try p.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, p.encode(r.encoder(), ref, ref, ref, ref, 17)); + } +} + +test "qwen27 quant: non-reg prompt key opts in only above 16 rows" { + for (wide_formats) |format| for ([_]u32{ 16, 17 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const sums: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 32 } else null; + try k.quant(r.encoder(), p, ref, sums, ref, ref, rows); + var entries = k.lanes.iterator(); + const entry = entries.next().?; + try t.expectEqual(rows > 16, entry.key_ptr.long); + try t.expectEqual(@as(u32, if (rows > 16) 2048 else 16), entry.value_ptr.max_rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(@as(usize, if (rows > 16) 2 else 1), r.launches[0].grid.height); + }; +} diff --git a/zig/runtime_sources/core.zig b/zig/runtime_sources/core.zig index dbc6cb491fd..6472609c1bc 100644 --- a/zig/runtime_sources/core.zig +++ b/zig/runtime_sources/core.zig @@ -46,10 +46,10 @@ fn projections(s: *Set) !void { try s.fragments("core-lane", "prepared", try lane.source(s.a, l)); }; const attention = @import("../src/core/shared_attention.zig"); - for (std.enums.values(attention.Rope)) |rope| for (std.enums.values(attention.Output)) |out| { - const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .rotary_dim = 64, .capacity = 4096, .rope = rope, .output = out }); + for (std.enums.values(attention.Output)) |out| { + const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .capacity = 4096, .output = out }); try s.add("core", "shared-attention", text, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); - }; + } const route = @import("../src/core/moe_route.zig"); for ([_]route.Shape{ .{ .hidden = 4096, .experts = 288, .topk = 8 }, .{ .hidden = 4096, .experts = 128, .topk = 8, .out_bf16 = true }, .{ .hidden = 6144, .experts = 128, .topk = 4 } }) |shape| try s.raw("core", "route", try route.source(s.a, shape)); diff --git a/zig/runtime_sources/qwen27.zig b/zig/runtime_sources/qwen27.zig index 83327917957..1cb40bfe3d3 100644 --- a/zig/runtime_sources/qwen27.zig +++ b/zig/runtime_sources/qwen27.zig @@ -3,7 +3,7 @@ const std = @import("std"); const ks = @import("qwen_runtime_sources"); const Set = @import("set.zig").Set; pub fn add(s: *Set) !void { - inline for (.{ "glue", "copy", "moe", "prompt_glue", "prompt_state" }) |name| + inline for (.{ "glue", "copy", "moe" }) |name| try s.raw("qwen27", name, @field(ks, name)); try s.raw("qwen27", "moe-rows", @import("../src/families/qwen27/moe_gpu.zig").rows_source); const quant = @import("../src/families/qwen27/quant_gpu.zig"); diff --git a/zig/src/core/contract/experts_bf16_records.zig b/zig/src/core/contract/experts_bf16_records.zig new file mode 100644 index 00000000000..7b4c4bf7784 --- /dev/null +++ b/zig/src/core/contract/experts_bf16_records.zig @@ -0,0 +1,37 @@ +//! The bf16 experts' CUDA launches (experts_bf16.cu): decode and prompt rows by epilogue, and the repack. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_up", .symbol = "tf_bf16_experts_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .symbol = "tf_bf16_experts_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .symbol = "tf_bf16_experts_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .symbol = "tf_bf16_experts_prompt_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .symbol = "tf_bf16_experts_prompt_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .symbol = "tf_bf16_experts_prompt_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_pack", .symbol = "tf_bf16_experts_pack", .image = "experts_bf16", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, +}; + +/// nvcc 13.2's prompt forms on sm_120 and sm_121 (164 registers where 13.0 takes 162); the rest compile as 13.0's. +pub const cuda_132 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, +}; + +/// nvcc 13.3's decode forms on sm_120 and sm_121 (70 registers where 13.0 and 13.2 take 72). +pub const cuda_133 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_up", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, +}; + +test "bf16 experts CUDA: every launch is in the live registry and fits every SM" { + const t = @import("std").testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(@import("std").mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/hip_records_item.zig b/zig/src/core/contract/hip_records_item.zig index e7d154bb843..265209c83c7 100644 --- a/zig/src/core/contract/hip_records_item.zig +++ b/zig/src/core/contract/hip_records_item.zig @@ -18,8 +18,8 @@ pub const records = [1]Record{ .launch = itemTile(4), .usage = .{ .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, }, }, diff --git a/zig/src/core/contract/lane_q8_records.zig b/zig/src/core/contract/lane_q8_records.zig index 16bdd8bd2c1..0e734861227 100644 --- a/zig/src/core/contract/lane_q8_records.zig +++ b/zig/src/core/contract/lane_q8_records.zig @@ -1,4 +1,4 @@ -//! M5 byte-lane sources retain their launch envelopes before and after removing code staging. +//! Byte-lane launch envelopes; the standalone sums compiled for M1/M2 even when tensor ops cannot. const M = @import("limits.zig").MetalLaunch; pub const Bad = enum { none, group, stride }; pub const Source = struct { name: []const u8, n: usize, k: usize, sk: usize, group: usize = 64, sums: bool = false, pf: usize = 1, f32_out: bool = false, bad: Bad = .none }; @@ -35,7 +35,24 @@ pub const records = blk: { const groups: usize = @max(@as(usize, s.sk) -| 1, @as(usize, 1)); const threads: usize = if (sum) 256 else 32 * s.sk; const memory: usize = if (sum) 0 else groups * 16 * 32 * 4; - out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .offline = false }; + out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .m1 = if (sum) 1024 else 0, .m2 = if (sum) 1024 else 0, .offline = sum }; } break :blk out; }; + +test "byte lane records: standalone sums have measured offline limits; tensor ops stay unmeasured" { + const std = @import("std"); + var count: usize = 0; + for (records) |r| { + const sum = std.mem.endsWith(u8, r.name, "tf_lane_sums"); + try std.testing.expectEqual(sum, r.offline); + try std.testing.expectEqual(@as(u32, if (sum) 1024 else 0), r.m1); + try std.testing.expectEqual(r.m1, r.m2); + if (sum) { + count += 1; + try std.testing.expectEqual(@as(u32, 256), r.threads); + try std.testing.expectEqual(@as(u32, 0), r.tg_memory); + } + } + try std.testing.expectEqual(@as(usize, 4), count); +} diff --git a/zig/src/core/contract/registry_joins.zig b/zig/src/core/contract/registry_joins.zig index 615190a1678..782618b89bd 100644 --- a/zig/src/core/contract/registry_joins.zig +++ b/zig/src/core/contract/registry_joins.zig @@ -13,6 +13,7 @@ pub const cuda_after = @import("registry_media.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("nvfp4_experts32_records.zig").cuda ++ + @import("experts_bf16_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @@ -30,6 +31,7 @@ pub const nvcc_132_after = @import("registry_learning.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ + @import("experts_bf16_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132; pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @@ -37,7 +39,8 @@ pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133 ++ @import("minimax_cuda_records.zig").cuda_133 ++ @import("nvfp4_experts32_records.zig").cuda_133 ++ - @import("paging_pending.zig").cuda_133; + @import("paging_pending.zig").cuda_133 ++ + @import("experts_bf16_records.zig").cuda_133; pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("exl3_prompt_pending.zig").records ++ @import("cuda_lane_pending.zig").records ++ diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index 743c4b1a8cf..43979dd1ab3 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -54,4 +54,5 @@ test { _ = @import("exl3_grouped_records.zig"); _ = @import("proposal_sampling_records.zig"); _ = @import("pocket_records.zig"); + _ = @import("experts_bf16_records.zig"); } diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e2ff1527b18..8ea652757ec 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -116,6 +116,8 @@ pub const FrameSink = vision_types.FrameSink; pub const PreparedVideo = vision_types.PreparedVideo; pub const VideoOffer = vision_types.VideoOffer; pub const VisionOpen = vision_types.VisionOpen; +pub const Size = vision_types.Size; +pub const ocr = @import("engine_ocr.zig"); pub const SpeechRequest = @import("speech_api.zig").SpeechRequest; pub const SpeechResult = @import("speech_api.zig").SpeechResult; @@ -223,6 +225,7 @@ pub const Info = struct { pictures: bool = false, /// Image input (--vision); null: text only. vision: ?Vision = null, + ocr: ?ocr.Ocr = null, // OCR reader for /v1/ocr; null: no OCR route /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, /// The learner reads conversation prompts at any length through taps. @@ -371,18 +374,12 @@ pub const Opened = struct { follow: ?*const fn (*anyopaque) anyerror!void = null, }; -pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; - /// A backend's words for its own refusals (a request it cannot serve); null: the error's name. pub const Explain = struct { ctx: ?*anyopaque = null, text: *const fn (ctx: ?*anyopaque, err: anyerror) ?[]const u8 }; -/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. -pub const MemorySource = struct { - ctx: ?*anyopaque = null, - read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, -}; - pub const HealthSource = @import("engine_status.zig").HealthSource; +pub const Memory = @import("engine_status.zig").Memory; +pub const MemorySource = @import("engine_status.zig").MemorySource; pub const Status = @import("engine_status.zig").Status; pub const SubmitError = error{ Closed, Busy, InvalidSpans, Stale, InvalidRequest, MemoryPressure }; // Stale: put-back has not reached the device diff --git a/zig/src/core/engine_api_test.zig b/zig/src/core/engine_api_test.zig index 74a4aadc51d..efadf7dc452 100644 --- a/zig/src/core/engine_api_test.zig +++ b/zig/src/core/engine_api_test.zig @@ -15,4 +15,5 @@ test { _ = @import("keepalive.zig"); _ = @import("originals_test.zig"); _ = @import("image_api_test.zig"); + _ = @import("engine_ocr.zig"); } diff --git a/zig/src/core/engine_ocr.zig b/zig/src/core/engine_ocr.zig new file mode 100644 index 00000000000..a00315b2d50 --- /dev/null +++ b/zig/src/core/engine_ocr.zig @@ -0,0 +1,48 @@ +//! An engine's OCR reader (``Info.ocr``): its prompt and the parse from generated text to Markdown blocks. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Size = @import("engine_vision.zig").Size; + +/// The HTML table subset readers and the server share (ocr_html.zig). +pub const html = @import("ocr_html.zig"); + +/// TL, TR, BR, BL of the text in its upright reading frame, in original-image pixels (rotated text keeps its order). +pub const Quad = [4][2]f32; + +/// What a block holds. +pub const Kind = enum { + text, + title, + section_header, + list, + table, + formula, + code, + caption, + footnote, + page_header, + page_footer, + picture, +}; + +/// One block of a page in reading order. +pub const OcrBlock = struct { + content: []const u8, // normalized Markdown (docs/api.md OCR: Reader Markdown); empty only for .picture + quad: Quad, + kind: Kind, +}; + +/// The generated text does not follow the model's grammar. +pub const Error = error{Unreadable}; + +/// OCR an engine serves: the prompt after the page image, and the reader of the generated text. +pub const Ocr = struct { + /// Text parts that follow the image in the single user turn (may be empty). + prompt: []const u8, + /// Text to blocks in pixels of ``original``; ``frame`` is the tower input size the prepare step used. + parse: *const fn (a: Allocator, text: []const u8, original: Size, frame: Size) anyerror![]OcrBlock, +}; + +test { + _ = @import("ocr_html_test.zig"); +} diff --git a/zig/src/core/engine_status.zig b/zig/src/core/engine_status.zig index 3a024ac4243..0700fd90a23 100644 --- a/zig/src/core/engine_status.zig +++ b/zig/src/core/engine_status.zig @@ -12,6 +12,14 @@ pub const FsrSource = struct { read: *const fn (*anyopaque) ?Status.FsrCache, }; +pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; + +/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. +pub const MemorySource = struct { + ctx: ?*anyopaque = null, + read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, +}; + pub const Status = struct { healthy: bool = true, /// Engine-thread publication; all retained states, including long-message interior cuts and spare storage. diff --git a/zig/src/core/engine_vision.zig b/zig/src/core/engine_vision.zig index a4d4e819da5..285b697bbe0 100644 --- a/zig/src/core/engine_vision.zig +++ b/zig/src/core/engine_vision.zig @@ -13,8 +13,13 @@ pub const PreparedImage = struct { hash: u64, /// The image as tower inputs in order, each image tokens between `before` and `after`; empty: `pixels` alone. pieces: []const Piece = &.{}, + /// Pixel size of the tower input the image was resized to; set by every family with an OCR reader. + frame: ?Size = null, }; +/// An image's pixel size. +pub const Size = struct { width: u32, height: u32 }; + /// One tower input of an image, and the tokens the placeholder expands to around its rows. pub const Piece = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, before: []const u32 = &.{}, after: []const u32 = &.{} }; diff --git a/zig/src/core/kv_pool_checks.zig b/zig/src/core/kv_pool_checks.zig index fb0678a4a73..a938dcc392d 100644 --- a/zig/src/core/kv_pool_checks.zig +++ b/zig/src/core/kv_pool_checks.zig @@ -1,6 +1,7 @@ //! The page pool's checks beside engine_api: under the paged prompt store and a lane host (kv_pool is its own module). test { _ = @import("kv_pool_store_test.zig"); + _ = @import("kv_pool_parts_test.zig"); _ = @import("kv_pool_host_test.zig"); _ = @import("kv_pool_ledger_test.zig"); } diff --git a/zig/src/core/kv_pool_parts_test.zig b/zig/src/core/kv_pool_parts_test.zig new file mode 100644 index 00000000000..70a567f8997 --- /dev/null +++ b/zig/src/core/kv_pool_parts_test.zig @@ -0,0 +1,411 @@ +//! The paged store's checkpoint parts (C3, #2444) over a real pool: made all or nothing, held once, reclaimed, pinned. +const std = @import("std"); +const kv = @import("kv_pool"); +const radix = @import("engine_api").prompt_radix; +const Pool = kv.Pool; +const Lease = kv.Lease; +const testing = std.testing; +const expect = testing.expect; +const expectEqual = testing.expectEqual; +const expectEqualSlices = testing.expectEqualSlices; + +const page = 4; // tokens a page of the prefix group +const page_bytes = 8; + +/// A family whose device memory is one word a page: a copy writes the live state's word, a restore reads it back. +const Family = struct { + pool: *Pool, + mem: [64]u32 = @splat(0), + live: [2]u32 = .{ 0, 0 }, // each group's live state + restored: [2]?u32 = .{ null, null }, + fences: u32 = 0, + copies: u32 = 0, + rows_copied: u32 = 0, + fail_copy: bool = false, + fail_rows: bool = false, + queued: []const u32 = &.{}, // the pages the last copy wrote (queued work, until a fence) + held_at_fence: bool = true, // every fence saw the queued pages still held: nothing went back before it + hold_on_restore: usize = 0, // pages the restore takes from the pool (its reclaim runs) + + fn of(p: *anyopaque) *Family { + return @ptrCast(@alignCast(p)); + } + fn bytes(_: *anyopaque, _: u32, _: u32) u64 { + return 20; // three pages of 8 bytes + } + fn copy(p: *anyopaque, _: ?*anyopaque, group: u32, _: u32, dst: radix.Parts.Part) anyerror!void { + const f = of(p); + for (dst.ids) |id| f.mem[id] = f.live[group]; + f.queued = dst.ids; + if (f.fail_copy) return error.CopyFailed; // after its work is queued + f.copies += 1; + } + fn restore(p: *anyopaque, _: ?*anyopaque, group: u32, src: radix.Parts.Part) anyerror!void { + const f = of(p); + if (f.hold_on_restore > 0) { + var ids: [8]u32 = undefined; + if (!f.pool.hold(ids[0..f.hold_on_restore])) return error.NoPages; + for (ids[0..f.hold_on_restore]) |id| f.pool.release(id); + } + f.restored[group] = f.mem[src.ids[0]]; + } + fn copyRows(p: *anyopaque, _: ?*anyopaque, _: u32, src: u32, dst: u32, rows: u32) anyerror!void { + const f = of(p); + if (f.fail_rows) return error.CopyFailed; + f.mem[dst] = f.mem[src]; + f.rows_copied += rows; + } + fn fence(p: *anyopaque, _: ?*anyopaque) void { + const f = of(p); + f.fences += 1; + for (f.queued) |id| f.held_at_fence = f.held_at_fence and f.pool.holders(id) > 0; + f.queued = &.{}; + } + fn hooks(f: *Family) radix.PartHooks { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .copy = copy, .restore = restore, .copy_rows = copyRows, .fence = fence } }; + } + + fn snapBytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + return error.NoLegacy; + } + fn snapRestore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: radix.Saved) void {} + fn snapshots(f: *Family) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = snapBytes, .save = save, .restore = snapRestore, .drop = drop } }; + } +}; + +const copied_state = [_]radix.Parts.PartSpec{.{ .kind = .state, .make = .copied }}; +const window3 = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained, .rows = 3 }}; +const both = [_]radix.Parts.PartSpec{ .{ .kind = .window, .make = .retained, .rows = 3 }, .{ .kind = .state, .make = .copied } }; + +const rules: radix.Rules = .{ .page = page, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }; + +const World = struct { + pool: Pool, + f: Family = undefined, + store: radix.Store = undefined, + reclaims: u32 = 0, // the pool asked the store + + fn init(w: *World, specs: []const radix.Parts.PartSpec, pages: usize) !void { + w.pool = try Pool.init(testing.allocator, pages, page_bytes); + w.f = .{ .pool = &w.pool }; + w.store = try radix.Store.initGroups(testing.allocator, w.f.snapshots(), radix.pagesOf(Pool, &w.pool, page_bytes), rules, .{ .parts = specs, .legacy = false }, w.f.hooks()); + w.store.limit(.{ .pages = pages, .snaps = 8 }); + try w.pool.joinReclaim(.{ .ctx = &w.store, .call = reclaim }); + } + fn deinit(w: *World) void { + w.pool.leaveReclaim(&w.store); + w.store.deinit(); + expect(w.pool.idle() and w.pool.available() == w.pool.total()) catch @panic("pages left after close"); + w.pool.deinit(); + } + fn reclaim(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + @as(*World, @fieldParentPtr("store", s)).reclaims += 1; + return s.reclaim(want); + } + /// A stream's lease of `n` pages. + fn lease(w: *World, n: usize) !Lease { + var l: Lease = .{}; + try w.pool.promise(&l, n); + try w.pool.cover(&l, n, null); + return l; + } +}; + +const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14 }; + +test "a copied part: held from the pool at keep, restored and fenced at a resume, back at eviction" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 77; + var path: [2]u32 = undefined; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(kept.held); + try expectEqual(@as(usize, 5), w.store.tree.heldPages()); // two path pages, three part pages + w.pool.finish(&l); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 8), plan.from); + try expectEqual(@as(?u32, 77), w.f.restored[0]); + try expect(w.f.fences >= 2); // the keep's copies, then the restore's + for (adopt.items) |id| w.pool.release(id); + const before = w.pool.available(); + w.store.limit(.{ .pages = 16, .snaps = 0 }); // the entry goes, its parts with it + try expectEqual(before + 5, w.pool.available()); // the part pages and the path the entry alone kept +} + +test "a retained window with other rows a page: whole pages shared, the partial one copied, mapped at their logical index" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); // the prefix group + var win = try w.lease(3); // the window: rows 0..8 on 3-row pages, the last one partial (2 rows) + for (win.ids.items, 0..) |id, i| w.f.mem[id] = @intCast(100 + i); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{.{ .base = 1, .ids = win.ids.items[1..3] }}; // its window omits page 0 + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + try expectEqual(@as(u32, 2), w.f.rows_copied); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[1])); // the stream and the store + try expectEqual(@as(u32, 1), w.pool.holders(win.ids.items[2])); // the stream's partial page: never the store's + const last_page = win.ids.items[2]; + w.f.mem[last_page] = 999; // the stream writes on: the checkpoint does not change + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 1), plan.parts.len); + try expectEqual(@as(u32, 1), plan.parts[0].base); + try expectEqual(@as(usize, 1), plan.parts[0].ids.len); // page 1: whole + try expectEqual(@as(usize, 1), plan.copies.len); + try expectEqual(@as(u32, 2), plan.copies[0].index); // logical page 2, rows 0..2 + try expectEqual(@as(u32, 2), plan.copies[0].rows); + try expectEqual(@as(u32, 102), w.f.mem[plan.copies[0].src]); // the rows as they were at the keep + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + const fences = w.f.fences; + w.store.finish(null, plan.copies); + try expectEqual(fences + 1, w.f.fences); +} + +test "a part copy that fails publishes nothing: the path stays, the part pages go back after a fence" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + for ([_]bool{ true, false }) |rows_fail| { + w.f.fail_rows = rows_fail; + w.f.fail_copy = !rows_fail; + const fences = w.f.fences; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + try expect(!kept.held); + try expectEqual(@as(usize, 2), kept.shared); // the path is the store's + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // and nothing else + try expectEqual(fences + 1, w.f.fences); + try expect(w.f.held_at_fence); // the fence came before the pages went back + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); + } + try expectEqual(@as(u64, 2), w.store.tree.counts.failed); +} + +test "two entries whose windows overlap hold each shared window page once" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(3); + defer w.pool.finish(&l); + var win = try w.lease(4); // rows 0..12 on 3-row pages + defer w.pool.finish(&win); + var path: [3]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items[0..1], &.{.{ .base = 0, .ids = win.ids.items[0..2] }}, path[0..1]).held); + try expect(w.store.keepParts(.of(&prompt), 12, null, l.ids.items[0..3], &.{.{ .base = 0, .ids = win.ids.items[0..4] }}, path[0..3]).held); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[0])); // named by both entries, held once by the store +} + +test "reclaim frees a copied part while its path and retained part are a stream's; the stream keeps its pages" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items[0..3] }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const before = w.pool.available(); + try expect(w.store.reclaim(before + 3)); // the copied part's 3 pages (and the partial copy) are store-only + try expect(w.pool.available() >= before + 3); + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (l.ids.items) |id| try expect(w.pool.holders(id) >= 1); // the stream's pages stay + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); +} + +test "a restore that takes pool pages runs reclaim, which never evicts the entry being restored" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 8); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 5; + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); + w.f.hold_on_restore = w.pool.available() + 1; // one more than free: the pool asks the store + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 0), plan.from); // the restore could not get its pages: a prefill from the start + try expectEqual(@as(?u32, null), w.f.restored[0]); + try expectEqual(@as(usize, 0), adopt.items.len); +} + +test "parts need hooks, a pool that holds pages, and rows for a retained group" { + var pool = try Pool.init(testing.allocator, 4, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + const pages = radix.pagesOf(Pool, &pool, page_bytes); + try testing.expectError(error.PartsNeedHooks, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &copied_state, .legacy = false }, null)); + var no_hold = pages; + var vt = pages.vtable.*; + vt.hold = null; + no_hold.vtable = &vt; + try testing.expectError(error.PartsNeedHold, radix.Store.initGroups(testing.allocator, f.snapshots(), no_hold, rules, .{ .parts = &copied_state, .legacy = false }, f.hooks())); + const zero = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained }}; + try testing.expectError(error.RetainedRows, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &zero, .legacy = false }, f.hooks())); +} + +test "an allocation failing at each step of a keep with parts: before the path, nothing changes; after it, the path stays alone" { + var k: usize = 0; + while (true) : (k += 1) { + var pool = try Pool.init(testing.allocator, 32, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + var fa = std.testing.FailingAllocator.init(testing.allocator, .{}); + var store = try radix.Store.initGroups(fa.allocator(), f.snapshots(), radix.pagesOf(Pool, &pool, page_bytes), rules, .{ .parts = &both, .legacy = false }, f.hooks()); + store.limit(.{ .pages = 32, .snaps = 8 }); + var l: Lease = .{}; + try pool.promise(&l, 2); + try pool.cover(&l, 2, null); + var win: Lease = .{}; + try pool.promise(&win, 3); + try pool.cover(&win, 3, null); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + fa.fail_index = fa.alloc_index + k; + const kept = store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + const failed = fa.has_induced_failure; + if (failed) { + try expect(!kept.held); + try expectEqual(@as(usize, 0), store.tree.snaps); + const path_pages: usize = if (kept.shared == 2) 2 else 0; // after the insert the path is the store's + try expectEqual(path_pages, store.tree.heldPages()); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), pool.holders(id)); // no part outlived the failure + } else try expect(kept.held); + pool.finish(&l); + pool.finish(&win); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); + if (!failed) break; + } + try expect(k > 3); // the insert, the parts, the copies and the prompt each failed once +} + +test "a retained group without its pages through the mark is refused, not counted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 16); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var path: [2]u32 = undefined; + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 3, .ids = &.{} }}, &path).held); + try expectEqual(@as(usize, 0), w.store.tree.heldPages()); +} + +test "a keep whose parts need more than the free pages reclaims, but never the path it is building" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 6); + defer w.deinit(); + var l = try w.lease(2); + var path: [2]u32 = undefined; + w.f.fail_copy = true; // a failed checkpoint leaves its path alone in the tree + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // the path is now the store's alone: idle, reclaimable + var other = try w.lease(2); // 6 pages: 2 the path's, 2 this stream's, 2 free; the part needs 3 + defer w.pool.finish(&other); + w.f.fail_copy = false; + const asked = w.reclaims; + const kept = w.store.keepParts(.of(&prompt), 8, null, other.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(w.reclaims > asked); // the part's hold ran the reclaim chain + try expect(!kept.held); // the only reclaimable pages were the path being built: refused, not pruned + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // the path is still the tree's + try expectEqual(@as(usize, 2), kept.shared); +} + +test "a plan's retained ids are its own: an eviction between begin and adoption leaves them valid" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const whole = [_]u32{ win.ids.items[0], win.ids.items[1] }; + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + w.store.limit(.{ .pages = 32, .snaps = 0 }); // the entry goes before the stream maps anything + try expectEqualSlices(u32, &whole, plan.parts[0].ids); + for (plan.parts[0].ids) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); // the stream's holds alone + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + for (adopt.items) |id| w.pool.release(id); +} + +test "a begin whose hand-out allocation fails gives back the pages it adopted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = win.ids.items }}, &path).held); + w.pool.finish(&l); + w.pool.finish(&win); + var k: usize = 0; + while (true) : (k += 1) { + var fa = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = k }); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const census = w.pool.census(); + const plan = w.store.begin(fa.allocator(), .of(&prompt), 14, &.{}, null, &adopt) catch { + try expectEqual(@as(usize, 0), adopt.items.len); + try expectEqual(census, w.pool.census()); // every hold this begin took went back + continue; + }; + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + plan.deinit(fa.allocator()); + break; + } + try expect(k >= 3); // the marks, the parts, their ids and the copies each failed once +} + +test "a mark whose path and parts pass the page limit is refused when planned, so it never evicts a checkpoint first" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + w.store.limit(.{ .pages = 5, .snaps = 8 }); + var l = try w.lease(1); + var path: [1]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // one prefix page and three copied: 4 of 5 + const other = [_]u32{ 9, 9, 9, 9, 8, 8, 8, 8, 7, 7, 7, 7, 6, 6 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&other), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 0), plan.marks.len); // its mark at 12 needs 3 path pages and 3 part pages: 6 of 5 + try expectEqual(@as(usize, 1), w.store.tree.snaps); // the kept checkpoint was not evicted for it +} diff --git a/zig/src/core/kv_pool_store_test.zig b/zig/src/core/kv_pool_store_test.zig index 7260b1836a7..7f3699267dc 100644 --- a/zig/src/core/kv_pool_store_test.zig +++ b/zig/src/core/kv_pool_store_test.zig @@ -79,7 +79,7 @@ const World = struct { fn admit(w: *World, a: Allocator, prompt: []const u32, history: u32, max_new: u32) !void { var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(w.gpa); - const plan = try w.store.begin(a, prompt, history, &.{}, null, &adopt); + const plan = try w.store.begin(a, .of(prompt), history, &.{}, null, &adopt); var l: Lease = .{}; errdefer w.pool.finish(&l); try w.pool.adopt(&l, adopt.items); @@ -89,7 +89,7 @@ const World = struct { for (plan.marks) |m| { var path: [64]u32 = undefined; const n = m / page; - const kept = w.store.keep(prompt, m, null, l.ids.items[0..n], path[0..n]); + const kept = w.store.keep(.of(prompt), m, null, l.ids.items[0..n], path[0..n]); for (path[0..kept.shared], 0..) |id, i| w.pool.share(&l, i, id); } try w.streams.append(w.gpa, l); @@ -212,3 +212,77 @@ test "page accounting: random admits, keeps, resumes, finishes, cancels and fail } try expect(refused > 0 and failed > 0 and adopted > 0); // the runs met full pools, failed uploads and resumes } + +/// A legacy family whose save holds pool pages for its copy (as the CUDA ring families' blobs do), one more than free. +const HoldingFamily = struct { + pool: *Pool, + fail: bool = false, + live: usize = 0, + + const State = struct { ids: [8]u32, n: usize }; + + fn of(ptr: *anyopaque) *HoldingFamily { + return @ptrCast(@alignCast(ptr)); + } + fn snapshots(f: *HoldingFamily) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = drop } }; + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(ptr: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + const f = of(ptr); + if (f.fail) return error.CopyFailed; + const st = try std.testing.allocator.create(State); + st.n = f.pool.available() + 1; // the pool must reclaim from the store to give them + if (!f.pool.hold(st.ids[0..st.n])) { + std.testing.allocator.destroy(st); + return error.NoCachePages; + } + f.live += 1; + return st; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(ptr: *anyopaque, saved: radix.Saved) void { + const f = of(ptr); + const st: *State = @ptrCast(@alignCast(saved)); + for (st.ids[0..st.n]) |id| f.pool.release(id); + std.testing.allocator.destroy(st); + f.live -= 1; + } +}; + +fn reclaimStore(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + return s.reclaim(want); +} + +test "a legacy save that reclaims pool pages while keeping onto a leaf with no entry never prunes that leaf" { + var pool = try Pool.init(std.testing.allocator, 6, 1); + defer pool.deinit(); + var f: HoldingFamily = .{ .pool = &pool }; + var store = try radix.Store.init(std.testing.allocator, f.snapshots(), radix.pagesOf(Pool, &pool, 1), .{ .page = page, .min_gap = 1, .min_prompt = 0 }); + store.limit(.{ .pages = 6, .snaps = 4 }); + try pool.joinReclaim(.{ .ctx = &store, .call = reclaimStore }); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + var path: [2]u32 = undefined; + var first: Lease = .{}; + try pool.promise(&first, 2); + try pool.cover(&first, 2, null); + f.fail = true; // a failed save leaves the path in the tree, a leaf with no entry + try expect(!store.keep(.of(&prompt), 8, null, first.ids.items, &path).held); + pool.finish(&first); // the leaf's pages are the store's alone: the only ones reclaim could free + try expectEqual(@as(usize, 2), store.tree.heldPages()); + var second: Lease = .{}; + try pool.promise(&second, 2); + try pool.cover(&second, 2, null); + f.fail = false; + const kept = store.keep(.of(&prompt), 8, null, second.ids.items, &path); // onto that leaf; its save reclaims + try expect(!kept.held); // the pinned leaf could not be taken, so the save found no page + try expectEqual(@as(usize, 2), store.tree.heldPages()); // the leaf survived the keep that was writing to it + try expectEqual(@as(usize, 0), f.live); + pool.finish(&second); + pool.leaveReclaim(&store); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); +} diff --git a/zig/src/core/kv_prompts.zig b/zig/src/core/kv_prompts.zig index 19518d71b40..9f5f696336d 100644 --- a/zig/src/core/kv_prompts.zig +++ b/zig/src/core/kv_prompts.zig @@ -5,6 +5,19 @@ const kv = @import("kv_pool.zig"); /// Where a pass resumes and the marks it keeps states at (`marks` the store's allocation, freed by `free`). pub const Plan = struct { from: u32 = 0, marks: []const u32 = &.{} }; +/// What a pass asks the store about (prompt_radix.Query as a family on the core pool gives it). +pub const Query = struct { + /// The words the store compares: the request's cache key, one word a model row; for text, the prompt. + key: []const u32, + /// The prompt's length; a key of another length reuses nothing. + prompt_len: u32, + + /// A query of `key` alone, as long as its prompt. + pub fn of(key: []const u32) Query { + return .{ .key = key, .prompt_len = @intCast(key.len) }; + } +}; + /// The store as a family sees it: a prompt pass begins, keeps at its marks, and reports once its prompt is in. pub const Prompts = struct { ctx: *anyopaque, @@ -12,20 +25,20 @@ pub const Prompts = struct { pub const VTable = struct { /// The deepest state `prompt` resumes: its pages into `adopt` (held once more), the rest restored into `owner`. - begin: *const fn (ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, + begin: *const fn (ctx: *anyopaque, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, /// At `at` the store takes the pages of `mine` it lacks and the state; `path`'s first n replace the stream's. - keep: *const fn (ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, + keep: *const fn (ctx: *anyopaque, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, /// One log line after a prompt pass: where it resumed and how many states it kept. report: *const fn (ctx: *anyopaque, prompt: usize, from: u32, kept: u64) void, free: *const fn (ctx: *anyopaque, marks: []const u32) void, }; - pub fn begin(p: Prompts, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - return p.vtable.begin(p.ctx, prompt, history_len, shared, owner, adopt); + pub fn begin(p: Prompts, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + return p.vtable.begin(p.ctx, q, history_len, shared, owner, adopt); } - pub fn keep(p: Prompts, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return p.vtable.keep(p.ctx, prompt, at, owner, mine, path); + pub fn keep(p: Prompts, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return p.vtable.keep(p.ctx, q, at, owner, mine, path); } pub fn report(p: Prompts, prompt: usize, from: u32, kept: u64) void { @@ -37,8 +50,8 @@ pub const Prompts = struct { } /// `begin`, its marks in `gpa` for a lane to own; a failure gives `adopt`'s pages back. - pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - const got = try p.begin(prompt, history_len, shared, owner, adopt); + pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + const got = try p.begin(q, history_len, shared, owner, adopt); defer p.free(got.marks); const marks = gpa.dupe(u32, got.marks) catch |e| { for (adopt.items) |id| pool.release(id); @@ -48,13 +61,13 @@ pub const Prompts = struct { } /// At mark `at` the store takes the stream's pages and state; pages it held replace `l`'s, the whole table written. - pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, prompt: []const u32, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { + pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, q: Query, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { const n = at / page_tokens; const path = try gpa.alloc(u32, n); defer gpa.free(path); const mine = try gpa.dupe(u32, l.ids.items); defer gpa.free(mine); - const shared = p.keep(prompt, @intCast(at), owner, mine[0..n], path); + const shared = p.keep(q, @intCast(at), owner, mine[0..n], path); if (std.mem.eql(u32, mine[0..shared], path[0..shared])) return shared == n; @memcpy(mine[0..shared], path[0..shared]); // the same tokens' rows: a resume from any row is exact try table.write(table.ctx, mine, 0); @@ -163,10 +176,10 @@ const Stub = struct { fn of(ctx: *anyopaque) *Stub { return @ptrCast(@alignCast(ctx)); } - fn begin(_: *anyopaque, _: []const u32, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { + fn begin(_: *anyopaque, _: Query, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { return .{ .from = 256, .marks = try testing.allocator.dupe(u32, &.{512}) }; } - fn keep(ctx: *anyopaque, _: []const u32, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { + fn keep(ctx: *anyopaque, _: Query, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { @memcpy(path, Stub.of(ctx).held[0..mine.len]); return mine.len; } @@ -203,7 +216,7 @@ test "keeping swaps the store's pages in and writes the whole table: a table of const p: Prompts = .{ .ctx = &stub, .vtable = &Stub.vtable }; var w: Written = .{}; var owner: u8 = 0; - const whole = try p.keepPages(testing.allocator, &pool, &l, &.{ 1, 2, 3 }, 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); + const whole = try p.keepPages(testing.allocator, &pool, &l, .of(&.{ 1, 2, 3 }), 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); try testing.expect(whole); try testing.expectEqual(@as(usize, 0), w.from); try testing.expectEqual(@as(usize, 4), w.cap); @@ -221,7 +234,7 @@ test "a plan's marks are the lane's own; a blob's pages come from the pool and g var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); var owner: u8 = 0; - const got = try p.plan(testing.allocator, &pool, &.{1}, 1, &.{}, &owner, &adopt); + const got = try p.plan(testing.allocator, &pool, .of(&.{1}), 1, &.{}, &owner, &adopt); defer testing.allocator.free(got.marks); try testing.expectEqual(@as(u32, 256), got.from); try testing.expectEqualSlices(u32, &.{512}, got.marks); diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index a12f9d52f2b..14e6029d370 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -1,4 +1,4 @@ -//! Affine lane projections for 1-16 bf16 rows on Metal. No family or model dependency. +//! Affine lane projections in independent 16-row blocks on Metal. No family or model dependency. const std = @import("std"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -36,8 +36,10 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, - /// Rows one reg encode may take, `reg_rows` to `max_reg_rows`: it sizes the sums scratch, never the kernel source. + /// Row capacity for reg layouts and opted-in row blocks; cooperative and scalar layouts take 16. rows: usize = reg_rows, + /// Non-reg prompt callers opt into independent 16-row blocks on the grid's y. + row_blocks: bool = false, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, @@ -50,7 +52,9 @@ pub const Layout = struct { if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; - if (l.rows != reg_rows and (!l.reg or l.rows < reg_rows or l.rows % 16 != 0 or l.rows > max_reg_rows)) return error.InvalidSchedule; + if (l.cooperative) { + if (l.rows != reg_rows) return error.InvalidSchedule; + } else if (l.rows < @as(usize, if (l.reg) reg_rows else 16) or l.rows % 16 != 0 or l.rows > max_reg_rows) return error.InvalidSchedule; const g = l.groupCut(); if (g[1] == 0 or g[0] >= l.k / l.format.group or g[1] > l.k / l.format.group - g[0]) return error.InvalidGroupCut; if (l.groups != null and l.output != .f32) return error.PartialOutputNeedsFloat; @@ -118,6 +122,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { } try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); + try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(l.row_blocks and !l.reg and !l.cooperative and l.rows > 16)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); @@ -238,13 +243,18 @@ pub const Projection = struct { sums_buffer: ?mtl.Buffer = null, threads_per_slice: usize = 32, reg: bool = false, - /// Rows one encode takes at most: the layout's rows for reg, 16 otherwise. + cooperative: bool = false, + /// Capacity is the reg or opted-in row count on tensor units, otherwise 16. max_rows: u32 = reg_rows, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); - if (!device.tensorUnits() and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; - const text = try source(a, l); + const units = device.tensorUnits(); + if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; + const capacity: usize = if (l.cooperative or !units or (!l.reg and !l.row_blocks)) 16 else l.rows; + var kernel_layout = l; + if (!units and !l.cooperative and !l.reg) kernel_layout.rows = capacity; + const text = try source(a, kernel_layout); defer a.free(text); const adapted = try frags.source(device, a, text); defer a.free(adapted); @@ -258,9 +268,9 @@ pub const Projection = struct { if (l.precompute_sums) { sums_pipe = try mtl.Pipeline.init(device, lib, "tf_lane_sums", false); errdefer sums_pipe.?.deinit(); - sums_buffer = try device.buffer(@as(usize, if (l.reg) l.rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + sums_buffer = try device.buffer(capacity * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .max_rows = @intCast(if (l.reg) l.rows else 16) }; + return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .cooperative = l.cooperative, .max_rows = @intCast(capacity) }; } pub fn deinit(p: Projection) void { @@ -306,7 +316,7 @@ pub const Projection = struct { e.setValue([2]i32{ @intCast(rows), @intCast(stride) }, 3); e.setBuffer(y.buf, y.off, 4); if (given) |g| e.setBuffer(g.ref.buf, g.ref.off, 5) else if (p.sums_buffer) |buffer| e.setBuffer(buffer, 0, 5); - e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.reg) blocks else 1, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); + e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.cooperative) 1 else blocks, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); if (given == null and p.sums_buffer != null) e.barrier(); } }; @@ -341,7 +351,7 @@ test "a reg layout's rows size its sums scratch, never its kernel source" { } var plain = base; plain.reg = false; - plain.rows = 256; + plain.rows = max_reg_rows + 16; try std.testing.expectError(error.InvalidSchedule, plain.validate()); } diff --git a/zig/src/core/ocr_html.zig b/zig/src/core/ocr_html.zig new file mode 100644 index 00000000000..39dab703d1a --- /dev/null +++ b/zig/src/core/ocr_html.zig @@ -0,0 +1,413 @@ +//! The HTML table subset OCR models write: cells and spans, and the table as GFM pipes, clean HTML or text rows. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// Most HTML bytes, grid rows and grid columns ``parse`` reads; a table past them is ``error.Malformed``. +pub const max_bytes = 64 * 1024; +pub const max_rows = 512; +pub const max_cols = 64; + +/// One cell at its top-left grid position, covering ``rows`` x ``cols``; ``text`` holds '\n' where the HTML had
. +pub const Cell = struct { text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool }; + +/// A table's grid and its cells in row-major order; ``spanned``: some cell covers more than one position. +pub const Table = struct { + rows: u32, + cols: u32, + cells: []Cell, + spanned: bool, + + /// Frees what ``parse`` allocated in ``a``. + pub fn deinit(t: Table, a: Allocator) void { + for (t.cells) |c| a.free(c.text); + a.free(t.cells); + } +}; + +pub const ParseError = error{ Malformed, OutOfMemory }; + +/// The one
in ``html``, with only whitespace and comments around it; cells and texts allocated in ``a``. +pub fn parse(a: Allocator, html: []const u8) ParseError!Table { + if (html.len > max_bytes) return error.Malformed; + var p: Parser = .{ .a = a, .src = html }; + defer p.text.deinit(a); + errdefer { + for (p.cells.items) |c| a.free(c.text); + p.cells.deinit(a); + } + try p.run(); + var cols: u32 = 0; + var spanned = false; + for (p.cells.items) |*c| { + c.rows = @min(c.rows, p.rows - c.row); // a rowspan past the last row ends there, as in HTML + cols = @max(cols, c.col + c.cols); + spanned = spanned or c.rows > 1 or c.cols > 1; + } + if (cols == 0) return error.Malformed; + return .{ .rows = p.rows, .cols = cols, .cells = try p.cells.toOwnedSlice(a), .spanned = spanned }; +} + +/// The table as a GFM pipe table: first row the header, cell text escaped to read literally (``cellText``); spans lost. +pub fn pipe(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + var k: usize = 0; + for (0..t.rows) |r| { + for (0..t.cols) |c| { + try out.appendSlice(a, if (c == 0) "| " else " | "); + try cellText(a, &out, if (topLeft(t, &k, r, c)) |cell| cell.text else ""); + } + try out.appendSlice(a, " |"); + if (r == 0) { + try out.append(a, '\n'); + for (0..t.cols) |_| try out.appendSlice(a, "| --- "); + try out.append(a, '|'); + } + if (r + 1 < t.rows) try out.append(a, '\n'); + } + return out.toOwnedSlice(a); +} + +/// The table as one clean
: rows of
/, rowspan and colspan only above 1, text entity-escaped. +pub fn normalized(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + try out.appendSlice(a, ""); + var k: usize = 0; + for (0..t.rows) |r| { + try out.appendSlice(a, ""); + while (k < t.cells.len and t.cells[k].row == r) : (k += 1) { + const cell = t.cells[k]; + const name = if (cell.head) "th" else "td"; + try out.print(a, "<{s}", .{name}); + if (cell.rows > 1) try out.print(a, " rowspan=\"{d}\"", .{cell.rows}); + if (cell.cols > 1) try out.print(a, " colspan=\"{d}\"", .{cell.cols}); + try out.append(a, '>'); + for (cell.text) |ch| switch (ch) { + '&' => try out.appendSlice(a, "&"), + '<' => try out.appendSlice(a, "<"), + '>' => try out.appendSlice(a, ">"), + '"' => try out.appendSlice(a, """), + '\n' => try out.appendSlice(a, "
"), + else => try out.append(a, ch), + }; + try out.print(a, "", .{name}); + } + try out.appendSlice(a, ""); + } + try out.appendSlice(a, "
"); + return out.toOwnedSlice(a); +} + +/// A cell's text in a pipe row: math as written (only `|` escaped), markup escaped elsewhere, to read back literally. +fn cellText(a: Allocator, out: *std.ArrayList(u8), text: []const u8) Allocator.Error!void { + var no_close = false; // a `$` found no closer, so no later `$` in the cell can (markdown.zig's rule) + var i: usize = 0; + while (i < text.len) { + const ch = text[i]; + if (ch == '$') { + if (mathEnd(text, i, &no_close)) |end| { + try mathText(a, out, text[i..end]); + i = end; + continue; + } + // a `$` that opens no math is escaped, so a reader never pairs it with another + const lone: usize = if (i + 1 < text.len and text[i + 1] == '$') 2 else 1; + for (0..lone) |_| try out.appendSlice(a, "\\$"); + i += lone; + continue; + } + switch (ch) { + '\n' => try out.append(a, ' '), + '|', '`', '*', '[', ']' => try out.appendSlice(a, &[_]u8{ '\\', ch }), + // CommonMark reads a backslash as an escape only before ASCII punctuation; elsewhere it is literal + '\\' => try out.appendSlice(a, if (i + 1 < text.len and isPunct(text[i + 1])) "\\\\" else "\\"), + // `_` between two word characters can neither open nor close emphasis, so snake_case stays readable + '_' => try out.appendSlice(a, if (i > 0 and i + 1 < text.len and isWord(text[i - 1]) and isWord(text[i + 1])) "_" else "\\_"), + '<' => try out.appendSlice(a, if (opensTag(text[i + 1 ..])) "<" else "<"), + '&' => try out.appendSlice(a, if (opensEntity(text[i + 1 ..])) "&" else "&"), + else => try out.append(a, ch), + } + i += 1; + } +} + +/// Math as written, each `|` as `\|`: GFM's table step removes exactly that backslash, whatever precedes it. +fn mathText(a: Allocator, out: *std.ArrayList(u8), m: []const u8) Allocator.Error!void { + for (m) |ch| switch (ch) { + '|' => try out.appendSlice(a, "\\|"), + '\n' => try out.append(a, ' '), + else => try out.append(a, ch), + }; +} + +/// The end of the `$$…$$` or `$…$` math opening at ``s[i]``, by markdown.zig's ``mathSpan`` rule; null when none. +fn mathEnd(s: []const u8, i: usize, no_close: *bool) ?usize { + if (i + 1 < s.len and s[i + 1] == '$') return if (std.mem.indexOfPos(u8, s, i + 2, "$$")) |j| j + 2 else null; + if (i + 1 >= s.len or isSpace(s[i + 1]) or no_close.*) return null; + var j = i + 2; + while (j < s.len) : (j += 1) { + if (s[j] != '$' or isSpace(s[j - 1]) or escaped(s, j)) continue; + if (j + 1 < s.len and std.ascii.isDigit(s[j + 1])) continue; + return j + 1; + } + no_close.* = true; + return null; +} + +/// Whether ``s[i]`` follows an odd run of backslashes. +fn escaped(s: []const u8, i: usize) bool { + var n: usize = 0; + while (n < i and s[i - 1 - n] == '\\') n += 1; + return n % 2 == 1; +} + +fn isSpace(c: u8) bool { + return c == ' ' or c == '\t' or c == '\n' or c == '\r'; +} + +fn isPunct(c: u8) bool { + return c < 0x80 and std.ascii.isPrint(c) and !std.ascii.isAlphanumeric(c) and c != ' '; +} + +/// A byte of a word for emphasis rules: neither whitespace nor ASCII punctuation (UTF-8 letters count). +fn isWord(c: u8) bool { + return !isSpace(c) and !isPunct(c); +} + +/// Whether text after a literal `<` would start raw HTML in Markdown: a tag, an end tag, a comment or a declaration. +fn opensTag(rest: []const u8) bool { + return rest.len > 0 and (std.ascii.isAlphabetic(rest[0]) or std.mem.indexOfScalar(u8, "/!?", rest[0]) != null); +} + +/// Whether text after a literal `&` reads as an entity reference in Markdown (`name;`, `#digits;` or `#xhex;`). +fn opensEntity(rest: []const u8) bool { + var n: usize = 0; + while (n < rest.len and n < 33 and (std.ascii.isAlphanumeric(rest[n]) or (n == 0 and rest[n] == '#'))) n += 1; + return n > 0 and n < rest.len and rest[n] == ';'; +} + +/// The cell whose top-left is (r, c), walking ``k`` through the row-major cells; null: covered or empty. +fn topLeft(t: Table, k: *usize, r: usize, c: usize) ?Cell { + if (k.* >= t.cells.len or t.cells[k.*].row != r or t.cells[k.*].col != c) return null; + k.* += 1; + return t.cells[k.* - 1]; +} + +const Tag = enum { table, tr, td, th, br, other }; + +fn tagKind(name: []const u8) Tag { + inline for (.{ Tag.table, Tag.tr, Tag.td, Tag.th, Tag.br }) |t| { + if (std.ascii.eqlIgnoreCase(name, @tagName(t))) return t; + } + return .other; +} + +/// A span attribute's value as HTML reads it: leading digits, and 1 when there are none or they read 0. +fn spanValue(v: []const u8) u32 { + var i: usize = 0; + while (i < v.len and std.ascii.isWhitespace(v[i])) i += 1; + var n: u32 = 0; + while (i < v.len and std.ascii.isDigit(v[i])) : (i += 1) n = @min(n * 10 + (v[i] - '0'), 1_000_000); + return @max(n, 1); +} + +/// The code point a named or numeric entity body (between '&' and ';') stands for; null: not one this reader knows. +fn entityPoint(name: []const u8) ?u21 { + const named = .{ .{ "amp", '&' }, .{ "lt", '<' }, .{ "gt", '>' }, .{ "quot", '"' }, .{ "apos", '\'' } }; + if (std.mem.eql(u8, name, "nbsp")) return 0xa0; + inline for (named) |e| if (std.mem.eql(u8, name, e[0])) return e[1]; + if (name.len < 2 or name[0] != '#') return null; + const hex = name[1] == 'x' or name[1] == 'X'; + const base: u8 = if (hex) 16 else 10; + const digits = name[if (hex) 2 else 1..]; + if (digits.len == 0 or digits.len > 7) return null; + var v: u32 = 0; + for (digits) |d| v = v * base + (std.fmt.charToDigit(d, base) catch return null); + return if (v == 0 or v > 0x10ffff or (v >= 0xd800 and v <= 0xdfff)) 0xfffd else @intCast(v); +} + +/// Whitespace or one of ``ends``: where an attribute's name or unquoted value stops. +fn stops(ch: u8, ends: []const u8) bool { + return std.ascii.isWhitespace(ch) or std.mem.indexOfScalar(u8, ends, ch) != null; +} + +const Parser = struct { + a: Allocator, + src: []const u8, + i: usize = 0, + where: enum { before, inside, after } = .before, + cells: std.ArrayList(Cell) = .empty, + cell: ?Cell = null, // the open cell; its text gathers in ``text`` until it closes + text: std.ArrayList(u8) = .empty, + space: bool = false, // whitespace met in the open cell since its last character + row_open: bool = false, + rows: u32 = 0, // rows opened so far; the open row is the last + col: u32 = 0, // the open row's next column to try + until: [max_cols]u32 = @splat(0), // each column is covered by earlier cells through row until - 1 + + fn run(p: *Parser) ParseError!void { + while (p.i < p.src.len) { + const ch = p.src[p.i]; + const next = if (p.i + 1 < p.src.len) p.src[p.i + 1] else 0; + const after = if (p.i + 2 < p.src.len) p.src[p.i + 2] else 0; + if (ch == '<' and (next == '!' or next == '?')) { + try p.skipMarkup(); + } else if (ch == '<' and std.ascii.isAlphabetic(if (next == '/') after else next)) { + try p.tag(); + } else if (ch == '&') { + try p.entity(); + } else { + try p.char(ch); + p.i += 1; + } + } + if (p.where != .after) return error.Malformed; + } + + /// A comment, a doctype or a processing instruction: skipped whole. + fn skipMarkup(p: *Parser) ParseError!void { + if (std.mem.startsWith(u8, p.src[p.i..], "") orelse return error.Malformed; + p.i = end + 3; + } else { + const end = std.mem.indexOfScalarPos(u8, p.src, p.i, '>') orelse return error.Malformed; + p.i = end + 1; + } + } + + fn entity(p: *Parser) ParseError!void { + const rest = p.src[p.i..]; + const semi = std.mem.indexOfScalar(u8, rest[0..@min(rest.len, 12)], ';'); + const point = if (semi) |s| entityPoint(rest[1..s]) else null; + const cp = point orelse { + p.i += 1; + return p.char('&'); + }; + p.i += semi.? + 1; + if (cp == 0xa0) return p.char(' '); + var buf: [4]u8 = undefined; + const n = std.unicode.utf8Encode(cp, &buf) catch unreachable; + for (buf[0..n]) |b| try p.char(b); + } + + /// One text byte: whitespace runs fold to one space between words; outside a cell only whitespace may stand. + fn char(p: *Parser, ch: u8) ParseError!void { + const blank = std.ascii.isWhitespace(ch); + if (p.cell == null) { + if (!blank) return error.Malformed; + return; + } + if (blank) { + p.space = true; + return; + } + if (p.space and p.text.items.len > 0 and p.text.getLast() != '\n') try p.text.append(p.a, ' '); + p.space = false; + try p.text.append(p.a, ch); + } + + /// A start or end tag at ``i``: its name, its span attributes in any quoting, then its effect. + fn tag(p: *Parser) ParseError!void { + const src = p.src; + var i = p.i + 1; + const close = src[i] == '/'; + if (close) i += 1; + const start = i; + while (i < src.len and (std.ascii.isAlphanumeric(src[i]) or std.mem.indexOfScalar(u8, "-_:.", src[i]) != null)) i += 1; + const kind = tagKind(src[start..i]); + var span = [2]u32{ 1, 1 }; // rowspan, colspan + while (true) { + while (i < src.len and (std.ascii.isWhitespace(src[i]) or src[i] == '/')) i += 1; + if (i >= src.len) return error.Malformed; + if (src[i] == '>') break; + const name_start = i; + while (i < src.len and !stops(src[i], "=>/")) i += 1; + const name = src[name_start..i]; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len or src[i] != '=') continue; + i += 1; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len) return error.Malformed; + var value: []const u8 = undefined; + if (src[i] == '"' or src[i] == '\'') { + const end = std.mem.indexOfScalarPos(u8, src, i + 1, src[i]) orelse return error.Malformed; + value = src[i + 1 .. end]; + i = end + 1; + } else { + const value_start = i; + while (i < src.len and !stops(src[i], ">")) i += 1; + value = src[value_start..i]; + } + if (std.ascii.eqlIgnoreCase(name, "rowspan")) span[0] = spanValue(value); + if (std.ascii.eqlIgnoreCase(name, "colspan")) span[1] = spanValue(value); + } + p.i = i + 1; + try p.apply(kind, close, span); + } + + fn apply(p: *Parser, kind: Tag, close: bool, span: [2]u32) ParseError!void { + switch (p.where) { + .before => { + if (kind != .table or close) return error.Malformed; + p.where = .inside; + return; + }, + .after => return error.Malformed, + .inside => {}, + } + switch (kind) { + .table => { + if (!close) return error.Malformed; // a table inside a cell + try p.endCell(); + p.row_open = false; + p.where = .after; + }, + .tr => { + try p.endCell(); + p.row_open = false; + if (!close) try p.openRow(); + }, + .td, .th => { + try p.endCell(); + if (!close) try p.openCell(kind == .th, span); + }, + .br => if (p.cell != null) { + p.space = false; + try p.text.append(p.a, '\n'); + }, + .other => {}, + } + } + + fn openRow(p: *Parser) ParseError!void { + if (p.rows == max_rows) return error.Malformed; + p.rows += 1; + p.row_open = true; + p.col = 0; + } + + /// A cell at the open row's first free column; its spans clipped to the grid and to columns earlier rows cover. + fn openCell(p: *Parser, head: bool, span: [2]u32) ParseError!void { + if (!p.row_open) return error.Malformed; + const row = p.rows - 1; + while (p.col < max_cols and p.until[p.col] > row) p.col += 1; + if (p.col == max_cols) return error.Malformed; + var cols: u32 = 0; + while (cols < span[1] and p.col + cols < max_cols and p.until[p.col + cols] <= row) cols += 1; + const rows = @min(span[0], max_rows - row); + for (p.until[p.col..][0..cols]) |*u| u.* = row + rows; + p.cell = .{ .text = "", .row = row, .col = p.col, .rows = rows, .cols = cols, .head = head }; + p.col += cols; + } + + fn endCell(p: *Parser) ParseError!void { + var cell = p.cell orelse return; + p.cell = null; + p.space = false; + defer p.text.clearRetainingCapacity(); + try p.cells.ensureUnusedCapacity(p.a, 1); + cell.text = try p.a.dupe(u8, std.mem.trim(u8, p.text.items, " \n")); + p.cells.appendAssumeCapacity(cell); + } +}; diff --git a/zig/src/core/ocr_html_test.zig b/zig/src/core/ocr_html_test.zig new file mode 100644 index 00000000000..6ec8517cbf8 --- /dev/null +++ b/zig/src/core/ocr_html_test.zig @@ -0,0 +1,252 @@ +//! ocr_html.zig against tables dots.mocr and GLM-OCR wrote, and the subset's edges. +const std = @import("std"); +const html = @import("ocr_html.zig"); +const t = std.testing; + +// dots.mocr, TABLE II of the layout oracle (run008), first body rows kept. +const dots_table = "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "
ControlsHodgkin lymphomaNon-Hodgkin lymphoma
CasesOR95% CICasesOR95% CI
Never Exposed (reference group)2,2733151.00—1,8231.00—
>1 year9070.840.37-1.91861.170.87-1.59
"; + +// GLM-OCR, oracle case table-007: no header row. +const glm_table = "" ++ + "
ItemQuantity
Apples3
Pears5
"; + +fn expectCell(got: html.Cell, text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool) !void { + const want: html.Cell = .{ .text = text, .row = row, .col = col, .rows = rows, .cols = cols, .head = head }; + try t.expectEqualDeep(want, got); +} + +test "ocr html: the dots.mocr table keeps its spans" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, dots_table); + try t.expectEqual(@as(u32, 4), table.rows); + try t.expectEqual(@as(u32, 8), table.cols); + try t.expect(table.spanned); + try expectCell(table.cells[0], "", 0, 0, 2, 1, true); + const nhl = html.Cell{ .text = "Non-Hodgkin lymphoma", .row = 0, .col = 5, .rows = 1, .cols = 3, .head = true }; + try t.expectEqualDeep(nhl, table.cells[3]); + try expectCell(table.cells[4], "Cases", 1, 2, 1, 1, true); + try t.expect(!table.cells[10].head); +} + +test "ocr html: the GLM-OCR table has no header row and becomes a pipe table with its first row the header" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, glm_table); + try t.expect(!table.spanned); + for (table.cells) |c| try t.expect(!c.head); + const want = "| Item | Quantity |\n| --- | --- |\n| Apples | 3 |\n| Pears | 5 |"; + try t.expectEqualStrings(want, try html.pipe(a, table)); + try t.expectEqualStrings("" ++ + "
ItemQuantity
Apples3
Pears5
", try html.normalized(a, table)); +} + +test "ocr html: entities decode and unknown ones stay as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "
a & b <c> "q" ABC x y z" ++ + " &bogus; & � � � � &
"; + const table = try html.parse(a, src); + const want = "a & b \"q\" ABC x y z &bogus; & \u{fffd} \u{fffd} \u{fffd} � &"; + try t.expectEqualStrings(want, table.cells[0].text); +} + +test "ocr html: other tags drop and keep their text; names and attributes in any case and quoting" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "\n " ++ + "
" ++ + "Bold and it2!
xy
\n"; + const table = try html.parse(a, src); + try t.expectEqual(@as(u32, 3), table.cols); + try expectCell(table.cells[0], "Bold and it2!", 0, 0, 1, 2, true); + try expectCell(table.cells[2], "y", 1, 1, 1, 2, false); +} + +test "ocr html:
is a line break in the cell and a space in pipe cells" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "

line one
line two
\n three
b
"; + const table = try html.parse(a, src); + try t.expectEqualStrings("line one\nline two\nthree", table.cells[0].text); + try t.expectEqualStrings("| line one line two three | b |\n| --- | --- |", try html.pipe(a, table)); + const clean = "
line one
line two
three		b
"; + try t.expectEqualStrings(clean, try html.normalized(a, table)); +} + +test "ocr html: pipe cells escape |" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
a|b
|c
"); + try t.expectEqualStrings("| a\\|b | |\n| --- | --- |\n| \\| | c |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal < that would start a tag, and keep other < as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
<em>x</em>a < b
"); + try t.expectEqualStrings("| <em>x</em> | a < b |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal & that would start an entity, and keep other & as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
&lt; &#38; &#x26;AT&T & b; &;
"); + try t.expectEqualStrings("| &lt; &#38; &#x26; | AT&T & b; &; |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: an unknown tag whose name starts like a table tag is dropped" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
xy
"); + try t.expectEqual(@as(usize, 1), table.cells.len); + try t.expectEqualStrings("xy", table.cells[0].text); +} + +test "ocr html: malformed tables and tables past the limits" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bad = [_][]const u8{ + "", // unclosed table + "
a
", // unclosed quote + "
a
a