", "The files are a.py and b.py."]
+
+
+def test_gemma_thought_channel_with_thinking_off_never_reaches_content():
+ app = _app(GEMMA, [4, 5, 6, 7, EOS])
+ try:
+ from tensorfold.server.text import CHANNEL_MARKERS
+
+ assert app.think_markers == CHANNEL_MARKERS
+ message, streamed, streamed_reasoning = _reply(app)
+ assert message["content"] == "The files are a.py and b.py."
+ assert not message.get("reasoning_content") # the block was empty
+ assert streamed == "The files are a.py and b.py." and streamed_reasoning == ""
+ finally:
+ app.close()
+
+
+HARMONY = ["", "", "", "<|return|>", "<|channel|>", "analysis", "<|message|>", "Think it over.", "<|end|>",
+ "<|start|>", "assistant", "final", "Hello there."]
+STREAMED_HARMONY = ("Hello there.", "Think it over.") # what v0.6.0 streams for it, measured before the change
+
+
+def test_harmony_reply_with_thinking_off_is_unchanged():
+ """gpt-oss replies with thinking off come out exactly as before the Gemma change: parse_harmony_output."""
+
+ from tensorfold.server.text import CHANNEL_MARKERS, parse_harmony_output
+
+ script = [4, 5, 6, 7, 8, 9, 10, 4, 11, 6, 12, EOS]
+ app = _app(HARMONY, script)
+ try:
+ assert app.think_markers != CHANNEL_MARKERS # no token: the Harmony path
+ text = "".join(HARMONY[t] for t in script[:-1])
+ assert parse_harmony_output(text) == ("Hello there.", "Think it over.")
+ message, streamed, streamed_reasoning = _reply(app)
+ assert (message["content"], message.get("reasoning_content")) == ("Hello there.", "Think it over.")
+ assert (streamed, streamed_reasoning) == STREAMED_HARMONY
+ finally:
+ app.close()
+
+
+def test_plain_reply_without_channels_is_unchanged_for_both_families():
+ for pieces in (GEMMA, HARMONY):
+ app = _app(pieces + ["Just an answer."], [len(pieces), EOS])
+ try:
+ message, streamed, _ = _reply(app)
+ assert message["content"] == "Just an answer." and streamed == "Just an answer."
+ assert not message.get("reasoning_content")
+ finally:
+ app.close()
+
From 17e6390f6e27609af02b0afde67ccd89cd22a50e Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 14:09:55 +0100
Subject: [PATCH 079/177] server: Flash Next's sparse-attention indexer is
priced by allocated capacity
---
src/tensorfold/server/memory_budget.py | 3 +--
tests/test_memory_budget.py | 25 +++++++++++++++++++++++++
2 files changed, 26 insertions(+), 2 deletions(-)
diff --git a/src/tensorfold/server/memory_budget.py b/src/tensorfold/server/memory_budget.py
index 90fef952e..7a1d03e10 100644
--- a/src/tensorfold/server/memory_budget.py
+++ b/src/tensorfold/server/memory_budget.py
@@ -162,8 +162,7 @@ def from_cache(cls, cache: Sequence[Any]) -> "CacheMemory":
spare = _array_bytes(getattr(item, "spare_keys", None))
spare += _array_bytes(getattr(item, "spare_values", None))
extra = max(0, held - main - spare)
- valid = max(1, min(positions, int(getattr(item, "offset", positions))))
- auxiliary = -(-extra // valid)
+ auxiliary = -(-extra // positions) # index keys and pooled blocks fill whole capacity steps
capacity = int(getattr(item, "max_size", 0) or 0)
if capacity:
fixed += max(held, (each + auxiliary) * capacity)
diff --git a/tests/test_memory_budget.py b/tests/test_memory_budget.py
index 660e49b0d..a75a24793 100644
--- a/tests/test_memory_budget.py
+++ b/tests/test_memory_budget.py
@@ -207,6 +207,31 @@ def state(self):
assert cache_nbytes(cache) > sum(array.nbytes for array in cache[0].state)
+def test_sparse_indexer_price_is_steady_not_inflated_by_a_short_request():
+ # Flash Next's sparse-attention indexer keys and pooled blocks are allocated in 256-position
+ # capacity steps (AttentionCache.update), so their per-token cost is a steady layer property.
+ # issue 95: pricing them per valid token made a 15-token request look ~3x the startup probe;
+ # observe_cache keeps the largest profile it has seen, so every later prompt inside the window
+ # was then refused.
+ def layer(capacity, offset):
+ return SimpleNamespace(
+ keys=Array((1, 2, capacity, 256)), values=Array((1, 2, capacity, 256)),
+ index_keys=Array((1, capacity, 128)), pooled=Array((1, capacity // 4, 128)),
+ offset=offset)
+
+ probe = CacheMemory.from_cache([layer(2304, 2112)]) # the startup probe: 2,112 tokens of a 2,304-position cache
+ short = CacheMemory.from_cache([layer(256, 15)]) # a 15-token request on a 256-position cache
+ full = CacheMemory.from_cache([layer(2304, 2304)]) # a long prompt that fills the 2,304-position cache
+
+ each = 2 * 2 * 256 * 2 # keys and values per position: 2 tensors x 2 heads x 256 dim x bf16
+ auxiliary = 1 * 128 * 2 + (1 * 128 * 2) // 4 # indexer key per position plus one pooled block per 4 positions
+ steady = each + auxiliary
+ assert probe.bytes_per_token == steady
+ assert short.bytes_per_token == steady
+ assert short.bytes_per_token == probe.bytes_per_token
+ assert full.bytes_per_token == steady # a full cache prices the same before and after the fix
+
+
def test_admission_reserves_reply_work_and_checkpoint_copies():
profile = CacheMemory(100, 2, 16)
projected = needed_bytes(profile, 65, resident_bytes=1000, working_bytes=100,
From 50dfe38aaee41edf49657d57d2378180c9553713 Mon Sep 17 00:00:00 2001
From: Dorian <9122189+SxMShaDoW@users.noreply.github.com>
Date: Thu, 1 Oct 2026 12:42:59 +0100
Subject: [PATCH 080/177] fix(server): a refused POST's body no longer reaches
the next request on its connection (#181)
Both servers answered a POST to an unknown route with 404 before they read
its body, and kept the connection open. The next request on that connection
began with those bytes and got 400 "Bad request syntax". Behind a pooling
proxy that request can come from any client. The MLX server did the same for
a body over 32 MiB, so its client waited until it timed out.
A refused body is now read when it is within the 32 MiB limit; otherwise the
connection closes. A JSON reply on a closing connection says
"Connection: close", so a pooling client does not reuse the socket.
---
src/tensorfold/cuda/http.py | 15 ++++++++
src/tensorfold/server/http.py | 16 +++++++++
src/tensorfold/server/responses.py | 2 ++
tests/test_server_refused_bodies.py | 54 +++++++++++++++++++++++++++++
4 files changed, 87 insertions(+)
create mode 100644 tests/test_server_refused_bodies.py
diff --git a/src/tensorfold/cuda/http.py b/src/tensorfold/cuda/http.py
index 7fe67e169..1cc436c29 100644
--- a/src/tensorfold/cuda/http.py
+++ b/src/tensorfold/cuda/http.py
@@ -50,11 +50,25 @@ def _json(self, code: int, payload: dict[str, Any]) -> None:
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(data)))
+ if self.close_connection: # so a pooling client does not reuse the socket
+ self.send_header("Connection", "close")
self.end_headers()
self.wfile.write(data)
except (BrokenPipeError, ConnectionResetError): # the client has gone
self.close_connection = True
+ def _discard_body(self) -> None:
+ """Read a refused request's body, so it cannot reach the next request on this connection."""
+
+ try:
+ length = int(self.headers.get("Content-Length", 0))
+ except ValueError:
+ length = -1
+ if 0 <= length <= 32 * 1024**2:
+ self.rfile.read(length)
+ else:
+ self.close_connection = True
+
def _stream_error(self, error: dict[str, Any]) -> None:
"""End an open stream with an error event and ``[DONE]``, as the MLX server does."""
@@ -90,6 +104,7 @@ def do_POST(self):
return responses.post(self, app)
chat = self.path.rstrip("/").endswith("/chat/completions")
if not chat and not self.path.rstrip("/").endswith("/completions"):
+ self._discard_body()
return self._json(404, {"error": "not found"})
try:
length = int(self.headers.get("Content-Length", 0))
diff --git a/src/tensorfold/server/http.py b/src/tensorfold/server/http.py
index eb17baf20..2f50012d6 100644
--- a/src/tensorfold/server/http.py
+++ b/src/tensorfold/server/http.py
@@ -93,9 +93,23 @@ def _send_json(self, payload: dict[str, Any], status: int = 200) -> None:
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
+ if self.close_connection: # so a pooling client does not reuse the socket
+ self.send_header("Connection", "close")
self.end_headers()
self.wfile.write(body)
+ def _discard_body(self) -> None:
+ """Read a refused request's body, so it cannot reach the next request on this connection."""
+
+ try:
+ length = int(self.headers.get("Content-Length", "0"))
+ except ValueError:
+ length = -1
+ if 0 <= length <= 32 * 1024**2:
+ self.rfile.read(length)
+ else:
+ self.close_connection = True
+
def _route(self) -> str:
# Tolerate query strings, trailing slashes and client URLs with or without the /v1 prefix.
return self.path.split("?", 1)[0].rstrip("/")
@@ -178,12 +192,14 @@ def do_POST(self) -> None:
is_chat_completion = route.endswith("/chat/completions")
is_text_completion = route.endswith("/completions") and not is_chat_completion
if not is_chat_completion and not is_text_completion:
+ self._discard_body()
self._send_json({"error": {"message": f"unknown path {self.path}"}}, status=404)
return
try:
length = int(self.headers.get("Content-Length", "0"))
if not 0 <= length <= 32 * 1024**2:
+ self.close_connection = True # the unread body must not reach the next request
raise RequestError("request body exceeds the 32 MiB limit")
body = parse_numbers(json.loads(self.rfile.read(length) or b"{}"))
validate_modalities(body)
diff --git a/src/tensorfold/server/responses.py b/src/tensorfold/server/responses.py
index 8cf8212b8..90179db65 100644
--- a/src/tensorfold/server/responses.py
+++ b/src/tensorfold/server/responses.py
@@ -111,6 +111,8 @@ def _send(handler: Any, status: int, payload: dict[str, Any]) -> None:
handler.send_response(status)
handler.send_header("Content-Type", "application/json")
handler.send_header("Content-Length", str(len(data)))
+ if handler.close_connection: # so a pooling client does not reuse the socket
+ handler.send_header("Connection", "close")
handler.end_headers()
handler.wfile.write(data)
except OSError: # the client has gone
diff --git a/tests/test_server_refused_bodies.py b/tests/test_server_refused_bodies.py
new file mode 100644
index 000000000..0d54c64da
--- /dev/null
+++ b/tests/test_server_refused_bodies.py
@@ -0,0 +1,54 @@
+"""A refused request's body never reaches the next request on its connection, on either server."""
+
+import http.client
+import json
+import socket
+
+import pytest
+
+pytest.importorskip("jinja2")
+
+from tests.test_cuda_admission import http_server
+from tests.test_cuda_server_errors import HI, app_for
+from tests.test_server_openai_compat import FakeApp, serve_fake
+
+
+@pytest.fixture(params=["cuda", "mac"])
+def port(request, tmp_path):
+ if request.param == "cuda":
+ with http_server(app_for(tmp_path)) as cuda_port:
+ yield cuda_port
+ else:
+ server = serve_fake(FakeApp())
+ yield server.server_port
+ server.shutdown()
+ server.server_close()
+
+
+def test_a_post_to_an_unknown_route_leaves_the_connection_usable(port):
+ connection = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
+ try:
+ connection.request("POST", "/v1/messages", json.dumps({"max_tokens": 16}),
+ {"Content-Type": "application/json"})
+ refused = connection.getresponse()
+ refused.read()
+ connection.request("POST", "/v1/chat/completions", json.dumps({"messages": HI, "max_tokens": 8}),
+ {"Content-Type": "application/json"})
+ answered = connection.getresponse()
+ answered.read()
+ finally:
+ connection.close()
+ assert (refused.status, answered.status) == (404, 200)
+
+
+@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/responses"])
+def test_an_oversized_body_is_refused_and_the_connection_closed_as_the_reply_says(port, path):
+ with socket.create_connection(("127.0.0.1", port), timeout=10) as client:
+ client.sendall(f"POST {path} HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\n"
+ "Content-Length: 40000000\r\n\r\n".encode())
+ reply = b""
+ while chunk := client.recv(65536): # the server closes once it has answered
+ reply += chunk
+ head = reply.split(b"\r\n\r\n", 1)[0].decode().lower().splitlines()
+ assert head[0].split()[1] == "400"
+ assert "connection: close" in head
From e624f7c4a77bee05314c78f78e933fdffb914886 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:27:16 +0100
Subject: [PATCH 081/177] perf(27b-cuda): streams plan on their measured round
cost, and the drafter drafts the depth the trees keep
---
.../families/qwen3_5/cuda/dflash2.py | 1 +
.../families/qwen3_5/cuda/forward.py | 6 +-
src/tensorfold/families/qwen3_5/cuda/multi.py | 64 ++++++++-
tests/test_qwen27_stream_depth.py | 134 ++++++++++++++++++
4 files changed, 198 insertions(+), 7 deletions(-)
create mode 100644 tests/test_qwen27_stream_depth.py
diff --git a/src/tensorfold/families/qwen3_5/cuda/dflash2.py b/src/tensorfold/families/qwen3_5/cuda/dflash2.py
index 61478bc1b..af380307a 100644
--- a/src/tensorfold/families/qwen3_5/cuda/dflash2.py
+++ b/src/tensorfold/families/qwen3_5/cuda/dflash2.py
@@ -198,6 +198,7 @@ def __init__(self, draft_dir: str | Path, target: Weights, bits: int = 4, block:
self.eps = float(cfg["rms_norm_eps"])
self.theta = float(cfg["rope_parameters"]["rope_theta"])
self.mask_id = int(cfg["dflash_config"]["mask_token_id"])
+ self.trained = int(cfg["dflash_config"].get("block_size", 8)) # its training block: the planner's floor
self.group_size = int(cfg["dflash_config"]["conv_group_size"])
self.layers = int(cfg["num_hidden_layers"])
self.window = int(cfg["sliding_window"]) - 1
diff --git a/src/tensorfold/families/qwen3_5/cuda/forward.py b/src/tensorfold/families/qwen3_5/cuda/forward.py
index 04e306256..65662e1d8 100644
--- a/src/tensorfold/families/qwen3_5/cuda/forward.py
+++ b/src/tensorfold/families/qwen3_5/cuda/forward.py
@@ -455,14 +455,18 @@ def _commit(states: Sequence[State], record: Sequence[Record], paths: Sequence[S
news = src.index_select(1, pick_t).unbind(0)
else: # many rows: gather a layer at a time, the record uncopied
news = [torch.cat([old, t.qkv]).index_select(0, pick_t) for old, t in zip(olds, items)]
+ dst, src = [], []
for j, ((i, _), new) in enumerate(zip(linear, news)):
for s, st in enumerate(states):
if replayed is not None:
st.rec[i] = replayed[s, j]
if in_place:
- st.conv[i].copy_(new[s * keep:(s + 1) * keep])
+ dst.append(st.conv[i])
+ src.append(new[s * keep:(s + 1) * keep])
else:
st.conv[i] = new[s * keep:(s + 1) * keep]
+ if dst:
+ torch._foreach_copy_(dst, src) # every layer's and stream's window in one launch
if att:
for st, path in zip(states, paths):
need = st.pos + len(path)
diff --git a/src/tensorfold/families/qwen3_5/cuda/multi.py b/src/tensorfold/families/qwen3_5/cuda/multi.py
index bd7e91ee9..5f4b0bfd7 100644
--- a/src/tensorfold/families/qwen3_5/cuda/multi.py
+++ b/src/tensorfold/families/qwen3_5/cuda/multi.py
@@ -26,6 +26,15 @@
STEP = 1024 # prompt rows a prefill step takes while other streams decode
GROW = 8192 # rows a stream's attention caches grow by at a time (one GPU)
GIB = 1024**3
+TIMED = 16 # the last rounds whose time beside the forward sets a stream count's overhead
+DEPTH_CHIPS = ((12, 0),) # where the drafter's block follows the trees; elsewhere it drafts every level
+
+
+def calibration_rows(streams: int) -> list[int]:
+ """Row counts the startup curve times: ``streams`` full windows and a point past each lane-matmul tile step."""
+
+ grid = (1, 2, 4, 8, 12, 16, 17, 24, 32, 33, 48, 64, 65, 96, 128, 129, 192, 256, 257, 384, 512)
+ return sorted({r for r in grid if r <= 16 * streams} | {16 * streams})
def private(st: State, rows: int) -> State:
@@ -79,6 +88,10 @@ class MultiDecoder:
"""The ``Scheduler``'s decoder on one GPU or as ``rank`` of two; a stream's window holds at most 16 rows."""
memory_gate: MemoryGate | None = None # one GPU: streams' caches grow by use (two ranks reserve up front)
+ block: int = 16 # rows of the drafter's next block with several streams (pending, masks)
+ depth: bool = True # whether the block follows the trees here (DEPTH_CHIPS)
+ spent: dict | None = None # streams -> the last rounds' ms beside the forward
+ last: tuple | None = None # (start, streams, rows) of the round before
def __init__(self, w: Weights, draft=None, *, max_rows: int = 16, allow_copy: bool = True, stop_eos: bool = True,
keep: int = 8, rank: int = 0, world: int = 1, context: int = 0, points=None, vision=None) -> None:
@@ -89,6 +102,7 @@ def __init__(self, w: Weights, draft=None, *, max_rows: int = 16, allow_copy: bo
self.context = context # prompt plus reply tokens a stream holds (0: no bound)
self.eos = tuple(w.config.eos) if stop_eos else ()
self.rank, self.world, self.device = rank, world, w.norm.device
+ self.depth = torch.cuda.is_available() and tuple(torch.cuda.get_device_capability(self.device)) in DEPTH_CHIPS
self.split = world == 2 and 2 * w.head.n == w.config.vocab # each rank holds half the head
self.drafts = draft is not None and (rank == 0 or getattr(draft, "world", 1) == 2)
self.streams: dict[int, Stream] = {} # decoding
@@ -98,7 +112,8 @@ def __init__(self, w: Weights, draft=None, *, max_rows: int = 16, allow_copy: bo
self.next_id = 0
self.broken: Exception | None = None
self.costs: list[tuple[int, float]] | None = None # (rows, ms) of the forward: tree widths by the curve
- self.overhead = (8.0, 1.5) # a round's other ms: fixed, and per stream
+ self.overhead = (8.0, 1.5) # a round's other ms until measured: fixed, per stream
+ self.block = max_rows
# one GPU: a stream's caches hold its prompt, then grow a step at a time while the gate has room
c, att = w.config, sum(1 for layer in getattr(w, "layers", ()) if not layer.linear)
self.layer_bytes = 2 * getattr(c, "kv_heads", 0) * getattr(c, "head_dim", 0) * 2 # a row of one layer
@@ -291,12 +306,15 @@ def round(self) -> list[Stream]:
"""A prefill step for the next queued prompt, then one round over the decoding streams; returns the finished."""
self._check()
+ self._timed(time.perf_counter())
done = self._fill() if self.filling else []
+ start = time.perf_counter()
live = [s for s in self.streams.values() if not s.done]
if self.memory_gate is not None:
done += self._make_room(live)
live = [s for s in live if not s.done and not s.waiting]
if not live:
+ self.last = None
return done
copied: dict[int, list[int]] = {}
plan = [(s.sid, self._mode(s, copied), s.out[-1], len(s.context)) for s in live]
@@ -309,6 +327,7 @@ def round(self) -> list[Stream]:
ends.append(end)
self._send([x for path in paths for x in (len(path), *path)])
self._commit(plan, wins, record, taps, starts, paths)
+ self.last = (start, len(plan), sum(len(t) for t, _ in wins)) if self.costs is not None else None
for s, (tokens, _), path, end in zip(live, wins, paths, ends):
new = [tokens[r] for r in path[1:]] + [end]
if s.constraint is not None and s.error is None:
@@ -320,8 +339,31 @@ def round(self) -> list[Stream]:
s.done, s.finished = True, time.perf_counter()
continue
s.take(new, self._ends(s))
+ if all(s.done for s in live):
+ self.last = None # the next round waits for requests: not this round's time
return done + [s for s in live if s.done]
+ def _timed(self, now: float) -> None:
+ """The round before: its time beside the forward (its start to this round's, less the curve's forward)."""
+
+ if self.last is None:
+ return
+ start, n, rows = self.last
+ self.last = None
+ if self.spent is None:
+ self.spent = {}
+ seen = self.spent.setdefault(n, [])
+ seen.append(max(0.0, 1e3 * (now - start) - self._cost(rows)))
+ del seen[:-TIMED]
+
+ def _overhead(self, n: int) -> float:
+ """A round's ms beside the forward at ``n`` streams: the median of the last rounds' (one stream: the prior)."""
+
+ seen = (self.spent or {}).get(n) if n > 1 else None
+ if not seen or len(seen) < 4:
+ return self.overhead[0] + self.overhead[1] * n
+ return sorted(seen)[len(seen) // 2]
+
def _ends(self, s: Stream) -> tuple[int, ...]:
"""The end tokens that end this stream: none when its request ignores them (rank 1 follows rank 0's paths)."""
@@ -355,7 +397,7 @@ def _windows(self, plan, copied, blocks) -> list[tuple[list[int], list[int]]]:
sids = list(trees)
fixed = sum(1 + (len(copied[sid]) if mode == COPY else 0) for sid, mode, _, _ in plan)
counts = allocate([trees[sid][2] for sid in sids], fixed, float(len(plan)), self._cost,
- self.overhead[0] + self.overhead[1] * len(plan))
+ self._overhead(len(plan)))
keep = dict(zip(sids, counts))
wins = []
for sid, mode, pending, _ in plan:
@@ -372,9 +414,7 @@ def calibrate(self, streams: int, reps: int = 3) -> None:
"""Time the forward at the row counts ``streams`` windows bring; every rank runs the same forwards."""
st, points = State(self.w), []
- rows = sorted({r for r in (1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512)
- if r <= 16 * streams} | {16 * streams})
- for r in rows:
+ for r in calibration_rows(streams):
n = -(-r // 16)
sizes = [r // n + (i < r % n) for i in range(n)]
wins = [([0] * k, list(range(-1, k - 1)), st) for k in sizes]
@@ -395,9 +435,10 @@ def _verify(self, plan, copied=None):
blocks = {}
tree = [(sid, pending) for sid, mode, pending, _ in plan if mode == TREE] if self.drafts else []
+ block = self.block if self.depth else self.max_rows # one stream too: DFlash2 drafts best near 8 rows
if tree: # every stream's block in one drafter pass
launched = self.draft.launch_blocks([self.streams[sid].snap for sid, _ in tree],
- [pending for _, pending in tree], self.max_rows - 1)
+ [pending for _, pending in tree], self.max_rows - 1, block)
blocks = {sid: block for (sid, _), block in zip(tree, launched)}
grammars = {}
if self.rank == 0:
@@ -407,6 +448,7 @@ def _verify(self, plan, copied=None):
else:
wins = _unflatten(_share(None, 1, self.device), pairs=True)
grammars = self._masks(plan, wins) if self.split else {}
+ self.block = self._deepest(plan, wins, block)
states = [self.streams[item[0]].st for item in plan]
taps_wanted = self.drafts and any(self.streams[item[0]].draft for item in plan)
logits, record, taps, starts = multi_tree_forward(
@@ -424,6 +466,16 @@ def _verify(self, plan, copied=None):
sampled = sample_streams(logits, starts, positions, samplings) if self.rank == 0 else [None] * len(plan)
return wins, record, taps, starts, sampled
+ def _deepest(self, plan, wins, block: int) -> int:
+ """The drafter's next block: past the deepest kept tree node, between its trained block and ``max_rows``."""
+
+ deepest = max((max(_paths(parents)[0]) for (_, mode, _, _), (_, parents) in zip(plan, wins) if mode == TREE),
+ default=-1)
+ if deepest < 0:
+ return self.block
+ floor = max(4, getattr(self.draft, "trained", 4)) # a shorter block cuts runs the drafter would have kept
+ return min(self.max_rows, max(floor, deepest + (5 if deepest >= block - 1 else 2)))
+
def _constrain(self, plan, wins) -> dict:
"""Rank 0: each constrained stream's window without the drafts its grammar rules out, and its rows' masks."""
diff --git a/tests/test_qwen27_stream_depth.py b/tests/test_qwen27_stream_depth.py
new file mode 100644
index 000000000..8a34ee293
--- /dev/null
+++ b/tests/test_qwen27_stream_depth.py
@@ -0,0 +1,134 @@
+"""The 27B's stream planner: measured round overhead, the drafter's block by tree depth, the startup curve's steps."""
+
+import importlib
+from types import SimpleNamespace
+
+import pytest
+
+from tensorfold.cuda.streams import Stream
+from tests.test_cuda_27b_ignore_eos import scripted_decoder
+from tests.test_cuda_geometry import allocations # noqa: F401 (fixture: fake triton, so the module imports)
+
+pytestmark = pytest.mark.torch
+
+CURVE = [(1, 14.0), (16, 16.3), (17, 23.5), (32, 23.9), (33, 33.0), (48, 34.0), (64, 38.8)]
+
+
+class Drafter:
+ """Records the block each pass drafts; a stream's tree is a chain as deep as ``depths`` says."""
+
+ def __init__(self, depths):
+ self.depths, self.blocks = depths, []
+
+ def launch_blocks(self, snaps, pendings, max_nodes, block=None):
+ self.blocks.append(block)
+ return list(range(len(snaps)))
+
+ def finish_tree(self, launched, length, max_nodes, sampling):
+ d = self.depths[launched]
+ return list(range(1, d + 1)), list(range(-1, d - 1)), [0.01 * (i + 1) for i in range(d)]
+
+
+def decoder(multi, monkeypatch, depths, costs=None):
+ dec = multi.MultiDecoder.__new__(multi.MultiDecoder)
+ dec.w, dec.max_rows, dec.costs, dec.overhead = None, 16, costs, (8.0, 1.5)
+ dec.rank, dec.world, dec.split, dec.drafts, dec.device = 0, 1, False, True, None
+ dec.draft = Drafter(depths)
+ dec.streams = {sid: SimpleNamespace(snap=None, sampling=None, st=SimpleNamespace(pos=10), draft=True,
+ constraint=None) for sid in range(len(depths))}
+ monkeypatch.setattr(multi, "multi_tree_forward", lambda w, wins, **kw: (None, None, None, None))
+ monkeypatch.setattr(multi, "sample_streams", lambda logits, starts, positions, samplings: [[]] * len(samplings))
+ return dec
+
+
+def plan(dec, modes=None):
+ return [(sid, (modes or {}).get(sid, 1), 5, 20) for sid in dec.streams] # 1: TREE
+
+
+def test_one_stream_drafts_the_block_its_trees_keep_and_keeps_the_prior(allocations, monkeypatch): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ dec = decoder(multi, monkeypatch, [3], CURVE)
+ dec.block, dec.spent = 6, {1: [1.0] * 8}
+ dec._verify(plan(dec), {})
+ assert dec.draft.blocks == [6] and dec.block == 5 and dec._overhead(1) == 9.5
+
+
+def test_streams_draft_a_level_below_the_deepest_kept_node(allocations, monkeypatch): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ dec = decoder(multi, monkeypatch, [3, 5, 2])
+ dec._verify(plan(dec), {})
+ assert dec.draft.blocks == [16] and dec.block == 7 # depth 5 kept: 6 levels and the pending row
+ dec.draft.depths = [6, 1, 1] # a chain to the block's last level: four more
+ dec._verify(plan(dec), {})
+ assert dec.draft.blocks == [16, 7] and dec.block == 11
+ dec.draft.depths = [1, 1, 1]
+ dec._verify(plan(dec), {})
+ assert dec.block == 4 # never below four rows
+ dec._verify(plan(dec, {0: 2, 1: 2, 2: 2}), {}) # no tree window this round: the block stays
+ assert dec.block == 4 and dec.draft.blocks == [16, 7, 11]
+
+
+def test_other_gpus_draft_every_level(allocations, monkeypatch): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ assert (12, 0) in multi.DEPTH_CHIPS and (12, 1) not in multi.DEPTH_CHIPS
+ dec = decoder(multi, monkeypatch, [3, 2])
+ dec.depth = False # a GB10 or an unmeasured GPU
+ dec._verify(plan(dec), {})
+ dec._verify(plan(dec), {})
+ assert dec.draft.blocks == [16, 16]
+
+
+def test_the_block_never_drops_under_the_drafters_training_block(allocations, monkeypatch): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ dec = decoder(multi, monkeypatch, [3, 2])
+ dec.draft.trained = 8 # DFlash2's block_size
+ dec._verify(plan(dec), {})
+ assert dec.block == 8 # depth 3 kept: 5 rows would do, 8 drafted
+ dec.draft.depths = [7, 1] # a chain to the block's last level: four more
+ dec._verify(plan(dec), {})
+ assert dec.draft.blocks == [16, 8] and dec.block == 12
+ dec.draft.depths = [9, 1]
+ dec._verify(plan(dec), {})
+ assert dec.block == 11
+
+
+def test_the_overhead_past_one_stream_is_the_median_of_the_last_rounds(allocations): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ dec = multi.MultiDecoder.__new__(multi.MultiDecoder)
+ dec.costs, dec.overhead = CURVE, (8.0, 1.5)
+ for i, ms in enumerate([33.0, 34.0, 60.0]):
+ dec.last = (0.0, 8, 32)
+ dec._timed(ms / 1e3)
+ assert dec._overhead(8) == 20.0 # three rounds: still the prior
+ dec.last = (0.0, 8, 32)
+ dec._timed(0.035)
+ assert dec._overhead(8) == pytest.approx(35.0 - 23.9) # the median of 9.1, 10.1, 11.1, 36.1
+ for _ in range(multi.TIMED):
+ dec.last = (0.0, 8, 32)
+ dec._timed(0.030)
+ assert dec._overhead(8) == pytest.approx(30.0 - 23.9) and len(dec.spent[8]) == multi.TIMED
+ assert dec._overhead(4) == 14.0 # unseen stream counts keep the prior
+
+
+def test_rounds_time_only_while_streams_go_on(allocations): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ dec = scripted_decoder(multi, list(range(10, 40))) # no end token: each stream decodes its count
+ dec.costs, dec.overhead = CURVE, (8.0, 1.5)
+ streams = [Stream([1, 2], 13), Stream([3, 4], 13)]
+ for s in streams:
+ dec.admit(s)
+ dec.finish([s for s in streams if s.done])
+ timed = []
+ while dec.live():
+ dec.finish(dec.round())
+ timed.append(dec.last and dec.last[1:])
+ # one stream while the second prompt fills, then both (three drafts a window), none after the last round
+ assert timed == [(1, 4), (2, 8), (2, 8), None]
+ assert {n: len(v) for n, v in dec.spent.items()} == {1: 1, 2: 2} # each round timed by the next one's start
+
+
+def test_the_curve_times_both_sides_of_each_row_step(allocations): # noqa: F811
+ multi = importlib.import_module("tensorfold.families.qwen3_5.cuda.multi")
+ assert multi.calibration_rows(1) == [1, 2, 4, 8, 12, 16]
+ rows = multi.calibration_rows(8)
+ assert {16, 17, 32, 33, 64, 65, 128} <= set(rows) and rows[-1] == 128 and 129 not in rows
From b3d873372e93dcf98c24a9a18e6dd6aabb99ea1f Mon Sep 17 00:00:00 2001
From: MiaAI-Lab <83042094+MiaAI-Lab@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:12:19 +0300
Subject: [PATCH 082/177] feat(server): reasoning_effort max: GLM-5.3's
template names it (its default level); xhigh where a template has no max
GLM-5.3's chat template sets its effort with `reasoning_effort if reasoning_effort in ['low', 'high'] else 'max'`:
Max is its default and top level. A request that asked for it by name ("max", top level or in
chat_template_kwargs) was refused with HTTP 400, although the template names it.
- EFFORTS accepts "max"; effort_levels reads 'max' from a template, so GLM-5.3's levels are low, high and max
- A template that names max hears max (GLM-5.3: the same prompt as an omitted effort, "Reasoning Effort: Max")
- A template without max hears the nearest named level. With no named levels, max is xhigh, as high already is
- medium is still heard as high on GLM-5.3, and xhigh still stays xhigh
Tests: tests/test_request_reasoning.py (max over HTTP without a template is xhigh; max on a template naming low,
high and max passes through, on Qwen3.8's names and on none it is xhigh; GLM-5.3's template renders max and xhigh as
Max); the refusal messages in tests/test_cuda_server_errors.py and tests/cuda/test_qwen27_server_errors.py;
tests/test_cuda_thinking_controls.py refuses "ultra" where it refused "max". docs/api.md and the GLM-5.3 recipe.
---
src/tensorfold/server/request_options.py | 14 +++++++-------
tests/cuda/test_qwen27_server_errors.py | 2 +-
tests/test_cuda_server_errors.py | 2 +-
tests/test_cuda_thinking_controls.py | 2 +-
tests/test_request_reasoning.py | 15 +++++++++++----
5 files changed, 21 insertions(+), 14 deletions(-)
diff --git a/src/tensorfold/server/request_options.py b/src/tensorfold/server/request_options.py
index fc139d748..c8f3e3ba2 100644
--- a/src/tensorfold/server/request_options.py
+++ b/src/tensorfold/server/request_options.py
@@ -38,15 +38,15 @@ def parse_numbers(fields: dict[str, Any]) -> dict[str, Any]:
return parsed
-EFFORTS = ("none", "minimal", "low", "medium", "high", "xhigh")
+EFFORTS = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
-_EFFORT_ORDER = ("xhigh", "high", "medium", "low", "minimal") # highest first
+_EFFORT_ORDER = ("max", "xhigh", "high", "medium", "low", "minimal") # highest first
def nearest_named_effort(effort: str, levels: frozenset[str]) -> str:
"""The nearest named level, ties going higher; effort itself when none is named or it is off the ladder."""
- # GLM-5.3 names low and high but no medium, so its template used to render medium as Max, its ceiling
+ # max sits above xhigh. A template that names it keeps max; any other template hears the nearest name
if not levels or effort not in _EFFORT_ORDER:
return effort
@@ -55,9 +55,9 @@ def nearest_named_effort(effort: str, levels: frozenset[str]) -> str:
def effort_levels(template: str | None) -> frozenset[str]:
- """The efforts a chat template names: Qwen3.8's low, medium and xhigh; GLM-5.3's low and high."""
+ """The efforts a chat template names: Qwen3.8's low, medium and xhigh; GLM-5.3's low, high and max."""
- return frozenset(re.findall(r"""['"](minimal|low|medium|high|xhigh)['"]""", template or ""))
+ return frozenset(re.findall(r"""['"](minimal|low|medium|high|xhigh|max)['"]""", template or ""))
def coerce_effort(effort: str | None, levels: frozenset[str] = frozenset()) -> str | None:
@@ -66,7 +66,7 @@ def coerce_effort(effort: str | None, levels: frozenset[str] = frozenset()) -> s
if effort is None:
return None
if not levels:
- if effort == "high":
+ if effort in ("high", "max"):
return "xhigh"
if effort == "minimal":
return "low"
@@ -92,7 +92,7 @@ def thinking_fields(body: dict[str, Any], levels: frozenset[str] = frozenset())
effort = kwargs.get("reasoning_effort") # where vLLM's clients put it
if effort is not None:
if not isinstance(effort, str) or effort not in EFFORTS:
- raise RequestError("reasoning_effort must be none, minimal, low, medium, high or xhigh")
+ raise RequestError("reasoning_effort must be none, minimal, low, medium, high, xhigh or max")
fields["reasoning_effort"] = coerce_effort(effort, levels)
fields["enable_thinking"] = effort != "none"
if isinstance(kwargs, dict) and "enable_thinking" in kwargs: # an explicit switch wins
diff --git a/tests/cuda/test_qwen27_server_errors.py b/tests/cuda/test_qwen27_server_errors.py
index 0d007adc9..d2612f95d 100644
--- a/tests/cuda/test_qwen27_server_errors.py
+++ b/tests/cuda/test_qwen27_server_errors.py
@@ -26,7 +26,7 @@
({"messages": [{"role": "tool", "content": "x"}]},
"the chat template rejected the request: No user query found in messages."),
({"messages": HI, "reasoning_effort": "extreme"},
- "reasoning_effort must be none, minimal, low, medium, high or xhigh"),
+ "reasoning_effort must be none, minimal, low, medium, high, xhigh or max"),
({"messages": HI, "chat_template_kwargs": "x"}, "chat_template_kwargs must be a JSON object or null"),
({"messages": HI, "temperature": "hot"}, "temperature must be a finite number or null"),
({"messages": HI, "temperature": 0.7, "seed": "abc"}, "seed must be an integer or null"),
diff --git a/tests/test_cuda_server_errors.py b/tests/test_cuda_server_errors.py
index ede583a04..8de9af9a8 100644
--- a/tests/test_cuda_server_errors.py
+++ b/tests/test_cuda_server_errors.py
@@ -114,7 +114,7 @@ def sampling_before(defaults, body, prompt):
"template raises on a kwarg": ({"messages": HI, "chat_template_kwargs": {"mode": "fancy"}},
"the chat template rejected the request: Unexpected mode fancy."),
"an unknown effort": ({"messages": HI, "chat_template_kwargs": {"reasoning_effort": "extreme"}},
- "reasoning_effort must be none, minimal, low, medium, high or xhigh"),
+ "reasoning_effort must be none, minimal, low, medium, high, xhigh or max"),
**{f"chat_template_kwargs {name}": ({"messages": HI, "chat_template_kwargs": value},
"chat_template_kwargs must be a JSON object or null")
for name, value in [("[]", []), ('""', ""), ("false", False), ("0", 0), # falsy: not read as absent
diff --git a/tests/test_cuda_thinking_controls.py b/tests/test_cuda_thinking_controls.py
index 155ea6c3b..8cb6ad940 100644
--- a/tests/test_cuda_thinking_controls.py
+++ b/tests/test_cuda_thinking_controls.py
@@ -114,7 +114,7 @@ def test_no_server_default_leaves_the_template_its_own(tmp_path):
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("fields", [{"reasoning_effort": "extreme"}, {"reasoning_effort": 3},
- {"chat_template_kwargs": {"reasoning_effort": "max"}},
+ {"chat_template_kwargs": {"reasoning_effort": "ultra"}},
{"thinking_budget": "lots"}, {"thinking_budget": 2.5}])
def test_a_bad_effort_or_budget_is_refused_before_the_stream(tmp_path, stream, fields):
engine = ChainEngine()
diff --git a/tests/test_request_reasoning.py b/tests/test_request_reasoning.py
index 289d518a6..74e528377 100644
--- a/tests/test_request_reasoning.py
+++ b/tests/test_request_reasoning.py
@@ -25,6 +25,7 @@ def chat(self, *args, sampling=None, **kwargs):
("medium", True, "medium"),
("high", True, "xhigh"),
("xhigh", True, "xhigh"),
+ ("max", True, "xhigh"),
],
)
def test_http_reasoning_controls(effort, thinking, normalized):
@@ -138,13 +139,17 @@ def checked_template(messages, **kwargs):
("{# 'xhigh' 'medium' 'low' #}", "low", "low"),
("{# 'xhigh' 'medium' 'low' #}", "medium", "medium"),
("{# 'xhigh' 'medium' 'low' #}", "xhigh", "xhigh"),
+ ("{# 'xhigh' 'medium' 'low' #}", "max", "xhigh"), # no max: the nearest named level is xhigh
("{# 'low' 'high' #}", "high", "high"), # GLM-5.3 names high: it renders High, not Max
("{# 'low' 'high' #}", "minimal", "low"),
("{# 'low' 'high' #}", "low", "low"),
("{# 'low' 'high' #}", "medium", "high"), # medium is not a GLM name; Max was the silent result
("{# 'low' 'high' #}", "xhigh", "xhigh"), # GLM renders xhigh as its own Max; leave the name
("{# 'low' 'high' #}", "none", "none"),
+ ("{# 'low' 'high' 'max' #}", "max", "max"), # the template names max, so the name stays
+ ("{# 'low' 'high' 'max' #}", "medium", "high"),
("", "high", "xhigh"),
+ ("", "max", "xhigh"),
])
def test_a_template_that_names_an_effort_is_given_that_effort(names, effort, want):
app = make_app(enable_thinking=True)
@@ -169,14 +174,14 @@ def test_glm_omitted_effort_stays_the_template_default_and_a_medium_default_is_h
source = ("{%- set effective_reasoning_effort = reasoning_effort if reasoning_effort is defined "
"and reasoning_effort in ['low', 'high'] else 'max' -%}{{ effective_reasoning_effort }}")
levels = effort_levels(source)
- assert levels == frozenset({"low", "high"})
+ assert levels == frozenset({"low", "high", "max"})
template = jinja2.Environment().from_string(source)
assert template.render() == "max"
assert template.render(reasoning_effort="medium") == "max"
- for effort in ("minimal", "low", "medium", "high", "xhigh"):
+ for effort in ("minimal", "low", "medium", "high", "xhigh", "max"):
heard = coerce_effort(effort, levels)
assert template.render(reasoning_effort=heard) == {"minimal": "low", "low": "low", "medium": "high",
- "high": "high", "xhigh": "max"}[effort]
+ "high": "high", "xhigh": "max", "max": "max"}[effort]
app = make_app(enable_thinking=True, reasoning_effort="medium")
app.tokenizer.chat_template = "{# 'low' 'high' #}"
@@ -289,14 +294,16 @@ def test_glm_and_qwen_templates_hear_the_nearest_named_level():
"{%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}"
"{{- raise_exception('Unexpected reasoning effort') }}{%- endif -%}{{ resolved_reasoning_effort }}")
glm_levels, qwen_levels = effort_levels(glm), effort_levels(qwen)
- assert glm_levels == frozenset({"low", "high"})
+ assert glm_levels == frozenset({"low", "high", "max"})
assert qwen_levels == frozenset({"low", "medium", "xhigh"})
glm_template = jinja2.Environment().from_string(glm)
qwen_template = jinja2.Environment().from_string(qwen)
assert glm_template.render(reasoning_effort=coerce_effort("medium", glm_levels)) == "high"
assert glm_template.render(reasoning_effort=coerce_effort("minimal", glm_levels)) == "low"
assert glm_template.render(reasoning_effort=coerce_effort("xhigh", glm_levels)) == "max"
+ assert glm_template.render(reasoning_effort=coerce_effort("max", glm_levels)) == "max"
assert qwen_template.render(reasoning_effort=coerce_effort("high", qwen_levels)) == "xhigh"
+ assert qwen_template.render(reasoning_effort=coerce_effort("max", qwen_levels)) == "xhigh"
assert qwen_template.render(reasoning_effort=coerce_effort("minimal", qwen_levels)) == "low"
assert qwen_template.render(reasoning_effort=coerce_effort("medium", qwen_levels)) == "medium"
assert heard_effort(None, "medium", glm_levels) == "high"
From bc96a9dc8e7093d167b559347f29ca3392825e86 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:36:56 +0100
Subject: [PATCH 083/177] test: an unnamed max uses the nearest named level
---
tests/test_request_reasoning.py | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/tests/test_request_reasoning.py b/tests/test_request_reasoning.py
index 74e528377..1763e49c0 100644
--- a/tests/test_request_reasoning.py
+++ b/tests/test_request_reasoning.py
@@ -221,6 +221,7 @@ def test_no_effort_leaves_the_template_its_own_default():
@pytest.mark.parametrize("names, effort, want", [
("{# 'low' 'high' #}", "medium", "high"), # GLM-5.3 names no medium: the higher of the two levels as near
+ ("{# 'low' 'high' #}", "max", "high"), # max is unnamed here, so the nearest name is high, not xhigh
("{# 'low' 'high' #}", "xhigh", "xhigh"), # the else-branch is the template's own ceiling: as sent
("{# 'low' 'high' #}", "low", "low"),
("{# 'low' 'high' #}", "high", "high"),
@@ -321,6 +322,9 @@ def test_nearest_named_effort_direct():
assert nearest("medium", frozenset()) == "medium" # no named levels: as sent
assert nearest("none", frozenset({"low", "high"})) == "none" # none is not a level: as sent
assert nearest("medium", frozenset({"low", "high"})) == "high" # the GLM-5.3 tie: upward
+ assert nearest("max", frozenset({"xhigh", "high"})) == "xhigh"
+ assert nearest("max", frozenset({"low", "high"})) == "high"
+ assert nearest("max", frozenset({"low", "high", "max"})) == "max"
def test_glm_medium_with_thinking_forced_off_reaches_no_effort():
From 861348862a0f108027575c6e529e527e0bfe754d Mon Sep 17 00:00:00 2001
From: Andrey Kolesnikov <15747200+akol1@users.noreply.github.com>
Date: Wed, 30 Sep 2026 09:35:10 -0700
Subject: [PATCH 084/177] glm5_next: route EXL3 experts through the universal
path for any bit width
The dedicated wrapper pins the 4-bit-mcg layout: words() requires
int16 [..., 64], so a 3-bit checkpoint refuses with 'only 4-bit EXL3
trellises ... are supported'. The universal grouped kernel
(tensorfold/cuda/exl3/experts.*) supports k2 2..16 with a bit width per
expert and is already tested bit-identical to the dedicated path at GLM's
shapes (test_glm_shaped_bit_identical: 288 experts, 4096 -> 1024/2048,
ACT_BF16, GLM tile configs) and against a float64 reference at every width
(test_mixed_k_rows_are_independent_and_match_the_reference).
Route GLM's experts through experts.prepare/routed:
- weights.moe_exl3 builds the layer via experts.prepare, which derives k2
from each tensor's own last dim -- no layout assumption. The result keeps
the same attribute names the rest of the module reads (dims, width,
count, gate_ptr ...), so nothing else changes.
- forward.moe_block calls exl3_generic.routed, which keeps its own
universal Scratch keyed by (rows, slots, expert count): the engine's
exl3_mm.Scratch lacks the generic one's construction, and the generic
Scratch needs the real expert count at construction time.
- The shared-expert contract is preserved: pick == E is skipped by the
universal kernels exactly as exl3_mm skips slot == slots - 1, and the
caller overwrites that slot from the BF16 shared MLP before combine.
The dedicated path remains untouched for 4-bit checkpoints; a checkpoint
whose width matches what exl3_mm reads can keep using it.
Exactness beyond the upstream tests: ext.dequant of a 3-bit mcg trellis
matches format.unpack bit-exactly (max |diff| = 0.0), and the fused
routed() path is deterministic and row-independent (same row, same output
regardless of batch composition).
Verified on a real GLM-5.3-Flash EXL3 3-bit pack (288 experts, D=4096,
I=1024 per rank): both ranks load, /health 200, completions finish,
125-128 tok/s decode with the default drafter policy.
tests/cuda/test_exl3_any_width_glm.py: dequant identity against
format.unpack, k2 derived from the checkpoint, determinism, row
independence, and the shared-expert skip.
---
.../families/glm5_next/cuda/exl3_generic.py | 26 ++++
.../families/glm5_next/cuda/forward.py | 5 +-
.../families/glm5_next/cuda/weights.py | 25 ++--
tests/cuda/test_exl3_any_width_glm.py | 122 ++++++++++++++++++
4 files changed, 165 insertions(+), 13 deletions(-)
create mode 100644 src/tensorfold/families/glm5_next/cuda/exl3_generic.py
create mode 100644 tests/cuda/test_exl3_any_width_glm.py
diff --git a/src/tensorfold/families/glm5_next/cuda/exl3_generic.py b/src/tensorfold/families/glm5_next/cuda/exl3_generic.py
new file mode 100644
index 000000000..07d3346de
--- /dev/null
+++ b/src/tensorfold/families/glm5_next/cuda/exl3_generic.py
@@ -0,0 +1,26 @@
+"""GLM routed experts through the generic mixed-width EXL3 path.
+
+Serves any bit width (k2 2..16) with bit-identical epilogue arithmetic
+(ACT_BF16, GLM tile configs). Uses the exl3_mm.Scratch's rows/slots (passed
+explicitly) to size its own generic Scratch once per key.
+
+The shared expert (pick == E) is excluded by the generic kernels, as in
+exl3_mm; the caller combines its slot from the BF16 shared MLP.
+"""
+from __future__ import annotations
+
+import torch
+
+from tensorfold.cuda.exl3 import experts as generic
+
+_scratch: dict = {}
+
+
+def routed(x, pick, ex, rows, limit):
+ """ey[pair] fp32 for routed pairs; the shared slot's row is stale (caller overwrites)."""
+ key = (rows, pick.shape[1], ex.count)
+ s = _scratch.get(key)
+ if s is None:
+ s = generic.Scratch(ex, rows, pick.shape[1], device=str(x.device))
+ _scratch[key] = s
+ return generic.routed(x, pick, None, ex, s, None, rows, limit, act_mode=generic.ACT_BF16)
diff --git a/src/tensorfold/families/glm5_next/cuda/forward.py b/src/tensorfold/families/glm5_next/cuda/forward.py
index 01a173cb1..f1234d77b 100644
--- a/src/tensorfold/families/glm5_next/cuda/forward.py
+++ b/src/tensorfold/families/glm5_next/cuda/forward.py
@@ -105,6 +105,7 @@ def __init__(self, w: Weights, rows: int, capacity: int = 2560, *, prefill: bool
sl = c.shared_width // w.world
self.exl3 = Scratch(rows, slots, D, ml, dev)
+ self.exl3_rows = rows
self.sgu = torch.empty((rows, 2 * sl), dtype=bf, device=dev)
self.sact = torch.empty((rows, sl), dtype=bf, device=dev)
self.sxs = torch.empty((rows, sl // 64), dtype=f32, device=dev)
@@ -398,7 +399,9 @@ def moe_block(layer: LayerW, w: Weights, b: Buffers, R: int) -> torch.Tensor:
# EXL3: the routed slots through the trellis kernels, the shared expert (last slot) through BF16 matmuls
from . import exl3_mm
- exl3_mm.routed(b.normed[:R], b.pick, b.plan, m.experts, b.exl3, b.ey.view(-1, c.hidden), R, c.limit)
+ from . import exl3_generic
+ ey_flat = exl3_generic.routed(b.normed[:R], b.pick, m.experts, b.exl3.rows, c.limit)
+ b.ey.view(-1, c.hidden)[:ey_flat.shape[0]].copy_(ey_flat)
s = m.shared
mm(b, b.normed[:R], s.gu, b.xs[:R], b.sgu[:R])
glue.swiglu(b.sgu[:R], b.sact[:R], b.sxs[:R], c.limit)
diff --git a/src/tensorfold/families/glm5_next/cuda/weights.py b/src/tensorfold/families/glm5_next/cuda/weights.py
index 48c4190ce..03771a39e 100644
--- a/src/tensorfold/families/glm5_next/cuda/weights.py
+++ b/src/tensorfold/families/glm5_next/cuda/weights.py
@@ -346,18 +346,19 @@ def on_device(tensors: list[torch.Tensor]) -> torch.Tensor:
return out
def moe_exl3(p: str) -> Exl3Experts:
- parts = {}
- for proj in ("gate_proj", "up_proj", "down_proj"):
- ts, us, vs = [], [], []
- for e in range(cfg.experts):
- name = PREFIX + p + f"experts.{e}.{proj}."
- ts.append(exl3_words(rd.get(name + "trellis")))
- us.append(rd.get(name + "suh"))
- vs.append(rd.get(name + "svh"))
- parts[proj] = (torch.stack(ts).to(dev), torch.stack(us).to(dev), torch.stack(vs).to(dev))
- del ts, us, vs
- (gt, sg, vg), (ut, su, vu), (dt, sd, vd) = parts["gate_proj"], parts["up_proj"], parts["down_proj"]
- return Exl3Experts(gt, ut, dt, sg, su, vg, vu, sd, vd, cfg.experts, int(vg.shape[1]), int(vd.shape[1]))
+ from tensorfold.cuda.exl3 import experts as generic
+ gate = generic.prepare(
+ [(rd.get(PREFIX + p + f"experts.{e}.gate_proj.trellis").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.gate_proj.suh").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.gate_proj.svh").to(dev)) for e in range(cfg.experts)],
+ [(rd.get(PREFIX + p + f"experts.{e}.up_proj.trellis").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.up_proj.suh").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.up_proj.svh").to(dev)) for e in range(cfg.experts)],
+ [(rd.get(PREFIX + p + f"experts.{e}.down_proj.trellis").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.down_proj.suh").to(dev),
+ rd.get(PREFIX + p + f"experts.{e}.down_proj.svh").to(dev)) for e in range(cfg.experts)],
+ "mcg", device=dev)
+ return gate
def moe(i: int) -> MoEW:
p = f"layers.{i}.mlp."
diff --git a/tests/cuda/test_exl3_any_width_glm.py b/tests/cuda/test_exl3_any_width_glm.py
new file mode 100644
index 000000000..30ec5904e
--- /dev/null
+++ b/tests/cuda/test_exl3_any_width_glm.py
@@ -0,0 +1,122 @@
+"""GLM-shaped 3-bit (k2 = 6 half-bits) experts through the universal path.
+
+The universal grouped kernel (`tensorfold/cuda/exl3/experts.*`) reads a bit
+width per expert and was tested against a float64 reference at every width
+(`test_mixed_k_rows_are_independent_and_match_the_reference`). GLM's own
+family does not route through it: `families/glm5_next/cuda/exl3_mm.py` pins
+the 4-bit-mcg layout (`words()` requires `int16 [..., 64]`), so a 3-bit
+checkpoint refuses with `only 4-bit EXL3 trellises ... are supported` even
+though the universal kernel serves it. These tests pin the universal path at
+GLM's shapes and bit width, and the two properties an integrator relies on
+before committing to it: determinism across calls, and a row's output being
+the same alone and inside a window.
+
+The layer fixture mirrors a real 3-bit pack's on-disk layout: trellis int16
+`[K/16, N/16, 16 * k2]` words (48 at k2 = 6), fp16 `suh`/`svh` scales, 288
+routed experts at D = 4096, I = 1024 per rank (the pack is rank-split), with
+`GLM_GATEUP`/`GLM_DOWN` tile configs and `ACT_BF16` — the same arithmetic the
+dedicated path runs, verified bit-identical by `test_glm_shaped_bit_identical`
+at 4-bit.
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+E, D, I = 288, 4096, 1024
+K2 = 6 # 3 bits = 6 half-bits a value
+
+
+def _layer(cb="mcg", device="cuda"):
+ from tensorfold.cuda.exl3 import experts
+
+ g = torch.Generator().manual_seed(11)
+
+ def trellis(k, n):
+ return torch.randint(-32768, 32768, (k // 16, n // 16, 8 * K2),
+ dtype=torch.int16, generator=g).to(device).contiguous()
+
+ def scale(n, mag):
+ sign = torch.randint(0, 2, (n,), generator=g).float() * 2 - 1
+ return (sign * (torch.rand((n,), generator=g) + 0.5) * mag).half().to(device)
+
+ gate, up, down = [], [], []
+ for _ in range(E):
+ gate.append((trellis(D, I), scale(D, 0.02), scale(I, 0.5)))
+ up.append((trellis(D, I), scale(D, 0.02), scale(I, 0.5)))
+ down.append((trellis(I, D), scale(I, 0.05), scale(D, 0.2)))
+ return experts.prepare(gate, up, down, cb, device=device)
+
+
+def test_dequant_identity_against_format_unpack():
+ """The kernels' own decode of a 3-bit trellis equals the format's reference
+ decoder, bit for bit. If the pack reads wrong here, everything downstream is
+ built on a wrong W_q and the rest of the file proves nothing."""
+ from tensorfold.cuda.exl3 import experts
+ from tensorfold.cuda.exl3.format import unpack as fmt_unpack
+
+ t = torch.randint(-32768, 32768, (D // 16, I // 16, 8 * K2), dtype=torch.int16).cuda()
+ ref = torch.from_numpy(fmt_unpack(t.cpu().numpy(), 3.0, "mcg")).cuda()
+ assert torch.equal(experts.dequant(t, "mcg"), ref)
+
+
+def test_prepared_layer_reports_k2_from_the_checkpoint():
+ """`prepare` derives k2 from each tensor's own last dim; a 3-bit pack must come
+ out as k2 = 6 half-bits, not the 4-bit assumption of the dedicated wrapper."""
+ from tensorfold.cuda.exl3 import experts
+
+ ex = _layer()
+ assert ex.k2_gu == (K2, K2)
+ assert ex.k2_d == (K2, K2)
+ assert ex.dims == D and ex.width == I and ex.count == E
+
+
+def test_routed_is_deterministic_across_calls():
+ from tensorfold.cuda.exl3 import experts
+
+ ex = _layer()
+ x = torch.randn(4, D, dtype=torch.bfloat16, device="cuda")
+ pick = torch.zeros((4, 4), dtype=torch.int32, device="cuda")
+ s = experts.Scratch(ex, 4, 4, device="cuda")
+ y1 = experts.routed(x, pick, None, ex, s, None, 4, act_mode=experts.ACT_BF16)
+ y2 = experts.routed(x, pick, None, ex, s, None, 4, act_mode=experts.ACT_BF16)
+ assert torch.equal(y1, y2), "repeated calls must be bit-identical"
+
+
+def test_rows_are_independent_of_batch_composition():
+ """A row alone must equal its output inside a window, whatever the other rows
+ hold: the grouped launch groups by expert and reads rows by index."""
+ from tensorfold.cuda.exl3 import experts
+
+ ex = _layer()
+ x = torch.randn(4, D, dtype=torch.bfloat16, device="cuda")
+ pick = torch.zeros((4, 4), dtype=torch.int32, device="cuda")
+
+ s_full = experts.Scratch(ex, 4, 4, device="cuda")
+ y_full = experts.routed(x, pick, None, ex, s_full, None, 4, act_mode=experts.ACT_BF16)
+ s_one = experts.Scratch(ex, 4, 4, device="cuda")
+ y_one = experts.routed(x[:1], pick[:1].contiguous(), None, ex, s_one, None, 1,
+ act_mode=experts.ACT_BF16)
+ assert torch.equal(y_full[0], y_one[0])
+ # duplicated rows with identical content give identical outputs
+ x_dup = torch.cat([x[:1], x[:1]])
+ s_dup = experts.Scratch(ex, 4, 4, device="cuda")
+ y_dup = experts.routed(x_dup, torch.cat([pick[:1], pick[:1]]), None, ex, s_dup, None, 2,
+ act_mode=experts.ACT_BF16)
+ assert torch.equal(y_dup[0], y_dup[1])
+
+
+def test_shared_expert_pairs_are_skipped():
+ """pick == E means the shared expert: the universal kernel leaves those pairs
+ unwritten, matching the dedicated path's `slot == slots - 1` guard. The caller
+ combines the slot from the BF16 shared MLP before reading it."""
+ from tensorfold.cuda.exl3 import experts
+
+ ex = _layer()
+ x = torch.randn(2, D, dtype=torch.bfloat16, device="cuda")
+ pick = torch.full((2, 4), E, dtype=torch.int32, device="cuda")
+ s = experts.Scratch(ex, 2, 4, device="cuda")
+ y = experts.routed(x, pick, None, ex, s, None, 2, act_mode=experts.ACT_BF16)
+ torch.cuda.synchronize()
+ assert torch.all(y == 0)
From 47bf822d88f4f8d60839a5cf0611d3a4d4295d25 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Wed, 30 Sep 2026 22:03:47 +0100
Subject: [PATCH 085/177] feat(cuda): RTX 40 cards (compute capability 8.9); a
kept GLM prompt keeps DFlash2's ring window (#134)
- The lane and prompt matmuls build for Ada (sm_89); sm_90 and newer compile the same kernels as before
- CUDA serves on 8.9, NVFP4 and FP8 checkpoints keep a 9.0 floor, and a GPU below its floor is refused at startup
- The 27B's startup estimate counts what it loads, and its verify scratch the rows a round takes
- The 27B's prompt chunk comes from the GPU's memory: the most rows, up to 4,096, whose arrays fit a 32nd of it
- Clones of a 27B state share one attention list, so a grow frees the buffers it replaces
- On one GPU the 27B's kept prompt states give way, oldest first, so admission counts one live window
- GLM-5.3 on CUDA: a prompt kept one token early keeps DFlash2's ring window, so resends resume with it on (#134)
- The Nemotron resume test expects prompts kept one token early (#98)
- The vision clone test expects the shared attention list and the carried byte budget
---
src/tensorfold/cuda/build.py | 36 ++++--
src/tensorfold/cuda/capacity.py | 29 ++++-
src/tensorfold/cuda/geometry.py | 47 +++++---
src/tensorfold/cuda/kernels/qmm.cpp | 3 +-
src/tensorfold/cuda/kernels/qmm.cu | 12 +-
src/tensorfold/cuda/kernels/qmm_frag.cuh | 5 +
src/tensorfold/cuda/nvfp4/linear.py | 4 +-
src/tensorfold/cuda/streams.py | 33 +++++-
.../families/glm5_next/cuda/decode.py | 12 +-
.../families/glm5_next/cuda/engine.py | 8 +-
.../families/qwen3_5/cuda/affine_memory.py | 38 ++++++-
.../families/qwen3_5/cuda/decode.py | 10 +-
.../families/qwen3_5/cuda/engine.py | 25 +++--
.../families/qwen3_5/cuda/forward.py | 36 +++---
src/tensorfold/families/qwen3_5/cuda/multi.py | 1 +
.../families/qwen3_5/cuda/prefill.py | 27 +----
.../families/qwen3_5/cuda/weights.py | 1 +
tests/cuda/test_glm_draft_ring.py | 106 ++++++++++++++++--
tests/cuda/test_nemotron_app.py | 2 +-
tests/cuda/test_qmm.py | 16 +++
tests/cuda/test_qwen27_kv_room.py | 32 ++++++
tests/cuda/test_qwen27_prompt_end_cache.py | 6 +-
tests/cuda/test_qwen27_shared_kv.py | 34 ++++++
tests/cuda_27b_headers.py | 81 +++++++++++++
tests/test_cuda_build.py | 11 +-
tests/test_cuda_capacity.py | 4 +-
tests/test_cuda_gpu_floor.py | 54 +++++++++
tests/test_cuda_kv_room.py | 98 ++++++++++++++++
tests/test_cuda_memory_plan.py | 73 ++++++++++++
tests/test_qwen27_prompt_end_cache_host.py | 10 +-
tests/test_vision_cuda.py | 7 +-
31 files changed, 734 insertions(+), 127 deletions(-)
create mode 100644 tests/cuda/test_qwen27_kv_room.py
create mode 100644 tests/cuda/test_qwen27_shared_kv.py
create mode 100644 tests/cuda_27b_headers.py
create mode 100644 tests/test_cuda_gpu_floor.py
create mode 100644 tests/test_cuda_kv_room.py
create mode 100644 tests/test_cuda_memory_plan.py
diff --git a/src/tensorfold/cuda/build.py b/src/tensorfold/cuda/build.py
index 5eb6d2bf5..ba1786d17 100644
--- a/src/tensorfold/cuda/build.py
+++ b/src/tensorfold/cuda/build.py
@@ -6,31 +6,47 @@
import threading
from typing import Any
-MIN_CAPABILITY = (9, 0) # the kernels use thread-block clusters and FP8 MMA
+MIN_CAPABILITY = (8, 9) # FP8 MMA and e4m3 conversions (Ada); kernels with clusters use them from 9.0
+CLUSTERS = (9, 0) # extensions built only on thread-block clusters (NVFP4) need Hopper or newer
# stop first: when the lock goes, a waiting start imports whatever module is there without building, even an old one
HINT = "if no other build is running, a killed build left it: stop this start, delete the lock and start again"
LOCK_WAIT_SECONDS = 60.0 # a start still waiting on the same lock this long says so again
-def arch_flags() -> list[str]:
- """nvcc flags for the current GPU alone; a GPU older than the kernels need is refused by name."""
+def arch_flags(need: tuple[int, int] = MIN_CAPABILITY) -> list[str]:
+ """nvcc flags for the current GPU alone; a GPU older than ``need`` is refused by name."""
import torch
major, minor = torch.cuda.get_device_capability()
- if (major, minor) < MIN_CAPABILITY:
- raise RuntimeError(f"TensorFold's CUDA kernels need compute capability {MIN_CAPABILITY[0]}.{MIN_CAPABILITY[1]} "
- f"or newer (thread-block clusters and FP8 MMA); this GPU ({torch.cuda.get_device_name()}) "
- f"is {major}.{minor}")
+ if (major, minor) < need:
+ why = "thread-block clusters" if need >= CLUSTERS else "FP8 MMA"
+ raise RuntimeError(f"TensorFold's CUDA kernels need compute capability {need[0]}.{need[1]} or newer ({why}"
+ f"{' for these weights' if need > MIN_CAPABILITY else ''}); this GPU "
+ f"({torch.cuda.get_device_name()}) is {major}.{minor}")
return [f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"]
-def load(name: str, sources: str | list[str], **kwargs: Any) -> Any:
+def refuse_old_gpu(need: tuple[int, int] = MIN_CAPABILITY) -> None:
+ """At startup, before any weight loads: a GPU older than ``need`` is refused by name (no GPU: skipped)."""
+
+ try:
+ import torch
+ except ImportError:
+ return
+ if torch.cuda.is_available():
+ try:
+ arch_flags(need)
+ except RuntimeError as exc:
+ raise ValueError(str(exc)) from None
+
+
+def load(name: str, sources: str | list[str], need: tuple[int, int] = MIN_CAPABILITY, **kwargs: Any) -> Any:
"""torch's JIT ``load`` for this GPU only (NVIDIA's containers list every architecture back to sm_80), with a line when it compiles or waits on a lock."""
from torch.utils import cpp_extension
- kwargs["extra_cuda_cflags"] = [*kwargs.get("extra_cuda_cflags", []), *arch_flags()]
+ kwargs["extra_cuda_cflags"] = [*kwargs.get("extra_cuda_cflags", []), *arch_flags(need)]
held = _announce(cpp_extension, name, sources, kwargs.get("build_directory"))
timer = None
if held is not None:
@@ -98,4 +114,4 @@ def _say(text: str) -> None:
print(f"[tensorfold] {text}", flush=True)
-__all__ = ["MIN_CAPABILITY", "arch_flags", "load"]
+__all__ = ["CLUSTERS", "MIN_CAPABILITY", "arch_flags", "load"]
diff --git a/src/tensorfold/cuda/capacity.py b/src/tensorfold/cuda/capacity.py
index 81584791c..e646bf9f9 100644
--- a/src/tensorfold/cuda/capacity.py
+++ b/src/tensorfold/cuda/capacity.py
@@ -182,6 +182,12 @@ def available_bytes(torch) -> int:
return host if unified(torch) else min(available, host)
+def total_bytes(torch) -> int:
+ """The GPU's memory (a GB10's is the host's): the same on every rank, so what it sizes agrees without a gather."""
+
+ return int(torch.cuda.mem_get_info()[1])
+
+
def page_room(torch) -> int | None:
"""What caches and mapped read-only tables share on a unified GPU (MemAvailable); None on a discrete GPU."""
@@ -246,6 +252,19 @@ def choose(plan: Plan, peers: list[list[int]] | None = None) -> int:
return min(target, fitting)
+def floor(model_dir: str | Path) -> tuple[int, int]:
+ """The compute capability a checkpoint's kernels need: NVFP4 and FP8 (ModelOpt, compressed-tensors) use clusters."""
+
+ from tensorfold.cuda import build
+ from tensorfold.cuda.nvfp4.format import is_quantized
+
+ try:
+ quantized = is_quantized(model_dir)
+ except (OSError, ValueError): # an unreadable config is named by the estimate below
+ quantized = False
+ return build.CLUSTERS if quantized else build.MIN_CAPABILITY
+
+
def admit(model_dir: str | Path, requested: int | None, explicit: bool | None, torch,
geometry: Geometry | Callable, transform: Callable, *, rank: int = 0, world: int = 1,
gather: Callable | None = None, draft_dir: Path | None = None,
@@ -253,9 +272,11 @@ def admit(model_dir: str | Path, requested: int | None, explicit: bool | None, t
extra_files: tuple[Path, ...] = (), files: list[Path] | None = None,
draft_transform: Callable | None = None,
draft_weights: Callable[[Path], Weights] | None = None) -> dict:
- """Reach the same refusal or capacity before either rank allocates model tensors. The draft model: ``draft_weights``
- from its folder, else its tensors through ``draft_transform`` (default: 4 bytes a value, or more)."""
+ """One refusal or capacity on both ranks before allocating; the draft model by ``draft_weights`` or a transform."""
+
+ from tensorfold.cuda import build
+ build.refuse_old_gpu(floor(model_dir)) # an old GPU is refused here, before any weight loads
error = None
plan = None
try:
@@ -271,7 +292,9 @@ def admit(model_dir: str | Path, requested: int | None, explicit: bool | None, t
draft = draft_weights(draft_dir) if draft_weights is not None else estimate_weights(
draft_dir, draft_transform or (lambda name, info: (math.prod(info["shape"]) * max(4, itemsize(info, name)),
0)))
- weights = Weights(weights.resident + draft.resident, weights.staging + draft.staging, weights.mapped)
+ # the drafter loads after the target: the peak is the larger of either load's
+ weights = Weights(weights.resident + draft.resident, max(weights.staging - draft.resident, draft.staging),
+ weights.mapped)
if draft_geometry is not None:
draft_geometry = draft_geometry(config(draft_dir)) if callable(draft_geometry) else draft_geometry
main = geometry
diff --git a/src/tensorfold/cuda/geometry.py b/src/tensorfold/cuda/geometry.py
index 79ede43ad..44d799421 100644
--- a/src/tensorfold/cuda/geometry.py
+++ b/src/tensorfold/cuda/geometry.py
@@ -8,6 +8,7 @@
from .capacity import Geometry, Weights, headers, itemsize
PREFILL_ROWS = 2048 # a prompt chunk's rows: Flash Next and GLM keep buffers of this many rows
+PROMPT_SHARE = 32 # a dense prompt chunk's arrays take at most this fraction of the GPU's memory
PREFILL_ATT_ROWS = 256 # Flash Next's prompt attention block
MLA_PROMPT_ATT_ROWS = 512 # GLM's prompt-chunk rows one dense latent attention call takes (forward.PROMPT_ATT_ROWS)
MLA_SELECT_ROWS = 512 # GLM's prompt-chunk rows whose pool scores are held at once (sparse.SELECT_ROWS)
@@ -136,8 +137,29 @@ def layer_counts(t: dict) -> tuple[int, int]:
return layers - layers // interval, layers // interval
+def prompt_row_bytes(t: dict, world: int = 1) -> int:
+ """Bytes one dense prompt row's arrays hold at once, drafter taps included (the 27B measured 306-315 KiB)."""
+
+ return 16 * (int(t["hidden_size"]) + int(t["intermediate_size"]) // world)
+
+
+def prompt_rows(total: int, row_bytes: int, most: int = 4096) -> int:
+ """Prompt chunk rows: a multiple of 512 up to ``most`` whose arrays fit a PROMPT_SHARE-th of ``total``."""
+
+ return max(512, min(most, total // PROMPT_SHARE // row_bytes // 512 * 512))
+
+
+def live_kv(t: dict, world: int, window: int) -> int:
+ """A dense stream's attention caches at ``window`` rows (1,024 at least) and one layer's buffer mid-grow."""
+
+ _, attention = layer_counts(t)
+ return (attention + 1) * max(1024, window) * int(t["num_key_value_heads"]) // world * int(t["head_dim"]) * 4
+
+
def gdn_geometry(t: dict, world: int, reserve: int, *, indexed: bool = False, mtp: bool = False,
- kv_bits: int = 16) -> Geometry:
+ kv_bits: int = 16, rows: int | None = None, prompt: int = 0, evicts: bool = False) -> Geometry:
+ """``rows``: widest verify; ``prompt``: chunk rows sharing its scratch; ``evicts``: only the live window counts."""
+
linear, attention = layer_counts(t)
d, h = int(t["hidden_size"]), int(t["num_attention_heads"]) // world
hk = int(t["num_key_value_heads"]) // world
@@ -151,13 +173,13 @@ def gdn_geometry(t: dict, world: int, reserve: int, *, indexed: bool = False, mt
# Persistent state, retained recurrent prefixes, rollback and row replay inputs.
fixed = linear * ((6 if indexed else 4) * nv * dk * dv * 4 +
4 * (conv - 1) * (2 * nk * dk + nv * dv) * 2)
- rows = 64 if indexed else 128
+ rows = rows or (64 if indexed else 128)
fixed += linear * rows * (width * 2 + nk * dk * 4 + nv * dv * 4 + nv * 8)
# Bound the concurrent activation arrays, MoE expert rows, logits and split-K scratch.
slots = int(t.get("num_experts_per_tok", 1)) + 1
intermediate = int(t.get("moe_intermediate_size", t.get("intermediate_size", d))) // world
extent = d * streams + int(t["vocab_size"]) // world + slots * (intermediate + d) + width + h * hd
- fixed += 16 * rows * extent * 4
+ fixed += max(16 * rows * extent * 4, prompt * prompt_row_bytes(t, world) if prompt else 0)
fixed += (2 if mtp else 1) * 32 * rows * 2560 * 4
if indexed:
fixed += 4 * (int(t.get("ple_conv_kernel_size", 4)) - 1) * int(t.get("ngram_size", 3)) * streams * d * 2
@@ -178,7 +200,7 @@ def bytes_at(capacity: int) -> int:
else:
# Bound two retained prefixes, current KV state and a growth copy; speculative rows use separate workspace.
rounded = 1 << (max(1024, capacity - reserve) - 1).bit_length()
- cache = 4 * attention * rounded * hk * hd * 4
+ cache = live_kv(t, world, capacity - reserve) if evicts else 4 * attention * rounded * hk * hd * 4
scratch = rows * h * (hd + 2) * ((capacity + 511) // 512) * 4
return fixed + cache + scratch
return Geometry(bytes_at, reserve)
@@ -263,10 +285,7 @@ def mla_chunk_scratch(t: dict, world: int, capacity: int, *, latent: bool) -> in
def draft_ring_rows(window: int, block: int, tile: int = 64) -> int:
- """Rows of GLM's DFlash2 context ring (``dflash2.Drafter``, TF_GLM_DRAFT_RING): a block pass at context end s
- reads keys from the ``tile``-row tile holding s - window (older rows are masked for every query; ``window`` is the
- drafter's, sliding_window - 1) through its own rows s .. s + block - 1, at most window + block + tile - 1 rows,
- here rounded up to whole tiles; a kept state's window rows (window + 1) fit as well."""
+ """Rows of GLM's DFlash2 context ring: window + block + a tile in whole tiles (a kept state's window fits too)."""
return -(-(window + block + tile - 1) // tile) * tile
@@ -288,10 +307,7 @@ def bytes_at(capacity: int) -> int:
def dflash2_geometry(t: dict, world: int, reserve: int, *, ring: bool) -> Geometry:
- """GLM's DFlash2 drafter (``dflash2.Drafter``) on each of ``world`` ranks: one context of keys and values, a ring
- of ``draft_ring_rows`` rows whatever the window (``ring``, TF_GLM_DRAFT_RING; a flat buffer when the window is
- smaller) or ``capacity`` + block rows, and a block pass's activations. Kept prompt states' copies of its window
- count in the kept-state budget (``snapshot_bytes``)."""
+ """GLM's DFlash2 drafter on each rank: one context (a ring, or capacity + block rows) and a block pass."""
layers = int(t["num_hidden_layers"])
heads = int(t["num_key_value_heads"]) // world
@@ -308,12 +324,7 @@ def bytes_at(capacity: int) -> int:
def dflash2_weights(draft_dir, world: int) -> Weights:
- """What GLM's DFlash2 drafter (``dflash2.Drafter``) holds on each of ``world`` ranks, not its checkpoint's BF16:
- 4-bit copies (``qmm.quantize4``: groups of 64 with BF16 scales and biases, rows padded to 128) of fc, the
- convolutions' kernel projections and this rank's attention and MLP rows (k and v twice: in qkv and in kv); BF16
- norms, base kernels and selector projection; the selector's float32 codebooks in host memory (the same memory on
- GB10). Staging: the largest matrix read and uploaded in BF16 with its quantization temporaries (float32 groups
- of up to 8,192 rows, the packer's int64 lanes)."""
+ """GLM's DFlash2 drafter as held on each rank (4-bit copies, bf16 norms, fp32 codebooks), and its staging."""
h = headers(draft_dir)
shape = {name: [int(x) for x in info["shape"]] for name, info in h.items()}
diff --git a/src/tensorfold/cuda/kernels/qmm.cpp b/src/tensorfold/cuda/kernels/qmm.cpp
index a150d68f3..81c10c751 100644
--- a/src/tensorfold/cuda/kernels/qmm.cpp
+++ b/src/tensorfold/cuda/kernels/qmm.cpp
@@ -1,6 +1,7 @@
#include
#include
+bool qmm_clusters(int, bool);
void qmm_cuda(const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&,
at::Tensor&, const at::Tensor&, int, int, int, int, bool, bool);
void qmm_prefill_cuda(const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, at::Tensor&, int,
@@ -34,7 +35,7 @@ void qmm(const at::Tensor& x, const at::Tensor& xs, const at::Tensor& w, const a
out.scalar_type() == (f32 ? at::kFloat : at::kBFloat16), "out: (M, n)");
c10::cuda::CUDAGuard guard(x.device());
at::Tensor p;
- if (sk > 1 && (sk > 8 || !reduce)) {
+ if (sk > 1 && !qmm_clusters(static_cast(sk), reduce)) {
p = part.has_value() ? *part : at::empty({sk, m, n}, x.options().dtype(at::kFloat));
TORCH_CHECK(p.is_cuda() && p.is_contiguous() && p.scalar_type() == at::kFloat && p.numel() >= sk * m * n,
"part: at least (SK, M, n) fp32");
diff --git a/src/tensorfold/cuda/kernels/qmm.cu b/src/tensorfold/cuda/kernels/qmm.cu
index a763387be..e8afd099d 100644
--- a/src/tensorfold/cuda/kernels/qmm.cu
+++ b/src/tensorfold/cuda/kernels/qmm.cu
@@ -171,6 +171,9 @@ __global__ void __launch_bounds__(WM * WN * 32) qmm_kernel(
wait<0>();
__syncthreads();
if constexpr (CLUSTER) {
+#if __CUDA_ARCH__ < 900
+ __trap(); // no clusters before sm_90: the host never launches this
+#else
// K slices of a tile form a cluster: slice 0 adds peers' partials in slice order, as reduce_kernel does
constexpr int E = T::MT * T::NT * 4;
auto cluster = cooperative_groups::this_cluster();
@@ -198,6 +201,7 @@ __global__ void __launch_bounds__(WM * WN * 32) qmm_kernel(
}
cluster.sync(); // peers keep their memory until slice 0 has read it
if (slice != 0) return;
+#endif
}
#pragma unroll
for (int i = 0; i < T::MT; ++i)
@@ -290,12 +294,16 @@ void dispatch(int bm, const at::Tensor& x, const at::Tensor& xs, const at::Tenso
} // namespace
+// Clusters hold up to 8 K slices (the portable size) from sm_90; more, unreduced slices or older GPUs use the buffer.
+bool qmm_clusters(int SK, bool reduce) {
+ return SK > 1 && SK <= 8 && reduce && at::cuda::getCurrentDeviceProperties()->major >= 9;
+}
+
void qmm_cuda(const at::Tensor& x, const at::Tensor& xs, const at::Tensor& w, const at::Tensor& scales,
const at::Tensor& biases, at::Tensor& out, const at::Tensor& part, int N, int SK, int gs, int bm,
bool f32, bool reduce) {
const int M = x.size(0);
- // clusters hold up to 8 K slices (the portable size); more, or unreduced slices, go through the slice buffer
- const bool cluster = SK > 1 && SK <= 8 && reduce;
+ const bool cluster = qmm_clusters(SK, reduce);
#define GO(G, F, C) dispatch(bm, x, xs, w, scales, biases, out, part, N, SK)
if (gs == 64) {
if (f32) { if (cluster) GO(64, true, true); else GO(64, true, false); }
diff --git a/src/tensorfold/cuda/kernels/qmm_frag.cuh b/src/tensorfold/cuda/kernels/qmm_frag.cuh
index e432e3bad..8dbd55ce8 100644
--- a/src/tensorfold/cuda/kernels/qmm_frag.cuh
+++ b/src/tensorfold/cuda/kernels/qmm_frag.cuh
@@ -60,7 +60,12 @@ __device__ __forceinline__ void mma0(float (&d)[4], const uint32_t (&a)[4], uint
__device__ __forceinline__ uint32_t pair(uint32_t w, int s) {
const uint32_t t = ((w >> s) & 0x000F000Fu) | 0x43004300u;
uint32_t r;
+#if __CUDA_ARCH__ >= 900
asm("sub.rn.bf16x2 %0, %1, %2;\n" : "=r"(r) : "r"(t), "r"(0x43004300u));
+#else
+ // no sub.bf16x2 before sm_90: t * 1 - 128 rounds the exact q once too (an exact value, the same bits)
+ asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(r) : "r"(t), "r"(0x3F803F80u), "r"(0xC300C300u));
+#endif
return r;
}
diff --git a/src/tensorfold/cuda/nvfp4/linear.py b/src/tensorfold/cuda/nvfp4/linear.py
index f09300e7f..9909819fa 100644
--- a/src/tensorfold/cuda/nvfp4/linear.py
+++ b/src/tensorfold/cuda/nvfp4/linear.py
@@ -15,11 +15,11 @@
@lru_cache(maxsize=1)
def _ext():
- from tensorfold.cuda.build import load
+ from tensorfold.cuda.build import CLUSTERS, load
here = Path(__file__).parent
return load(name="tensorfold_nvfp4_v3", sources=[str(here / "qmmf.cpp"), str(here / "qmmf.cu"),
- str(here / "experts.cu")],
+ str(here / "experts.cu")], need=CLUSTERS,
extra_include_paths=[str(here)], extra_cuda_cflags=["-O3"], verbose=False)
diff --git a/src/tensorfold/cuda/streams.py b/src/tensorfold/cuda/streams.py
index c93e757f4..07492c801 100644
--- a/src/tensorfold/cuda/streams.py
+++ b/src/tensorfold/cuda/streams.py
@@ -147,12 +147,13 @@ def add(self, ids: list[int], state: Any, snap: Any) -> None:
while len(self.entries) > self.keep:
self._drop(self.entries[:-1])
- def evict(self) -> bool:
- """Memory is short: drop the entry ``add`` would drop next; False when none is left."""
+ def evict(self, among: list | None = None) -> bool:
+ """Memory is short: drop the entry ``add`` would drop next (of ``among``); False when none is left."""
- if not self.entries:
+ among = self.entries if among is None else among
+ if not among:
return False
- self._drop(self.entries)
+ self._drop(among)
return True
def _drop(self, among: list) -> None:
@@ -160,3 +161,27 @@ def _drop(self, among: list) -> None:
gone = cold[0] if cold else among[0]
self.entries = [e for e in self.entries if e is not gone]
self.hit &= {tuple(e[0]) for e in self.entries}
+
+
+class KVRoom:
+ """One GPU's attention-cache bytes: a grow first evicts kept entries on other buffers, least recently used."""
+
+ def __init__(self, cache: PrefixCache, budget: int) -> None:
+ self.cache, self.budget = cache, int(budget)
+
+ def __call__(self, st: Any, extra: int) -> None:
+ while self.held(st) + extra > self.budget:
+ if not self.cache.evict([e for e in self.cache.entries if e[1].kv is not st.kv]):
+ return # only this conversation is left: the window was admitted for it
+
+ def held(self, st: Any) -> int:
+ """Bytes of every distinct attention buffer the state and the kept entries hold."""
+
+ seen, total = set(), 0
+ for kv in [st.kv, *(e[1].kv for e in self.cache.entries)]:
+ for pair in kv:
+ for t in pair or ():
+ if t.data_ptr() not in seen:
+ seen.add(t.data_ptr())
+ total += t.untyped_storage().nbytes()
+ return total
diff --git a/src/tensorfold/families/glm5_next/cuda/decode.py b/src/tensorfold/families/glm5_next/cuda/decode.py
index e5f0d4000..f8b4ec4c1 100644
--- a/src/tensorfold/families/glm5_next/cuda/decode.py
+++ b/src/tensorfold/families/glm5_next/cuda/decode.py
@@ -258,8 +258,7 @@ def take_snapshot(e: Engine, ids: Sequence[int], pending: torch.Tensor | None, *
def _ring_slots(drafter, n: int) -> torch.Tensor:
- """Where a ring drafter (``drafter.ring`` rows) holds the window rows a block pass at context end n reads
- (positions n - window - 1 .. n - 1, one spare): those positions modulo the ring."""
+ """The ring slots of the window rows a block pass at context end n reads (n - window - 1 .. n - 1)."""
lo = max(0, n - drafter.window - 1)
return torch.arange(lo, n, device=drafter.kc[0].device) % drafter.ring
@@ -311,8 +310,7 @@ def row_bytes(e: Engine, snap: Snapshot) -> int:
def snapshot_bytes(snap: Snapshot) -> int:
- """Device memory a kept snapshot holds: its KDA states, conv windows, pending MTP rows, a ring drafter's window
- and any saved rows."""
+ """A kept snapshot's device bytes: KDA states, conv windows, pending MTP rows, a ring window, saved rows."""
held = [snap.rec, snap.conv] + ([snap.pending] if snap.pending is not None else []) + (snap.drafter_rows or [])
return sum(t.numel() * t.element_size() for t in held) + (snap.nbytes if snap.rows is not None else 0)
@@ -387,6 +385,12 @@ def prefill(e: Engine, prompt: Sequence[int], sampling: Sampling | None, *, mtp:
keep_at - 1 if use_mtp else -1, keep_at if drafter is not None else -1)
if drafter is not None:
drafter.add_taps(e.tap_rows(R, b))
+ if 0 < point <= R and getattr(drafter, "ring", 0):
+ # a ring keeps the kept point's window unless this chunk wrote past it by more than the ring's slack
+ if R - point < drafter.ring - drafter.window:
+ kept.drafter_rows = _ring_window(drafter, keep_at)
+ else:
+ kept.drafter_end = -1
if use_mtp:
nxt = list(prompt[start + 1:start + R + 1])
if nxt:
diff --git a/src/tensorfold/families/glm5_next/cuda/engine.py b/src/tensorfold/families/glm5_next/cuda/engine.py
index 53beb20f0..b98d846ce 100644
--- a/src/tensorfold/families/glm5_next/cuda/engine.py
+++ b/src/tensorfold/families/glm5_next/cuda/engine.py
@@ -17,9 +17,7 @@
GRAPH_ROWS = (1, 2, 3, 4, 5, 6) # verify windows captured as CUDA graphs
MAX_ROWS = 8 # the widest verify window (a pending token and up to 7 drafts)
DENSE_CAPACITY = 2560 # cache slots while DSA attention stays dense (contexts up to 2,051 tokens)
-# TF_GLM_DRAFT_RING=0: DFlash2 keeps its context in a flat buffer of the whole window (10,240 bytes a slot on each
-# rank) instead of a ring of its 2,048-row sliding window and block (2,176 rows, 21.25 MiB; ``dflash2.Drafter``);
-# the drafts are the same bits either way
+# TF_GLM_DRAFT_RING=0: DFlash2's context in a flat buffer of the whole window, not a 2,176-row ring (the same drafts)
DRAFT_RING = os.environ.get("TF_GLM_DRAFT_RING", "1").strip() != "0"
@@ -369,9 +367,9 @@ def _drop(self, snap) -> None:
self.cache.remove(snap)
def _remember(self, snap) -> None:
- for c in [c for c in self.cache if c.ids == snap.ids]:
+ for c in [c for c in self.cache if c.ids == snap.ids and c is not snap]:
self._drop(c)
- self.cache.append(snap)
+ self.cache[:] = [c for c in self.cache if c is not snap] + [snap] # a resumed prompt kept again moves last
dropped = False
while len(self.cache) > 1 and (len(self.cache) > self.cache_entries or self._held_bytes() > self.cache_bytes):
self._drop(self.cache[0])
diff --git a/src/tensorfold/families/qwen3_5/cuda/affine_memory.py b/src/tensorfold/families/qwen3_5/cuda/affine_memory.py
index 053a150e5..bb6bb6d49 100644
--- a/src/tensorfold/families/qwen3_5/cuda/affine_memory.py
+++ b/src/tensorfold/families/qwen3_5/cuda/affine_memory.py
@@ -8,7 +8,9 @@
from pathlib import Path
-def weight_transform(model_dir):
+def weight_transform(model_dir, *, one_gpu: bool = False):
+ """Tensor bytes; the head counts twice (the drafter's copy), once on ``one_gpu`` where the drafter views it."""
+
from tensorfold.cuda.geometry import linear_weights, size
from tensorfold.cuda.capacity import headers
from tensorfold.quantization import resolve_affine
@@ -18,9 +20,16 @@ def metadata():
path = Path(model_dir)
return json.loads((path / "config.json").read_text()), headers(path)
- def transform(name, info):
- if name.startswith("vision_tower") or ".mtp." in name or name.startswith("mtp."):
- return 0, 0
+ def tiled(path: str) -> bool:
+ """4-bit words in groups of 64 with bf16 scales and biases: what loading tiles."""
+
+ config, tensors = metadata()
+ spec = resolve_affine(config, path)
+ scales, biases = (tensors.get(path + suffix) for suffix in (".scales", ".biases"))
+ return (spec is not None and scales is not None and biases is not None and
+ (spec.bits, spec.group_size, scales["dtype"], biases["dtype"]) == (4, 64, "BF16", "BF16"))
+
+ def stored(name, info):
if info["dtype"] not in ("U32", "I32") or not name.endswith(".weight"):
return linear_weights(name, info)
config, tensors = metadata()
@@ -34,9 +43,17 @@ def transform(name, info):
from tensorfold.quantization import validate_shapes
validate_shapes(info["shape"], scales["shape"], biases["shape"], spec)
- if (spec.bits, spec.group_size, scales["dtype"], biases["dtype"]) == (4, 64, "BF16", "BF16"):
+ if tiled(path):
return linear_weights(name, info)
return size(info) * (2 if "lm_head." in name else 1), 0
+
+ def transform(name, info):
+ if name.startswith("vision_tower") or ".mtp." in name or name.startswith("mtp."):
+ return 0, 0
+ amount, host = stored(name, info)
+ if one_gpu and "lm_head." in name and tiled(name.rsplit(".", 1)[0]):
+ amount //= 2 # the drafter's rows are views of the target's head
+ return amount, host
return transform
@@ -47,6 +64,17 @@ def packed_draft(name: str, shape) -> bool:
and shape[0] * shape[1] >= 1 << 20)
+def draft_weights(draft_dir):
+ """The drafter's bytes and load peak: quantize4 holds 14 bytes an element of its largest tensor while packing."""
+
+ from tensorfold.cuda.capacity import Weights, estimate_weights, headers
+
+ held = estimate_weights(draft_dir, draft_bytes)
+ largest = max((math.prod(info["shape"]) for name, info in headers(draft_dir).items()
+ if packed_draft(name, info["shape"])), default=0)
+ return Weights(held.resident, max(held.staging, 14 * largest), held.mapped)
+
+
def draft_bytes(name: str, info: dict) -> tuple[int, int]:
"""GPU bytes of a draft tensor as the 4-bit drafter holds it (q4 words and scales; k and v twice)."""
diff --git a/src/tensorfold/families/qwen3_5/cuda/decode.py b/src/tensorfold/families/qwen3_5/cuda/decode.py
index 6bef6f06c..e6b696060 100644
--- a/src/tensorfold/families/qwen3_5/cuda/decode.py
+++ b/src/tensorfold/families/qwen3_5/cuda/decode.py
@@ -16,14 +16,14 @@
def clone_state(st: State) -> State:
- """The committed tensors are immutable; commits replace their list entries."""
+ """Commits replace list entries; the attention list is shared, so a grow reaches every clone and frees the old."""
other = object.__new__(State)
other.pos, other.limit = st.pos, st.limit
- other.rope_delta = st.rope_delta
+ other.rope_delta, other.room = st.rope_delta, st.room
other.conv = st.conv.copy()
other.rec = st.rec.copy()
- other.kv = st.kv.copy()
+ other.kv = st.kv
return other
@@ -49,7 +49,7 @@ def prefill_stops(w: Weights, prompt: Sequence[int], st: State, draft=None, *, s
@torch.no_grad()
def prefill(w: Weights, prompt: Sequence[int], sampling: Sampling | None,
draft=None, *, state: State | None = None, limit: int = 0, stops: Sequence[int] = (),
- keep: Callable | None = None, keep_at: int | None = None, vision=None, constraint=None):
+ keep: Callable | None = None, keep_at: int | None = None, vision=None, constraint=None, room=None):
"""Commit the prompt and sample the first token; resuming a kept ``state`` gives a fresh prefill's bits (``keep_at`` adds a third item: the state after prompt[:keep_at] and the drafter's snapshot there)."""
from .forward import _mm
@@ -59,6 +59,8 @@ def prefill(w: Weights, prompt: Sequence[int], sampling: Sampling | None,
st = clone_state(state) if state is not None else State(w)
if state is None:
st.limit = limit # a fresh state's attention caches stop here; a resumed one keeps its own
+ if room is not None:
+ st.room = room # before a grow, the engine frees other conversations' kept buffers
if st.pos >= len(prompt):
raise ValueError("a reused state must leave at least one prompt token to process")
out = prefill_stops(w, prompt, st, draft, stops=stops, keep=keep, keep_at=keep_at, vision=vision)
diff --git a/src/tensorfold/families/qwen3_5/cuda/engine.py b/src/tensorfold/families/qwen3_5/cuda/engine.py
index aa079260b..779ad9b8a 100644
--- a/src/tensorfold/families/qwen3_5/cuda/engine.py
+++ b/src/tensorfold/families/qwen3_5/cuda/engine.py
@@ -22,6 +22,7 @@ class Qwen27Engine:
"""Qwen3.8-27B on one GPU or two ranks (rank 0 here), DFlash2 drafting, prefix reuse."""
tree_rows: int | None = None # a lone stream's tree rows on one GPU (None: max_rows, as in 0.5.0)
+ room = None # one GPU's attention-cache budget (streams.KVRoom), None on two
def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 12, tp: int = 1,
rank: int = 0, master: str = "", port: int = 29551, split_head: bool = False,
@@ -43,9 +44,10 @@ def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 1
raise ValueError("image input on CUDA is tested on the MLX checkpoint only: drop --vision for an NVFP4 "
"checkpoint, or serve Vontra/Qwen3.8-27B-MLX-4bit")
from .weights import load
- from tensorfold.cuda.capacity import admit, gather_ints
- from tensorfold.cuda.geometry import draft_geometry, gdn_geometry, stream_geometry
- from .affine_memory import draft_bytes, weight_transform
+ from tensorfold.cuda.capacity import admit, config, gather_ints, total_bytes
+ from tensorfold.cuda.geometry import (draft_geometry, gdn_geometry, live_kv, prompt_row_bytes, prompt_rows,
+ stream_geometry)
+ from .affine_memory import draft_weights, weight_transform
from tensorfold.vision.qwen_cuda import capacity_geometry, weight_transform as vision_weights
self.torch = torch
@@ -81,10 +83,12 @@ def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 1
else:
gather = None
many = streams > 1
+ # prompt chunks sized to the card (4096 rows from 80 GB), the verify scratch to the rows a round takes
+ chunk = prompt_rows(total_bytes(torch), prompt_row_bytes(config(model_dir), tp))
geometry = ((lambda text: stream_geometry(text, tp, streams, keep, first=256 if tp == 1 else None)) if many
- else (lambda text: gdn_geometry(text, tp, max_rows)))
+ else (lambda text: gdn_geometry(text, tp, max_rows, rows=max_rows, prompt=chunk, evicts=tp == 1)))
# an affine checkpoint's packed words at their stored precision; an EXL3 pack's by its own format
- tensor_bytes = weight_transform(model_dir)
+ tensor_bytes = weight_transform(model_dir, one_gpu=tp == 1)
if exl3:
geometry, tensor_bytes = admission(geometry)
elif nvfp4:
@@ -97,7 +101,7 @@ def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 1
vision_weights(tensor_bytes, vision, rank),
rank=rank, world=tp, gather=gather,
draft_dir=draft_dir if rank == 0 or tp_draft else None,
- draft_transform=draft_bytes,
+ draft_weights=draft_weights,
draft_geometry=lambda text: draft_geometry(text, tp if tp_draft else 1, max_rows,
bounded=True, streams=streams,
kept=keep + 1 if many else 0),
@@ -109,6 +113,7 @@ def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 1
else:
full = load(model_dir, tiled=True)
self.w = full
+ self.w.prompt_rows = chunk
self.draft = None
if draft_dir is not None and (rank == 0 or (tp == 2 and tp_draft)):
from .dflash2 import DFlash2
@@ -128,6 +133,12 @@ def __init__(self, model_dir: Path, draft_dir: Path | None, *, max_rows: int = 1
self.model_dir = Path(model_dir)
self.points = resume_points(model_dir) # message starts a prefill keeps states at
self.cache = PrefixCache(KEEP_ONE) # (committed ids, state, drafter snapshot)
+ if tp == 1 and not many: # one GPU: kept buffers stay inside the bytes admission left the caches
+ from tensorfold.cuda.streams import KVRoom
+
+ plan = self.capacity_plan
+ spare = plan["budget_bytes"] - plan["weight_bytes_estimate"] - plan["cache_workspace_bytes_estimate"]
+ self.room = KVRoom(self.cache, spare + live_kv(config(model_dir), 1, self.context_window))
# ``streams`` > 1: up to that many requests decoded together, their windows verified in one forward
self.concurrent = streams > 1
self.multi = self.scheduler = None
@@ -216,7 +227,7 @@ def generate(self, prompt: list[int], max_tokens: int, sampling, on_tokens: Call
drafter.restore(([None] * drafter.layers, [None] * drafter.layers, 0, 0))
stops, keep = self._stops(prompt, hit, draft) if vision is None else ((), None)
end = entry_end(prompt) if draft and vision is None and self._ends(prompt, stops) else None
- st, pending, *kept = prefill(self.w, prompt, sampling, drafter, state=hit[1] if hit else None,
+ st, pending, *kept = prefill(self.w, prompt, sampling, drafter, state=hit[1] if hit else None, room=self.room,
limit=self.context_window, stops=stops, keep=keep, keep_at=end, vision=encoded,
**grammar)
if end is not None:
diff --git a/src/tensorfold/families/qwen3_5/cuda/forward.py b/src/tensorfold/families/qwen3_5/cuda/forward.py
index 7aead9d45..e83c64906 100644
--- a/src/tensorfold/families/qwen3_5/cuda/forward.py
+++ b/src/tensorfold/families/qwen3_5/cuda/forward.py
@@ -98,9 +98,10 @@ class AttentionRecord:
class State:
- """Cloned states share growable KV buffers whose rows below ``pos`` remain committed, so writes preserve shorter clones but invalidate longer cached extensions."""
+ """Clones share one list of KV buffers: writes keep shorter clones' rows and invalidate longer cached ones."""
rope_delta = 0 # an image prompt's rotary shift past its tokens; text has none
+ room = None # called (state, bytes) before a grow: one GPU's engine frees kept buffers
def __init__(self, w: Weights):
c = w.config
@@ -162,14 +163,28 @@ def stage(w: Weights, st: State, width: int, context: int) -> Staged:
_cache_offsets([st], softmax, device), _conv_windows(parents, c.conv_kernel - 1).to(device), host, dev)
+def grow(st: State, i: int, need: int, *, exact: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+ """Layer ``i``'s cache at ``need`` rows (doubled to ``st.limit``, or exact): every row moves, the old ones free."""
+
+ k, v = st.kv[i]
+ if k.shape[0] < need:
+ cap = need if exact else max(need, 2 * k.shape[0], 1024)
+ if st.limit and not exact:
+ cap = max(need, min(cap, st.limit))
+ if st.room is not None:
+ st.room(st, cap * (k.stride(0) * k.element_size() + v.stride(0) * v.element_size()))
+ grown = k.new_empty((cap, *k.shape[1:])), v.new_empty((cap, *v.shape[1:]))
+ grown[0][:k.shape[0]], grown[1][:v.shape[0]] = k, v
+ st.kv[i] = grown
+ return st.kv[i]
+
+
def reserve(st: State, rows: int) -> None:
"""Grow every attention cache to ``rows`` now, so later commits never move a buffer (graphs keep addresses)."""
for i, kv in enumerate(st.kv):
- if kv is not None and kv[0].shape[0] < rows:
- k, v = kv[0].new_empty((rows, *kv[0].shape[1:])), kv[1].new_empty((rows, *kv[1].shape[1:]))
- k[:st.pos], v[:st.pos] = kv[0][:st.pos], kv[1][:st.pos]
- st.kv[i] = (k, v)
+ if kv is not None:
+ grow(st, i, rows, exact=True)
st.limit = rows
@@ -450,16 +465,7 @@ def _commit(states: Sequence[State], record: Sequence[Record], paths: Sequence[S
for st, path in zip(states, paths):
need = st.pos + len(path)
for i, _ in att:
- kbuf, vbuf = st.kv[i]
- if kbuf.shape[0] < need:
- cap = max(need, 2 * kbuf.shape[0], 1024)
- if st.limit:
- cap = max(need, min(cap, st.limit))
- grown_k = kbuf.new_empty((cap, *kbuf.shape[1:]))
- grown_v = vbuf.new_empty((cap, *vbuf.shape[1:]))
- grown_k[:st.pos] = kbuf[:st.pos]
- grown_v[:st.pos] = vbuf[:st.pos]
- st.kv[i] = (grown_k, grown_v)
+ grow(st, i, need)
# every stream's accepted key/value rows: one gather a layer, one multi-tensor copy into every cache
take = takes[0] if len(takes) == 1 else torch.cat(list(takes))
dst, src = [], []
diff --git a/src/tensorfold/families/qwen3_5/cuda/multi.py b/src/tensorfold/families/qwen3_5/cuda/multi.py
index 31787d1bd..bd7e91ee9 100644
--- a/src/tensorfold/families/qwen3_5/cuda/multi.py
+++ b/src/tensorfold/families/qwen3_5/cuda/multi.py
@@ -32,6 +32,7 @@ def private(st: State, rows: int) -> State:
"""A copy of a committed state with its own attention caches of ``rows`` rows (rows below ``pos`` copied in)."""
other = clone_state(st)
+ other.kv = list(st.kv) # its own list: the buffers reserved next are this stream's
reserve(other, rows) # new buffers now: prefill never grows or reallocates them
return other
diff --git a/src/tensorfold/families/qwen3_5/cuda/prefill.py b/src/tensorfold/families/qwen3_5/cuda/prefill.py
index 340565f51..a41463810 100644
--- a/src/tensorfold/families/qwen3_5/cuda/prefill.py
+++ b/src/tensorfold/families/qwen3_5/cuda/prefill.py
@@ -14,7 +14,7 @@
from . import glue
from . import prefill_bf16, prefill_glue
from .decode import clone_state
-from .forward import State
+from .forward import State, grow as _grow
from .qmm_fast import matmul, matmul_partial, tile
from .weights import QLinear, Weights
@@ -43,19 +43,6 @@ def _row_mm(x, w: QLinear, tp: bool) -> torch.Tensor:
return gather_rank_partials(_mm(x, w)) # bf16 partials: half the bytes of fp32 over the link
-def _grow(st: State, i: int, need: int) -> tuple[torch.Tensor, torch.Tensor]:
- kbuf, vbuf = st.kv[i]
- if kbuf.shape[0] < need:
- cap = max(need, 2 * kbuf.shape[0], 1024)
- if st.limit:
- cap = max(need, min(cap, st.limit))
- grown_k, grown_v = kbuf.new_empty((cap, *kbuf.shape[1:])), vbuf.new_empty((cap, *vbuf.shape[1:]))
- grown_k[:st.pos] = kbuf[:st.pos]
- grown_v[:st.pos] = vbuf[:st.pos]
- st.kv[i] = (grown_k, grown_v)
- return st.kv[i]
-
-
@torch.no_grad()
def prefill_chunk(w: Weights, tokens: torch.Tensor, st: State, *, tp: bool = False, capture_taps: bool = False,
last: bool = True, every: bool = False, cut: int = 0, vision=None):
@@ -80,9 +67,7 @@ def prefill_chunk(w: Weights, tokens: torch.Tensor, st: State, *, tp: bool = Fal
x = replace_rows(x, vision, p0, p0 + W)
pending: torch.Tensor | None = None
taps: list[torch.Tensor] = []
- part = clone_state(st) if cut else None
- if part is not None:
- part.kv = [] # the chunk's final buffers, set below: these would outlive a grow that replaces them
+ part = clone_state(st) if cut else None # its attention buffers are the chunk's, through the shared list
for i, layer in enumerate(w.layers):
x, h = pg.add_rmsnorm(x, pending, layer.input_norm, c.eps)
if layer.linear:
@@ -146,7 +131,7 @@ def prefill_chunk(w: Weights, tokens: torch.Tensor, st: State, *, tp: bool = Fal
taps_out = torch.cat(taps, dim=-1) if capture_taps else None
if part is None:
return normed, taps_out
- part.pos, part.kv = p0 + cut, st.kv.copy() # the chunk's buffers: their rows below part.pos stay as committed
+ part.pos = p0 + cut # the chunk's buffers: their rows below part.pos stay as committed
return normed, taps_out, part
@@ -159,7 +144,7 @@ def chunks(start: int, end: int, size: int = CHUNK) -> list[tuple[int, int]]:
@torch.no_grad()
def prefill_state(w: Weights, prompt: Sequence[int], st: State, *, tp: bool = False, draft=None,
- size: int = CHUNK, keep_at: int | None = None, vision=None):
+ size: int | None = None, keep_at: int | None = None, vision=None):
"""Commit prompt[st.pos:] into ``st``, tapping the drafter's window; ``keep_at``: ``(normed, (state, snapshot))``, the state after prompt[:keep_at] from a cut chunk."""
dev = w.norm.device
@@ -177,11 +162,10 @@ def prefill_state(w: Weights, prompt: Sequence[int], st: State, *, tp: bool = Fa
if draft is not None and end - draft.window > base:
tap_from = end - draft.window
draft.skip(tap_from - base)
- spans = chunks(base, n, size)
+ spans = chunks(base, n, size or getattr(w, "prompt_rows", CHUNK)) # stand-in weights take 4096
for j, (a, b) in enumerate(spans):
if keep_at == a:
kept = (clone_state(st), draft.snapshot() if draft is not None else None)
- kept[0].kv = [] # the final buffers, set below, as for a cut
cut = keep_at - a if keep_at is not None and a < keep_at < b else 0
want = draft is not None and b > tap_from
normed, taps, *part = prefill_chunk(w, ids[a - base:b - base], st, tp=tp, capture_taps=want,
@@ -201,5 +185,4 @@ def prefill_state(w: Weights, prompt: Sequence[int], st: State, *, tp: bool = Fa
return normed
if keep_at == n:
kept = (clone_state(st), draft.snapshot() if draft is not None else None)
- kept[0].kv = st.kv.copy() # grown buffers copy the committed rows: never hold the old ones
return normed, kept
diff --git a/src/tensorfold/families/qwen3_5/cuda/weights.py b/src/tensorfold/families/qwen3_5/cuda/weights.py
index affecb414..ab0192c92 100644
--- a/src/tensorfold/families/qwen3_5/cuda/weights.py
+++ b/src/tensorfold/families/qwen3_5/cuda/weights.py
@@ -209,6 +209,7 @@ class Weights:
head: Any # QLinear, Exl3, or an NVFP4 checkpoint's linear
inv_freq: torch.Tensor | None = None # (rope_dims/2,) fp32
quant: str = "mlx" # "exl3": an EXL3 pack (prompt glue then stays in bf16); "nvfp4"
+ prompt_rows: int = 4096 # a prompt chunk's rows, sized to the GPU: any count, the same bits
@cached_property
def fast_prefill(self) -> bool:
diff --git a/tests/cuda/test_glm_draft_ring.py b/tests/cuda/test_glm_draft_ring.py
index d4775249a..802c7290a 100644
--- a/tests/cuda/test_glm_draft_ring.py
+++ b/tests/cuda/test_glm_draft_ring.py
@@ -1,10 +1,4 @@
-"""DFlash2's block attention from the window's first tile, and its context ring (TF_GLM_DRAFT_RING), against the
-loop over every context tile from 0 and the flat buffer, bit for bit: the attention kernel over contexts far past the
-window (the ring wrapped many times, its other rows garbage), and the synthetic drafter of test_glm_engine end to end
-(eager and CUDA graphs): the same candidates, logits and selector rows at every round over contexts longer than the
-window, after a kept state is taken, the ring overwritten and the state restored.
-
-Small: one-layer drafter, no engine (no admission), a few MiB of caches."""
+"""DFlash2's windowed block attention and context ring against the loop from tile 0 and a flat buffer, bitwise."""
from __future__ import annotations
@@ -89,8 +83,7 @@ def pair(request, tmp_path_factory):
def _drive(pair, seed: int) -> int:
- """Prompt chunks and decode rounds far past the window on both drafters, a kept state taken, the context run on
- past it (the ring overwritten), restored and run on again; every round's candidates compared bit for bit."""
+ """Rounds far past the window on both drafters, a kept state taken, overwritten, restored; candidates compared."""
from tensorfold.families.glm5_next.cuda import decode
@@ -150,3 +143,98 @@ def test_the_ring_drafts_the_flat_buffers_bits_in_cuda_graphs(pair):
finally:
for d in pair:
d.block_graph, d.tap_graphs = None, {}
+
+
+@pytest.fixture(scope="module")
+def engine_ring(tmp_path_factory):
+ """The synthetic GLM engine with DFlash2 in a 128-row ring (window 48) and 128-row prompt chunks."""
+
+ from tensorfold.families.glm5_next.cuda.engine import GlmEngine
+ from test_glm_engine import _TwoCopies, _checkpoint
+
+ path = tmp_path_factory.mktemp("glm_ring")
+ _checkpoint(path / "model")
+ _drafter(path / "dflash2")
+ (path / "dflash2" / "config.json").write_text(json.dumps(dict(DRAFT, sliding_window=49)))
+ engine = GlmEngine(path / "model", rank=0, master="", port=0, drafter=path / "dflash2", comm=_TwoCopies(),
+ prefill_rows=128)
+ assert engine.drafter.ring == draft_ring_rows(48, engine.drafter.block) == 128
+ yield engine
+ del engine
+ torch.cuda.empty_cache()
+
+
+def _rounds(drafter, pending: int, seed: int, k: int = 12) -> list:
+ """k rounds of the drafter's candidates, the same taps between them for every caller with this seed."""
+
+ rng = np.random.default_rng(seed)
+ gen = torch.Generator(device="cuda").manual_seed(seed)
+ out = []
+ for _ in range(k):
+ out.append(drafter.candidates(pending, int(rng.integers(1, drafter.block))))
+ n = int(rng.integers(1, 9))
+ drafter.add_taps((torch.randn((n, drafter.tap_in.shape[1]), generator=gen, device="cuda") * 0.5).bfloat16())
+ pending = int(rng.integers(0, 1000))
+ return out
+
+
+def _same(a: list, b: list) -> bool:
+ return len(a) == len(b) and all(all(np.array_equal(x, y) for x, y in zip(r, s)) for r, s in zip(a, b))
+
+
+def test_a_kept_prompt_point_resumes_the_rings_drafts_round_by_round(engine_ring):
+ """A prompt kept one token early (the engine's len - 1) in a ring wrapped past its window: an identical resend and
+ a next turn resumed from it draft, round by round, what a fresh prefill of the same prompt drafts."""
+
+ from tensorfold.families.glm5_next.cuda import decode
+
+ e, drafter = engine_ring.e, engine_ring.drafter
+ prompt = [int(t) for t in np.random.default_rng(41).integers(0, 1000, size=300)] # 128 + 128 + 44 rows
+ kept = []
+ first = decode.prefill(e, prompt, None, drafter=drafter, keep_at=len(prompt) - 1, keep=kept.append)
+ fresh = _rounds(drafter, first, 1)
+ snap = kept[-1]
+ assert snap.drafter_end == len(prompt) - 1 and snap.drafter_rows is not None
+ again = decode.prefill(e, prompt, None, drafter=drafter, resume=snap, keep_at=len(prompt) - 1, keep=kept.append)
+ assert again == first and _same(_rounds(drafter, again, 1), fresh) # identical resend
+ turn = prompt[:-1] + [271, 77, 78]
+ first = decode.prefill(e, turn, None, drafter=drafter, resume=kept[-1], keep_at=len(turn) - 1, keep=kept.append)
+ resumed = _rounds(drafter, first, 2)
+ assert decode.prefill(e, turn, None, drafter=drafter) == first # next turn, fresh
+ assert _same(_rounds(drafter, first, 2), resumed)
+
+
+def test_a_kept_point_its_chunk_wrote_past_keeps_no_ring_window(engine_ring):
+ """Past the ring's slack behind the chunk's end the kept point's window is gone: DFlash2 won't resume from it."""
+
+ from tensorfold.families.glm5_next.cuda import decode
+
+ e, drafter = engine_ring.e, engine_ring.drafter
+ rows, slack = e.prefill_rows, drafter.ring - drafter.window
+ prompt = [int(t) for t in np.random.default_rng(43).integers(0, 1000, size=rows)]
+ kept = []
+ decode.prefill(e, prompt, None, drafter=drafter, keep_at=rows - slack, keep=kept.append)
+ assert kept[-1].drafter_end == -1 and kept[-1].drafter_rows is None
+ decode.prefill(e, prompt, None, drafter=drafter, keep_at=rows - slack + 1, keep=kept.append)
+ assert kept[-1].drafter_end == rows - slack + 1 and kept[-1].drafter_rows is not None
+
+
+def test_identical_resends_and_a_turn_keep_the_rings_window_through_the_engine(engine_ring):
+ """Each identical resend keeps the state it resumed from again (the same snapshot), which must keep its window."""
+
+ from test_glm_engine import _forget, _generate
+
+ _forget(engine_ring)
+ prompt = [int(t) for t in np.random.default_rng(47).integers(0, 1000, size=300)]
+ fresh, stats = _generate(engine_ring, prompt, None, policy="fc5:0.3", tokens=24)
+ want = (fresh, stats["rounds"], stats["min_rows"])
+ for _ in range(3):
+ again, stats = _generate(engine_ring, prompt, None, policy="fc5:0.3", tokens=24)
+ assert stats["cached"] == len(prompt) - 1 and (again, stats["rounds"], stats["min_rows"]) == want
+ turn = prompt[:-1] + [271, 77, 78]
+ resumed, stats = _generate(engine_ring, turn, None, policy="fc5:0.3", tokens=24)
+ assert stats["cached"] == len(prompt) - 1
+ got = (resumed, stats["rounds"], stats["min_rows"])
+ _forget(engine_ring)
+ fresh, stats = _generate(engine_ring, turn, None, policy="fc5:0.3", tokens=24)
+ assert stats["cached"] == 0 and (fresh, stats["rounds"], stats["min_rows"]) == got
diff --git a/tests/cuda/test_nemotron_app.py b/tests/cuda/test_nemotron_app.py
index dd5e9b964..18facbd29 100644
--- a/tests/cuda/test_nemotron_app.py
+++ b/tests/cuda/test_nemotron_app.py
@@ -52,7 +52,7 @@ def test_resumed_prompts_equal_fresh(engine):
for prompt in (first + _ids(" Also give its complexity."), first + reply + _ids(" Now in C.")):
resumed: list[int] = []
stats = engine.generate(prompt, 32, sampling, lambda new: resumed.extend(new))
- assert stats["cached"] == len(first) # prompt ends only: a reply prefills again
+ assert stats["cached"] == len(first) - 1 # kept one token early (#98): a reply prefills again
engine.cache = [] # fresh: nothing to resume from
fresh: list[int] = []
engine.generate(prompt, 32, sampling, lambda new: fresh.extend(new))
diff --git a/tests/cuda/test_qmm.py b/tests/cuda/test_qmm.py
index 9417582b8..5284451c4 100644
--- a/tests/cuda/test_qmm.py
+++ b/tests/cuda/test_qmm.py
@@ -93,3 +93,19 @@ def test_strided_rows_buffers_and_unreduced_slices():
for s in range(1, sk):
total = total + slices[s]
assert torch.equal(total, qmm.matmul(x, q, f32=True))
+
+
+@pytest.mark.parametrize("n,k,gs", [(128, 5120, 64), (200, 1024, 32)])
+def test_every_nibble_decodes_exactly(n, k, gs):
+ """One-hot rows read back each stored nibble (scale 1, bias 0): pair() and every K split are exact on any GPU."""
+
+ words = _weights(n, k, gs, 29)[0]
+ ones = torch.ones((n, k // gs), device="cuda").bfloat16()
+ zeros = torch.zeros_like(ones)
+ q = qmm.pack(words, ones, zeros, gs)
+ want = _dequant(words, ones, zeros, gs).T.contiguous() # (k, n): every q, 0-15
+ eye = torch.eye(k, device="cuda").bfloat16()
+ assert qmm.split_k(n, k, gs) > 1 and set(want.unique().tolist()) == set(range(16))
+ for f32 in (False, True):
+ assert torch.equal(qmm.matmul(eye, q, f32=f32).float(), want), f32
+ assert torch.equal(qmm.prefill_matmul(eye, q, f32=f32).float(), want), f32
diff --git a/tests/cuda/test_qwen27_kv_room.py b/tests/cuda/test_qwen27_kv_room.py
new file mode 100644
index 000000000..820aafadc
--- /dev/null
+++ b/tests/cuda/test_qwen27_kv_room.py
@@ -0,0 +1,32 @@
+"""A serial run evicts a kept conversation only when its grow needs the room; either way the next prompt is exact."""
+
+import pytest
+import torch
+
+if not torch.cuda.is_available():
+ pytest.skip("CUDA only", allow_module_level=True)
+
+from tensorfold.cuda.streams import KVRoom, PrefixCache # noqa: E402
+from tensorfold.families.qwen3_5.cuda.decode import prefill # noqa: E402
+
+from test_qwen27_prefill import _model, _prompt, _same_state # noqa: E402
+
+BUFFER = 1024 * 128 * 2 * 2 # the tiny model's 1,024-row attention keys and values
+
+
+@pytest.mark.parametrize("budget,kept", [(4 * BUFFER, True), (BUFFER + BUFFER // 2, False)])
+def test_a_serial_run_evicts_the_drafted_conversation_only_for_room(budget, kept):
+ w = _model()
+ cache = PrefixCache(4)
+ room = KVRoom(cache, budget)
+ drafted, serial = _prompt(600, seed=31), _prompt(600, seed=32)
+ st, _ = prefill(w, drafted[:500], None, room=room)
+ cache.add(drafted[:500], st, None) # the drafted conversation's kept entry
+ prefill(w, serial, None, room=room) # a serial run on buffers of its own
+ assert bool(cache.entries) == kept
+ hit = cache.longest(drafted)
+ assert (hit is not None) == kept
+ resumed, first = prefill(w, drafted, None, state=hit[1] if hit else None, room=room)
+ fresh, first_fresh = prefill(w, drafted, None)
+ _same_state(fresh, resumed)
+ assert first == first_fresh
diff --git a/tests/cuda/test_qwen27_prompt_end_cache.py b/tests/cuda/test_qwen27_prompt_end_cache.py
index de4674008..819c89e54 100644
--- a/tests/cuda/test_qwen27_prompt_end_cache.py
+++ b/tests/cuda/test_qwen27_prompt_end_cache.py
@@ -255,9 +255,7 @@ def test_longer_and_repeated_prompts_resume_exactly(monkeypatch, w, sampling, n)
@pytest.mark.parametrize("cached,n,point", [(0, 1025, 1024), (1000, 1025, 1024), (1000, 1030, 1024),
(2000, 2049, 2048)])
def test_the_kept_state_holds_the_buffers_grown_after_the_point(monkeypatch, w, cached, n, point):
- """Key/value buffers grow to at least 1,024 rows and double; one stream sets no limit. A prefill resumed at 1,000
- or 2,000 tokens grows them inside the chunk that holds the point: the kept state then holds the grown buffers,
- whose first rows are the old ones' copies, and no buffer of its own."""
+ """A prefill resumed at 1,000 or 2,000 tokens grows the buffers in the kept chunk; kept state and prefix share."""
prompt = _prompt(n, seed=15)
# two separately built prefixes: states resumed from one prefix share its key/value buffers
@@ -268,7 +266,7 @@ def test_the_kept_state_holds_the_buffers_grown_after_the_point(monkeypatch, w,
assert _same_bits(logits, ref_logits) and pending == ref_pending
_assert_same_state(st, ref)
if cached:
- assert prefix_b.kv[3][0].shape[0] < n <= st.kv[3][0].shape[0] # grown inside the chunk
+ assert prefix_b.kv is st.kv and n <= st.kv[3][0].shape[0] # grown inside the chunk, for both
assert _shares_kv(kept, st)
(fresh_prefix, _), _ = _prefill(monkeypatch, w, prompt[:point])
_assert_same_state(kept, fresh_prefix)
diff --git a/tests/cuda/test_qwen27_shared_kv.py b/tests/cuda/test_qwen27_shared_kv.py
new file mode 100644
index 000000000..bb73dce0c
--- /dev/null
+++ b/tests/cuda/test_qwen27_shared_kv.py
@@ -0,0 +1,34 @@
+"""Clones share one attention list: a grow moves every clone and frees the old buffers; a resume stays exact."""
+
+import weakref
+
+import pytest
+import torch
+
+if not torch.cuda.is_available():
+ pytest.skip("CUDA only", allow_module_level=True)
+
+from tensorfold.families.qwen3_5.cuda.decode import clone_state, draft_decode, prefill # noqa: E402
+
+from test_qwen27_prefill import _model, _prompt, _same_state # noqa: E402
+
+ATT = 1 # the tiny model's attention layer
+
+
+@pytest.mark.parametrize("grow_by", ["prompt", "decode"])
+def test_a_grow_moves_every_clone_and_frees_the_old_buffers(grow_by):
+ w = _model()
+ prompt = _prompt(1300, seed=9)
+ st, first = prefill(w, prompt[:1000], None)
+ kept = clone_state(st)
+ old = weakref.ref(st.kv[ATT][0])
+ assert st.kv[ATT][0].shape[0] == 1024 and kept.kv is st.kv
+ if grow_by == "prompt":
+ prefill(w, prompt, None, state=st)
+ else:
+ draft_decode(w, st, prompt[:1000], first, 40, None, None, stop_eos=False)
+ assert kept.kv[ATT][0].shape[0] == 2048 and old() is None # one buffer for every clone
+ resumed, first_resumed = prefill(w, prompt[:1200], None, state=kept)
+ fresh, first_fresh = prefill(w, prompt[:1200], None)
+ _same_state(fresh, resumed)
+ assert first_fresh == first_resumed
diff --git a/tests/cuda_27b_headers.py b/tests/cuda_27b_headers.py
new file mode 100644
index 000000000..9f7c89e7f
--- /dev/null
+++ b/tests/cuda_27b_headers.py
@@ -0,0 +1,81 @@
+"""Header-only checkpoints shaped like Qwen3.8-27B (MLX 4-bit) and its DFlash2 drafter, for the startup estimates."""
+
+import json
+import math
+import struct
+from pathlib import Path
+
+SIZES = {"U32": 4, "BF16": 2}
+TEXT = {"hidden_size": 5120, "num_hidden_layers": 64, "full_attention_interval": 4, "num_attention_heads": 24,
+ "num_key_value_heads": 4, "head_dim": 256, "linear_num_key_heads": 16, "linear_num_value_heads": 48,
+ "linear_key_head_dim": 128, "linear_value_head_dim": 128, "linear_conv_kernel_dim": 4,
+ "intermediate_size": 17408, "vocab_size": 248320, "max_position_embeddings": 262144}
+DRAFT = {"hidden_size": 5120, "num_hidden_layers": 5, "num_attention_heads": 32, "num_key_value_heads": 8,
+ "head_dim": 128, "intermediate_size": 17408, "sliding_window": 2048, "dflash_config": {"block_size": 8}}
+
+
+def _write(folder: Path, config: dict, tensors: list[tuple[str, str, list[int]]]) -> Path:
+ folder.mkdir(parents=True, exist_ok=True)
+ (folder / "config.json").write_text(json.dumps(config))
+ entries, offset = {}, 0
+ for name, dtype, shape in tensors:
+ size = math.prod(shape) * SIZES[dtype]
+ entries[name] = {"dtype": dtype, "shape": shape, "data_offsets": [offset, offset + size]}
+ offset += size
+ raw = json.dumps(entries).encode()
+ (folder / "model.safetensors").write_bytes(struct.pack(" list[tuple[str, str, list[int]]]:
+ return [(name + ".weight", "U32", [n, k // 8]), (name + ".scales", "BF16", [n, k // 64]),
+ (name + ".biases", "BF16", [n, k // 64])]
+
+
+def target(folder: Path) -> Path:
+ """The 27B's tensors as stored (vision tower left out)."""
+
+ d, p = TEXT["hidden_size"], "language_model."
+ tensors = [*_q4(p + "lm_head", 248320, d), *_q4(p + "model.embed_tokens", 248320, d),
+ (p + "model.norm.weight", "BF16", [d])]
+ for i in range(TEXT["num_hidden_layers"]):
+ at = f"{p}model.layers.{i}."
+ tensors += [(at + "input_layernorm.weight", "BF16", [d]), (at + "post_attention_layernorm.weight", "BF16", [d]),
+ *_q4(at + "mlp.gate_proj", 17408, d), *_q4(at + "mlp.up_proj", 17408, d),
+ *_q4(at + "mlp.down_proj", d, 17408)]
+ if i % 4 == 3:
+ tensors += [*_q4(at + "self_attn.q_proj", 12288, d), *_q4(at + "self_attn.k_proj", 1024, d),
+ *_q4(at + "self_attn.v_proj", 1024, d), *_q4(at + "self_attn.o_proj", d, 6144),
+ (at + "self_attn.q_norm.weight", "BF16", [256]),
+ (at + "self_attn.k_norm.weight", "BF16", [256])]
+ else:
+ la = at + "linear_attn."
+ tensors += [*_q4(la + "in_proj_qkv", 10240, d), *_q4(la + "in_proj_z", 6144, d),
+ *_q4(la + "in_proj_a", 48, d), *_q4(la + "in_proj_b", 48, d), *_q4(la + "out_proj", d, 6144),
+ (la + "conv1d.weight", "BF16", [10240, 4, 1]), (la + "A_log", "BF16", [48]),
+ (la + "dt_bias", "BF16", [48]), (la + "norm.weight", "BF16", [128])]
+ config = {"model_type": "qwen3_5", "quantization": {"group_size": 64, "bits": 4}, "text_config": TEXT}
+ return _write(folder, config, tensors)
+
+
+def drafter(folder: Path) -> Path:
+ """DFlash2's bf16 tensors: the candidate selector, the fc over five taps, five layers."""
+
+ d, n = DRAFT["hidden_size"], DRAFT["intermediate_size"]
+ tensors = [("candidate_selector.hidden_projection.weight", "BF16", [256, d]),
+ ("candidate_selector.predecessor_codebook", "BF16", [248320, 256]),
+ ("candidate_selector.successor_codebook", "BF16", [248320, 256]),
+ ("fc.weight", "BF16", [d, 5 * d]), ("hidden_norm.weight", "BF16", [d]), ("norm.weight", "BF16", [d])]
+ for i in range(DRAFT["num_hidden_layers"]):
+ at, sa = f"layers.{i}.", f"layers.{i}.self_attn."
+ tensors += [(at + "attention_conv.base_kernel", "BF16", [2, 2, d]),
+ (at + "attention_conv.kernel_projection.weight", "BF16", [1280, d]),
+ (at + "mlp_conv.base_kernel", "BF16", [2, 2, d]),
+ (at + "mlp_conv.kernel_projection.weight", "BF16", [1280, d]),
+ (at + "input_layernorm.weight", "BF16", [d]), (at + "post_attention_layernorm.weight", "BF16", [d]),
+ (at + "mlp.gate_proj.weight", "BF16", [n, d]), (at + "mlp.up_proj.weight", "BF16", [n, d]),
+ (at + "mlp.down_proj.weight", "BF16", [d, n]), (sa + "q_norm.weight", "BF16", [128]),
+ (sa + "k_norm.weight", "BF16", [128]), (sa + "q_proj.weight", "BF16", [4096, d]),
+ (sa + "k_proj.weight", "BF16", [1024, d]), (sa + "v_proj.weight", "BF16", [1024, d]),
+ (sa + "o_proj.weight", "BF16", [d, 4096])]
+ return _write(folder, DRAFT, tensors)
diff --git a/tests/test_cuda_build.py b/tests/test_cuda_build.py
index 7cd4cebf7..7d8fb21f8 100644
--- a/tests/test_cuda_build.py
+++ b/tests/test_cuda_build.py
@@ -39,10 +39,19 @@ def test_the_flags_name_only_this_gpu(monkeypatch):
@pytest.mark.torch
def test_an_older_gpu_is_refused_by_name(monkeypatch):
_gpu(monkeypatch, (8, 6), "NVIDIA GeForce RTX 3090")
- with pytest.raises(RuntimeError, match=r"capability 9\.0 or newer.*RTX 3090.*is 8\.6"):
+ with pytest.raises(RuntimeError, match=r"capability 8\.9 or newer \(FP8 MMA\).*RTX 3090.*is 8\.6"):
build.arch_flags()
+@pytest.mark.torch
+def test_ada_builds_all_but_the_cluster_only_extensions(monkeypatch):
+ _gpu(monkeypatch, (8, 9), "NVIDIA GeForce RTX 4090")
+ assert build.arch_flags() == ["-gencode=arch=compute_89,code=sm_89"]
+ with pytest.raises(RuntimeError, match=r"capability 9\.0 or newer \(thread-block clusters for these weights\)"
+ r".*RTX 4090.*is 8\.9"):
+ build.arch_flags(build.CLUSTERS)
+
+
@pytest.mark.torch
def test_the_container_list_adds_nothing(monkeypatch, tmp_path):
import torch.utils.cpp_extension as ext
diff --git a/tests/test_cuda_capacity.py b/tests/test_cuda_capacity.py
index e059aca1a..2b647d6f4 100644
--- a/tests/test_cuda_capacity.py
+++ b/tests/test_cuda_capacity.py
@@ -126,6 +126,7 @@ def call(*a, **kw):
monkeypatch.setattr(torch, "empty", cpu(original_empty))
monkeypatch.setattr(torch.cuda, "set_device", lambda *a: None)
monkeypatch.setattr(capacity, "available_bytes", lambda t: 16 * capacity.GIB)
+ monkeypatch.setattr(capacity, "total_bytes", lambda t: 128 * capacity.GIB) # a GB10: 4096-row prompt chunks
calls = []
def load(*a, **kw):
calls.append(kw)
@@ -290,7 +291,8 @@ def test_actual_distributed_startup_agrees_on_smaller_rank_before_loading(tmp_pa
checkpoint(tmp_path, small_config(), HEAD)
calls, capacity = fake_runtime
geom = (mla_geometry(small_config(), 2, 8, latent=LATENT) if family == "mla" else
- gdn_geometry(small_config(), 2, 1 if family == "indexed" else 12, indexed=family == "indexed"))
+ gdn_geometry(small_config(), 2, 1, indexed=True) if family == "indexed" else
+ gdn_geometry(small_config(), 2, 12, rows=12, prompt=4096)) # the 27B engine's, prompt chunks on a GB10
transform = split_weights(rule) if family == "mla" else indexed_weights(2, False) if family == "indexed" else linear_weights
weights = capacity.estimate_weights(tmp_path, transform)
if family == "linear":
diff --git a/tests/test_cuda_gpu_floor.py b/tests/test_cuda_gpu_floor.py
new file mode 100644
index 000000000..82a482b90
--- /dev/null
+++ b/tests/test_cuda_gpu_floor.py
@@ -0,0 +1,54 @@
+"""The CUDA server refuses a GPU below the checkpoint's kernels' compute capability at startup, before any weights."""
+
+import json
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from tensorfold.cuda import build, capacity # noqa: E402
+
+GPUS = [((8, 6), "NVIDIA GeForce RTX 3090"), ((8, 9), "NVIDIA GeForce RTX 4090"), ((9, 0), "NVIDIA H100"),
+ ((12, 1), "NVIDIA GB10")]
+
+
+def _gpu(monkeypatch, capability, name):
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a: capability)
+ monkeypatch.setattr(torch.cuda, "get_device_name", lambda *a: name)
+
+
+@pytest.mark.parametrize("capability,name", GPUS)
+@pytest.mark.parametrize("need", [build.MIN_CAPABILITY, build.CLUSTERS])
+def test_the_startup_check_names_the_gpu(monkeypatch, capability, name, need):
+ _gpu(monkeypatch, capability, name)
+ if capability < need:
+ want = f"compute capability {need[0]}.{need[1]} or newer.*{name}.*{capability[0]}.{capability[1]}"
+ with pytest.raises(ValueError, match=want):
+ build.refuse_old_gpu(need)
+ else:
+ build.refuse_old_gpu(need)
+
+
+def test_no_gpu_leaves_it_to_the_engine(monkeypatch):
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+ build.refuse_old_gpu()
+
+
+@pytest.mark.parametrize("quantization,need", [({"group_size": 64, "bits": 4}, build.MIN_CAPABILITY),
+ ({"quant_method": "modelopt", "quant_algo": "NVFP4"}, build.CLUSTERS),
+ ({"quant_method": "compressed-tensors"}, build.CLUSTERS)])
+def test_a_checkpoint_sets_the_floor(tmp_path, quantization, need):
+ key = "quantization" if "bits" in quantization else "quantization_config"
+ (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen3_5", key: quantization}))
+ assert capacity.floor(tmp_path) == need
+
+
+def test_an_nvfp4_checkpoint_on_ada_is_refused_before_any_weight_loads(monkeypatch, tmp_path):
+ _gpu(monkeypatch, (8, 9), "NVIDIA GeForce RTX 4090")
+ config = {"model_type": "qwen3_5", "quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4"}}
+ (tmp_path / "config.json").write_text(json.dumps(config))
+ fail = lambda *a, **k: pytest.fail("admission read the checkpoint on a GPU it refuses") # noqa: E731
+ monkeypatch.setattr(capacity, "estimate_weights", fail)
+ with pytest.raises(ValueError, match=r"compute capability 9\.0 or newer.*RTX 4090.*is 8\.9"):
+ capacity.admit(tmp_path, None, None, torch, fail, fail)
diff --git a/tests/test_cuda_kv_room.py b/tests/test_cuda_kv_room.py
new file mode 100644
index 000000000..f29a4d372
--- /dev/null
+++ b/tests/test_cuda_kv_room.py
@@ -0,0 +1,98 @@
+"""One GPU's attention-cache budget: other conversations' buffers go oldest first, and only for room."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tensorfold.cuda import capacity, geometry
+from tensorfold.cuda.streams import KVRoom, PrefixCache
+from tensorfold.families.qwen3_5.cuda.affine_memory import draft_weights, weight_transform
+
+from cuda_27b_headers import DRAFT, TEXT, drafter, target
+
+ROW = 16 * 2 * 4 * 256 * 2 # the 27B's key and value bytes a row, every attention layer
+
+
+class Buffer:
+ """A stand-in tensor with storage of its own."""
+
+ made = 0
+
+ def __init__(self, nbytes: int) -> None:
+ Buffer.made += 1
+ self.ptr, self.bytes = Buffer.made, nbytes
+
+ def data_ptr(self) -> int:
+ return self.ptr
+
+ def untyped_storage(self):
+ return SimpleNamespace(nbytes=lambda: self.bytes)
+
+
+def _state(rows: int) -> SimpleNamespace:
+ return SimpleNamespace(kv=[pair for _ in range(16) for pair in (None, (Buffer(rows * ROW // 32),
+ Buffer(rows * ROW // 32)))])
+
+
+def _cache(*states) -> PrefixCache:
+ cache = PrefixCache(4)
+ for i, st in enumerate(states):
+ cache.add([i], st, None)
+ return cache
+
+
+def test_nothing_goes_while_it_fits():
+ live, a, b = _state(4096), _state(8192), _state(1024)
+ cache = _cache(a, b)
+ KVRoom(cache, (4096 + 8192 + 1024 + 8192) * ROW)(live, 8192 * ROW)
+ assert [e[1] for e in cache.entries] == [a, b]
+
+
+def test_other_buffers_go_least_recently_used_first_and_the_resumed_list_stays():
+ live, a, b, c = _state(4096), _state(8192), _state(8192), _state(1024)
+ shared = SimpleNamespace(kv=live.kv) # the entry the live state resumed: same list
+ cache = _cache(a, b, shared, c)
+ cache.longest([2, 9]) # resuming touches it: newest, and hit
+ room = KVRoom(cache, (4096 + 8192 + 1024 + 8192) * ROW)
+ room(live, 8192 * ROW) # 8,192 rows past the budget: a, the oldest, goes
+ assert [e[1] for e in cache.entries] == [b, c, shared]
+ room(live, 17000 * ROW) # then b, then c; never the shared entry
+ assert [e[1] for e in cache.entries] == [shared]
+ room(live, 10 ** 6 * ROW) # only this conversation is left: nothing more to free
+ assert [e[1] for e in cache.entries] == [shared]
+
+
+def test_a_gb10_budget_evicts_nothing_today_keeps(tmp_path):
+ """A GB10's 100 GiB keeps three native-window conversations beside the live one (the old estimate, two)."""
+
+ folder, draft = target(tmp_path / "target"), drafter(tmp_path / "drafter")
+ weights = capacity.estimate_weights(folder, weight_transform(folder, one_gpu=True))
+ side = draft_weights(draft)
+ main = geometry.gdn_geometry(TEXT, 1, 128, rows=128, prompt=4096, evicts=True)
+ rider = geometry.draft_geometry(DRAFT, 1, 128, bounded=True)
+ both = capacity.Geometry(lambda slots: main.bytes_at(slots) + rider.bytes_at(slots), 128)
+ budget = 100 * capacity.GIB
+ plan = capacity.make_plan(262144, None, False, budget, capacity.Weights(weights.resident + side.resident, 0),
+ both)
+ assert plan.fitting == 262144
+ spare = budget - weights.resident - side.resident - both.needed(plan.fitting)
+ room = KVRoom(_cache(*(_state(262144) for _ in range(3))), spare + geometry.live_kv(TEXT, 1, 262144))
+ before = list(room.cache.entries)
+ room(_state(262144 - 1024), 1024 * ROW // 32) # the live one's last layer grows to the window
+ assert room.cache.entries == before
+
+
+@pytest.mark.parametrize("gib,window", [(19.16, 39936), (100, 262144)])
+def test_one_live_window_is_admitted(gib, window, tmp_path):
+ """An RTX 4090's default budget admits a 40k window with DFlash2 (four copies gave 8,192); a GB10 its native one."""
+
+ folder, draft = target(tmp_path / "target"), drafter(tmp_path / "drafter")
+ weights = capacity.estimate_weights(folder, weight_transform(folder, one_gpu=True))
+ side = draft_weights(draft)
+ rows = 128 if gib == 100 else 12
+ main = geometry.gdn_geometry(TEXT, 1, rows, rows=rows, prompt=4096 if gib == 100 else 2048, evicts=True)
+ rider = geometry.draft_geometry(DRAFT, 1, rows, bounded=True)
+ both = capacity.Geometry(lambda slots: main.bytes_at(slots) + rider.bytes_at(slots), rows)
+ plan = capacity.make_plan(262144, None, False, int(gib * capacity.GIB),
+ capacity.Weights(weights.resident + side.resident, 0), both)
+ assert plan.fitting // 1024 * 1024 == window
diff --git a/tests/test_cuda_memory_plan.py b/tests/test_cuda_memory_plan.py
new file mode 100644
index 000000000..7fb1a2864
--- /dev/null
+++ b/tests/test_cuda_memory_plan.py
@@ -0,0 +1,73 @@
+"""The 27B's startup memory plan on any GPU: exact bytes for the weights and their load, sized from the budget."""
+
+import pytest
+
+from tensorfold.cuda import capacity, geometry
+from tensorfold.families.qwen3_5.cuda.affine_memory import draft_bytes, draft_weights, weight_transform
+
+from cuda_27b_headers import TEXT, drafter, target
+
+GIB = capacity.GIB
+HEAD = 248320 * 640 * 4 + 2 * 248320 * 80 * 2 # the 4-bit head: words, scales and biases
+
+
+@pytest.fixture(scope="module")
+def model(tmp_path_factory):
+ root = tmp_path_factory.mktemp("27b")
+ return target(root / "target"), drafter(root / "drafter")
+
+
+def test_one_gpu_counts_the_tiled_head_once(model):
+ folder, _ = model
+ twice = capacity.estimate_weights(folder, weight_transform(folder))
+ once = capacity.estimate_weights(folder, weight_transform(folder, one_gpu=True))
+ assert twice.resident - once.resident == HEAD # the drafter reads the head's rows as views
+ assert round(twice.resident / GIB, 2) == 14.76 # the real checkpoint's figure, measured 14.20 held
+ assert once.staging == 3 * 248320 * 640 * 4 # the largest tensor as loaded: the head's words
+
+
+def test_the_drafter_load_peak_is_its_largest_quantize(model):
+ _, folder = model
+ held = capacity.estimate_weights(folder, draft_bytes)
+ peak = draft_weights(folder)
+ assert peak.resident == held.resident and round(held.resident / GIB, 2) == 0.97
+ assert peak.staging == 14 * 5120 * 25600 # fc: bf16 upload, fp32 copy, two fp32 temporaries
+
+
+def test_the_drafter_loads_after_the_target(model, monkeypatch):
+ folder, draft = model
+ monkeypatch.setattr(capacity, "available_bytes", lambda torch: 100 * GIB)
+ monkeypatch.setattr(capacity, "page_room", lambda torch: None)
+ geometry = capacity.Geometry(lambda slots: slots * 1024, 12)
+ main = capacity.estimate_weights(folder, weight_transform(folder, one_gpu=True))
+ side = draft_weights(draft)
+ receipt = capacity.admit(folder, 8192, True, None, geometry, weight_transform(folder, one_gpu=True),
+ draft_dir=draft, draft_weights=draft_weights)
+ assert receipt["weight_bytes_estimate"] == main.resident + side.resident
+ assert receipt["loading_bytes_estimate"] == max(main.staging - side.resident, side.staging)
+ # measured on an RTX 4090: 16.73 GiB allocated at the load's peak (17.95 reserved: the budget's reserve covers it)
+ assert round((receipt["weight_bytes_estimate"] + receipt["loading_bytes_estimate"]) / GIB, 2) == 16.78
+
+
+@pytest.mark.parametrize("gib,rows", [(128, 4096), (80, 4096), (23.54, 2048), (16, 1024), (12, 1024), (8, 512)])
+def test_a_prompt_chunk_is_sized_to_the_card(gib, rows):
+ assert geometry.prompt_rows(int(gib * GIB), geometry.prompt_row_bytes(TEXT)) == rows
+
+
+@pytest.mark.parametrize("slots", [0, 4108, 65536 + 128, 262144 + 128])
+def test_a_gb10_keeps_todays_scratch(slots):
+ """One stream on a GB10 verifies 128 rows and prompts 4096: the scratch and every cache size are today's."""
+
+ gb10 = geometry.gdn_geometry(TEXT, 1, 128, rows=128, prompt=geometry.prompt_rows(128 * GIB, 360448))
+ assert gb10.bytes_at(slots) == geometry.gdn_geometry(TEXT, 1, 128).bytes_at(slots)
+ assert gb10.bytes_at(262144 + 128) == 74040369152
+
+
+def test_twelve_rows_bound_the_scratch_and_a_prompt_chunk_shares_it():
+ extent = 5120 + 248320 + 2 * (17408 + 5120) + 16480 + 24 * 256
+ verify = geometry.gdn_geometry(TEXT, 1, 12, rows=12)
+ prompt = geometry.gdn_geometry(TEXT, 1, 12, rows=12, prompt=2048)
+ wide = geometry.gdn_geometry(TEXT, 1, 12, rows=128)
+ replay = 48 * (128 - 12) * (16480 * 2 + 16 * 128 * 4 + 48 * 128 * 4 + 48 * 8) + 32 * (128 - 12) * 2560 * 4
+ assert wide.bytes_at(4108) - verify.bytes_at(4108) >= replay + 16 * (128 - 12) * extent * 4
+ assert prompt.bytes_at(4108) - verify.bytes_at(4108) == 2048 * 360448 - 16 * 12 * extent * 4
diff --git a/tests/test_qwen27_prompt_end_cache_host.py b/tests/test_qwen27_prompt_end_cache_host.py
index 1147871b0..8e6925f16 100644
--- a/tests/test_qwen27_prompt_end_cache_host.py
+++ b/tests/test_qwen27_prompt_end_cache_host.py
@@ -87,7 +87,7 @@ def __init__(self):
self.prefills = []
def prefill(self, w, prompt, sampling, drafter=None, *, state=None, keep_at=None, rank=None, limit=0, stops=(),
- keep=None, vision=None):
+ keep=None, vision=None, room=None):
start = state.pos if state is not None else 0
if state is not None:
assert list(prompt[:start]) == state.ids, "resumed from a state that is not a prefix of the prompt"
@@ -831,7 +831,7 @@ def run(keep, until=n):
assert [(p0, rows) for p0, rows, *_ in got] == [(a, b - a) for a, b in spans] == [(p0, r) for p0, r, *_ in ref_calls]
assert [cut for _, _, cut, _, _ in got] == [keep_at - a if a < keep_at < b else 0 for a, b in spans]
assert normed == ref == ("normed", n) and st.pos == n
- assert kept.pos == keep_at and kept.kv == st.kv and kept.kv is not st.kv
+ assert kept.pos == keep_at and kept.kv is st.kv # one list: a later grow frees the old buffers
assert kept.rec[0] == ("rec at", keep_at)
if draft is None:
assert snap is None
@@ -1114,9 +1114,7 @@ def test_the_same_prompt_again_resumes_from_its_kept_state_with_one_token(cpu):
@pytest.mark.parametrize("cached,length,point", [(0, 1025, 1024), (1000, 1025, 1024), (1000, 1030, 1024)])
def test_the_kept_state_holds_the_buffers_grown_after_the_point(cpu, cached, length, point):
- """Key/value buffers grow to at least 1,024 rows and double. A prefill resumed at 1,000 tokens grows them past
- 1,024 inside the chunk that holds the point: the kept state then holds the grown buffers, whose first rows are
- the old ones' copies, and no buffer of its own."""
+ """A prefill resumed at 1,000 tokens grows the buffers in the kept chunk; the kept state and prefix share them."""
torch = cpu.torch
prompt = _prompt(length, 50)
@@ -1127,7 +1125,7 @@ def test_the_kept_state_holds_the_buffers_grown_after_the_point(cpu, cached, len
assert _bits_equal(torch, logits, ref_logits) and pending == ref_pending
_assert_same_state(torch, st, ref)
if cached:
- assert prefix_b.kv[3][0].shape[0] == 1024 < st.kv[3][0].shape[0] # grown inside the chunk
+ assert prefix_b.kv is st.kv and st.kv[3][0].shape[0] > 1024 # grown inside the chunk, for both
assert _shares_kv(kept, st)
(fresh, _), _ = cpu.run(prompt[:point])
_assert_same_state(torch, kept, fresh)
diff --git a/tests/test_vision_cuda.py b/tests/test_vision_cuda.py
index f70ca7de8..f95adb033 100644
--- a/tests/test_vision_cuda.py
+++ b/tests/test_vision_cuda.py
@@ -170,11 +170,12 @@ def test_state_clone_preserves_image_offset_without_importing_cuda():
namespace = {"State": state_type}
exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), namespace)
state = state_type()
- state.pos, state.limit, state.rope_delta = 257, 1024, -192
+ state.pos, state.limit, state.rope_delta, state.room = 257, 1024, -192, object()
state.conv, state.rec, state.kv = [1], [2], [3]
cloned = namespace["clone_state"](state)
- assert (cloned.pos, cloned.rope_delta, cloned.limit) == (257, -192, 1024)
- assert cloned.kv == state.kv and cloned.kv is not state.kv
+ assert (cloned.pos, cloned.rope_delta, cloned.limit, cloned.room) == (257, -192, 1024, state.room)
+ assert cloned.conv == state.conv and cloned.conv is not state.conv
+ assert cloned.kv is state.kv # one attention list: a grow reaches every clone
def test_a_meta_built_tower_matches_a_normally_built_one_in_the_installed_transformers():
From 50588fbbe4c73bf18bc88cc1829f702f9a824d88 Mon Sep 17 00:00:00 2001
From: Kevin Gilbertson <880013+gilby@users.noreply.github.com>
Date: Wed, 30 Sep 2026 19:52:33 +0000
Subject: [PATCH 086/177] feat(lane_fuse): stack 4-bit groups of 32 too
(oQ-style g32 checkpoints keep their fused projections)
lane_fuse refused any member not in groups of 64, so a 4-bit g32 checkpoint (MLX 4-bit g32, oMLX oQ4 bases)
ran every zba/kv/gu projection as separate lane calls even though lane_matmul, tile_weight and pack_scales all
take groups of 32. A stack now takes the members' shared group size (mixed group sizes stay separate, like mixed
widths), tiles its tail in that group and passes it to the lane matmul. Same kernels, same K split: each member's
bits are unchanged. Tests: the build and whole-round fused==unfused cases gain 4-bit g32.
---
.../kernels/qwen/dense/v1/lane_fuse.py | 23 +++++++-----
tests/test_lane_fuse.py | 36 +++++++++++--------
2 files changed, 35 insertions(+), 24 deletions(-)
diff --git a/src/tensorfold/kernels/qwen/dense/v1/lane_fuse.py b/src/tensorfold/kernels/qwen/dense/v1/lane_fuse.py
index 9bea544b5..f9b8859d0 100644
--- a/src/tensorfold/kernels/qwen/dense/v1/lane_fuse.py
+++ b/src/tensorfold/kernels/qwen/dense/v1/lane_fuse.py
@@ -23,13 +23,15 @@
class _Group:
- """A stacked projection: ``weight`` (sum N, K*bits/32) and ``sbt`` (K/64, sum N, 2) for one lane matmul."""
+ """A stacked projection: ``weight`` (sum N, K*bits/32) and ``sbt`` (K/group, sum N, 2) for one lane matmul."""
- __slots__ = ("weight", "sbt", "tiled", "sk", "k", "sizes", "added", "members", "held", "sbts", "nt", "rotate")
+ __slots__ = ("weight", "sbt", "tiled", "sk", "k", "sizes", "added", "members", "held", "sbts", "nt", "rotate",
+ "group")
def __init__(self, weight: Any, sbt: Any, tiled: bool, sk: int, k: int, sizes: tuple[int, ...], added: int,
- members: tuple[Any, ...], nt: int = 32, rotate: Any = None) -> None:
+ members: tuple[Any, ...], nt: int = 32, rotate: Any = None, group: int = 64) -> None:
self.weight, self.sbt, self.tiled, self.sk, self.k, self.sizes = weight, sbt, tiled, sk, k, sizes
+ self.group = group # the members' shared group size
self.rotate = rotate # the members' shared input transform
self.nt = nt # the stack's tile width (lane_qmm)
self.added = added # bytes not shared with the modules
@@ -82,10 +84,11 @@ def _build(parent: Any, kind: str) -> _Group | _Unfusable:
rotate = outer[0].rotate if hasattr(outer[0], "rotate") else None
for m in members:
w = m["weight"]
- if not lane_qmm.takes(m) or m.group_size != 64 or "bias" in m or w.dtype != mx.uint32 or w.ndim != 2:
+ if not lane_qmm.takes(m) or m.group_size not in (32, 64) or "bias" in m or w.dtype != mx.uint32 or w.ndim != 2:
return no
- bits = members[0].bits
- if any(m.bits != bits for m in members): # one kernel a stack: members of mixed widths stay separate calls
+ bits, group = members[0].bits, int(members[0].group_size)
+ # one kernel a stack: members of mixed widths or group sizes stay separate calls
+ if any(m.bits != bits or int(m.group_size) != group for m in members):
return no
kw = int(members[0]["weight"].shape[1])
k = kw * 32 // bits
@@ -111,7 +114,8 @@ def _build(parent: Any, kind: str) -> _Group | _Unfusable:
if any(tiled[j:]) or tail % lane_qmm.NT or sum(m["weight"].nbytes for m in members[j:]) > _SMALL_TAIL:
return no
parts = [m["weight"] for m in members[:j]]
- parts.append(lane_qmm.tile_weight(mx.concatenate([m["weight"] for m in members[j:]], axis=0), bits=bits))
+ parts.append(lane_qmm.tile_weight(mx.concatenate([m["weight"] for m in members[j:]], axis=0), group=group,
+ bits=bits))
viewed = members[:j]
stacked_tiled = True
copied = parts[-1].nbytes
@@ -136,7 +140,7 @@ def _build(parent: Any, kind: str) -> _Group | _Unfusable:
offset += n
mx.eval(views)
return _Group(weight, sbt, stacked_tiled, sk, k, sizes, sbt.nbytes + copied, members,
- nt if stacked_tiled else lane_qmm.NT, rotate)
+ nt if stacked_tiled else lane_qmm.NT, rotate, group)
def _group(parent: Any, kind: str, *, build: bool | None = None) -> _Group | None:
@@ -176,7 +180,8 @@ def _project(parent: Any, kind: str, x: mx.array) -> mx.array | None:
return None
if group.rotate is not None:
x = group.rotate(x)
- return lane_qmm.lane_matmul(x, group.weight, group.sbt, tiled=group.tiled, sk=group.sk, nt=group.nt)
+ return lane_qmm.lane_matmul(x, group.weight, group.sbt, tiled=group.tiled, sk=group.sk, nt=group.nt,
+ group=group.group)
def gdn_in(gdn: Any, x: mx.array) -> mx.array | None:
diff --git a/tests/test_lane_fuse.py b/tests/test_lane_fuse.py
index bf854a53a..5f4d0cc13 100644
--- a/tests/test_lane_fuse.py
+++ b/tests/test_lane_fuse.py
@@ -158,7 +158,7 @@ def __init__(self, names_sizes):
setattr(self, name, nn.Linear(K, n, bias=False))
-def _real_groups(bits=4):
+def _real_groups(bits=4, group=64):
root = nn.Module()
root.linear_attn = _Holder([("in_proj_z", 6144), ("in_proj_b", 48), ("in_proj_a", 48)])
root.self_attn = _Holder([("k_proj", 1024), ("v_proj", 1024)])
@@ -167,15 +167,19 @@ def _real_groups(bits=4):
if isinstance(module, nn.Linear):
mx.random.seed(300 + i)
module.weight = (mx.random.normal(module.weight.shape) * 0.02).astype(mx.bfloat16)
- nn.quantize(root, group_size=64, bits=bits)
+ nn.quantize(root, group_size=group, bits=bits)
mx.eval(root.parameters())
return root
-@pytest.mark.parametrize("bits", [4, 3, 2, 5, 6, 8])
-def test_build_keeps_the_weights_and_adds_only_the_scales(bits):
+# every width in groups of 64; 4-bit in groups of 32 too (the lane matmul reads 4-bit g32 checkpoints)
+WIDTHS = [(4, 64), (3, 64), (2, 64), (5, 64), (6, 64), (8, 64), (4, 32)]
+
+
+@pytest.mark.parametrize("bits,group", WIDTHS)
+def test_build_keeps_the_weights_and_adds_only_the_scales(bits, group):
_needs_tensor_units()
- root = _real_groups(bits)
+ root = _real_groups(bits, group)
members = {kind: [getattr(parent, n) for n in lane_fuse.GROUPS[kind]]
for kind, parent in (("zba", root.linear_attn), ("kv", root.self_attn), ("gu", root.mlp))}
originals = {id(m): m["weight"] for ms in members.values() for m in ms} # MLX's layout
@@ -198,12 +202,14 @@ def test_build_keeps_the_weights_and_adds_only_the_scales(bits):
expect = sum(sum(lane_qmm.pack_scales(m["scales"], m["biases"]).nbytes for m in ms) for ms in members.values())
expect += 96 * (K * bits // 32) * 4
assert added == expect
- # scales are 1/(2 * bits) of the weights' bytes (1/8 at 4 bits): a second copy of the weights would add them all
- assert abs(grown - expect) < 1024**2 and grown < 1.6 / (2 * bits) * weight_bytes, (grown, expect, weight_bytes)
+ # scales are 64 / (2 * bits * group) of the weights' bytes (1/8 at 4 bits g64): a second copy would add them all
+ assert abs(grown - expect) < 1024**2 and grown < 1.6 * 64 / (2 * bits * group) * weight_bytes, \
+ (grown, expect, weight_bytes)
for ms in members.values():
for m in ms:
w = m["weight"]
- seen = lane_qmm.untile_weight(w, bits=bits) if getattr(m, "_lane_tiled", False) else w
+ seen = lane_qmm.untile_weight(w, int(getattr(m, "_lane_nt", lane_qmm.NT)), group, bits=bits) \
+ if getattr(m, "_lane_tiled", False) else w
assert _same(seen, originals[id(m)]), "a member's weight changed"
# each member's own call (on its view of the stack) keeps its bits, and the stack gives them too
lane_fuse.enabled = True
@@ -305,7 +311,7 @@ def test_groups_of_mixed_widths_stay_separate():
lane_fuse.clear(root)
-def _tiny_model(bits=4):
+def _tiny_model(bits=4, group=64):
from mlx_lm.models.qwen3_5 import TextModel, TextModelArgs
# Qwen3.8's head shapes and MLP width on a 1024-wide residual: the gate/up stack alone would
@@ -318,21 +324,21 @@ def _tiny_model(bits=4):
mx.random.seed(21)
model = TextModel(args)
model.set_dtype(mx.bfloat16)
- nn.quantize(model, group_size=64, bits=bits)
+ nn.quantize(model, group_size=group, bits=bits)
mx.eval(model.parameters())
return model
-@pytest.mark.parametrize("bits", [4, 3, 2, 5, 6, 8])
-def test_tree_forward_fused_equals_unfused(bits):
+@pytest.mark.parametrize("bits,group", WIDTHS)
+def test_tree_forward_fused_equals_unfused(bits, group):
"""Whole lane-decoder rounds: prompt chain, draft tree, commit, chain, one row; every bit the same."""
_needs_tensor_units()
- _check_fused_rounds(bits, pipeline_layers=2)
+ _check_fused_rounds(bits, pipeline_layers=2, group=group)
-def _check_fused_rounds(bits, *, pipeline_layers):
- model = _tiny_model(bits)
+def _check_fused_rounds(bits, *, pipeline_layers, group=64):
+ model = _tiny_model(bits, group)
core, head = model.model, model.lm_head
mx.random.seed(5)
prompt = [int(t) for t in mx.random.randint(0, 512, (40,)).tolist()] # gate/up stacked (> 32 rows)
From e3ac0eaad8d359c70b874fe5720d52d1ec380ef9 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Wed, 30 Sep 2026 22:18:42 +0100
Subject: [PATCH 087/177] chore: Apache-2.0 from 0.6.0
- LICENSE holds the Apache License 2.0, and NOTICE credits TensorFold contributors
- Releases up to 0.5.0 stay MIT, and code written before 0.6.0 keeps its MIT notice in LICENSES/MIT.txt
- pyproject's license, classifier and license files, and the README's license section, say so
---
LICENSE | 223 ++++++++++++++++++++++++++++++++++++++++++-----
LICENSES/MIT.txt | 21 +++++
NOTICE | 8 ++
README.md | 3 +-
pyproject.toml | 6 +-
5 files changed, 236 insertions(+), 25 deletions(-)
create mode 100644 LICENSES/MIT.txt
create mode 100644 NOTICE
diff --git a/LICENSE b/LICENSE
index a8e21673d..d64569567 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,21 +1,202 @@
-MIT License
-
-Copyright (c) 2026 TensorFold contributors
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/LICENSES/MIT.txt b/LICENSES/MIT.txt
new file mode 100644
index 000000000..a8e21673d
--- /dev/null
+++ b/LICENSES/MIT.txt
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 TensorFold contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/NOTICE b/NOTICE
new file mode 100644
index 000000000..50e45b9ac
--- /dev/null
+++ b/NOTICE
@@ -0,0 +1,8 @@
+TensorFold
+Copyright 2026 TensorFold contributors
+https://github.com/ashhart/TensorFold
+
+TensorFold is licensed under the Apache License, Version 2.0 (LICENSE) from 0.6.0.
+Releases up to 0.5.0 were MIT-licensed, and code written before 0.6.0 keeps its
+MIT notice in LICENSES/MIT.txt. Third-party code and its licenses are listed in
+THIRD_PARTY_NOTICES.md.
diff --git a/README.md b/README.md
index ac966f285..f71cf47de 100644
--- a/README.md
+++ b/README.md
@@ -268,5 +268,6 @@ every release. The first time a new version serves, it prints one line linking t
Family interfaces, kernel layout and verification requirements are in the [recipe book](docs/recipes/README.md),
[family map](src/tensorfold/families/README.md) and [kernel map](src/tensorfold/kernels/README.md).
-MIT; see [LICENSE](LICENSE) and [third-party notices](THIRD_PARTY_NOTICES.md).
+Apache-2.0 from 0.6.0; see [LICENSE](LICENSE), [NOTICE](NOTICE) and [third-party notices](THIRD_PARTY_NOTICES.md).
+Releases up to 0.5.0 were MIT, and code written before 0.6.0 keeps its [MIT notice](LICENSES/MIT.txt).
Model weights keep their own licenses.
diff --git a/pyproject.toml b/pyproject.toml
index 5deec7f2f..ae9dd3ba3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,13 +8,13 @@ dynamic = ["version"]
description = "Fast, exact LLM decoding on Apple Silicon (MLX) and NVIDIA GPUs (CUDA) behind an OpenAI-compatible endpoint"
readme = "README.md"
requires-python = ">=3.11"
-license = { text = "MIT" }
+license = { text = "Apache-2.0" }
authors = [{ name = "TensorFold contributors" }]
keywords = ["llm", "mlx", "apple-silicon", "cuda", "dgx-spark", "inference", "speculative-decoding", "openai"]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
- "License :: OSI Approved :: MIT License",
+ "License :: OSI Approved :: Apache Software License",
"Operating System :: MacOS",
"Operating System :: POSIX :: Linux",
"Environment :: GPU :: NVIDIA CUDA",
@@ -45,7 +45,7 @@ grammar = ["xgrammar>=0.2.8,<0.3"]
tensorfold = "tensorfold.cli:main"
[tool.setuptools]
-license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md", "LICENSES/*"]
+license-files = ["LICENSE", "NOTICE", "THIRD_PARTY_NOTICES.md", "LICENSES/*"]
[tool.setuptools.packages.find]
where = ["src"]
From c4646171139ee8a3c38103eaa1699dad226ec12b Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Wed, 30 Sep 2026 22:18:52 +0100
Subject: [PATCH 088/177] release: TensorFold 0.6.0
---
CHANGELOG.md | 50 ++++++++++++++++++++++++++++++++++++++
src/tensorfold/__init__.py | 2 +-
2 files changed, 51 insertions(+), 1 deletion(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b2078f098..9ac3a6f3e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,56 @@
`tensorfold update` prints the sections below that are newer than the version you had. Each release's page on
GitHub has the full notes and the measurements behind them.
+## 0.6.0 (30 Sep 2026)
+
+- **RTX 40 cards.** CUDA now runs on compute capability 8.9 (Ada). On one RTX 4090 the 27B serves with DFlash2 in a
+ 40,182-token window, exact: drafted replies equal serial ones, and a resume or resend gives a fresh run's reply.
+ Prompts fill at 2.4-2.6k tok/s from 2k to 32k tokens, and replies decode at 64-86 tok/s. On one GPU the 27B's kept
+ prompt states give way, oldest first, when a live reply needs the room, so admission counts one live window.
+- **Prompts fill inside the decode rounds.** With `--parallel` on CUDA, Flash Next prefills a queued prompt in the
+ same forward as the live replies instead of stopping them: on one DGX Spark, first tokens came 2.8-3.1x sooner than
+ on 0.5.0. On Macs several prompts fill side by side, the fewest tokens left first: on an M3 Ultra, short requests
+ queued behind a long prompt got their first token in a median 9.9 s instead of 113 s (the slowest 11.7 s, not 130).
+- **Conversations resume on more engines.** An identical resend or the next thinking turn now resumes from the kept
+ prompt state on Flash Next, Qwen3.6, Nemotron and GLM-5.3 on CUDA, as the 27B did, with a fresh run's reply: an
+ 18.7k-token Flash Next resend went from 8.3 s to 0.08 s. On Macs the prompt cache keeps each conversation's newest
+ checkpoint and grows into memory the model leaves idle.
+- **CUDA prompts run at bf16 by default.** Against an fp32 reference the 27B's prompt rows are 20x closer than with
+ FP8 (KL 0.0031 against 0.0624). `--prefill-fp8` keeps 0.5.0's faster FP8 prompts for those who want them.
+- **Tool calls for agents.** The CUDA server streams tool-call arguments as the model writes them (the longest
+ silence in a long call fell from about 30 s to half a second), Python-spelled values like `False` and `None`
+ decode to their schema types, Gemma 4's bare tool calls parse (#121), and a prompt past the context window gets
+ OpenAI's `context_length_exceeded`, so clients compact instead of retrying.
+- **More checkpoints.** GLM-5.3 on Macs reads 8-bit and Q8_0 GGUF checkpoints and takes images, and has an opt-in
+ float32 activation mode (bf16 stays the default). On Macs, Flash Next loads oMLX's oQ checkpoints with scaled
+ n-gram tables; on CUDA it reads consolidated EXL3 tables and block-scaled FP8 linears.
+- **Faster.** Flash Next NVFP4 decodes 4.9-7.1% faster and fills prompts 15-16% faster at 2k-16k, bit-identical. On an
+ M5 Ultra the 27B with DFlash2 on an oQ4e checkpoint went from 33 to 131-160 tok/s. Before M5, a lone stream's copy
+ windows widen to 128 rows (an M3 Ultra edit ran 190 -> 233 tok/s), and Flash Next's concurrent rounds use the matrix
+ units. GLM-5.3's DFlash2 reads only its sliding window: 8% faster at 52k, and 3.8 GiB lighter a rank on two Sparks.
+- **Operations.** Prometheus `/metrics` on both servers, `TENSORFOLD_MEMORY_RESERVE_GIB` for the CUDA startup
+ reserve, `--checkpoint-slots` for the 27B's concurrent decoder on CUDA, and an idle GLM rank no longer spins.
+- **Apache-2.0.** TensorFold is licensed under the Apache License 2.0 from this release, which adds an explicit patent
+ grant from contributors. Releases up to 0.5.0 stay MIT, and code written before 0.6.0 keeps its MIT notice in
+ `LICENSES/MIT.txt`.
+- **Community pull requests.**
+ - Resumed resends and thinking turns on Flash Next (#124). Thanks to @benthecarman, and to @Arminova for the
+ GB10 measurements and the third-resend test.
+ - Streamed tool-call arguments on CUDA (#114), `--alias` on CUDA (#111) and Python-spelled tool parameters
+ (#135). Thanks to @olexale, @philip-pentatonic and @outcastofmusic.
+ - GLM-5.3 on Macs: image input (#101), 8-bit and Q8_0 GGUF checkpoints (#119), float32 activations (#118) and
+ the backbone wiring notes (#120). Thanks to @mgoldwasser and @feni6.
+ - GLM-5.3 on two Sparks: lighter prompt buffers, the EXL3 estimate, the MTP setting, the idle rank and DFlash2's
+ ring (#128, #129, #131, #132, #134), visible-pool selection (#140), and the CUDA startup reserve setting (#133).
+ Thanks to @MiaAI-Lab and @mikolaj92.
+ - Flash Next: scaled n-gram tables (#148, reported in #142) and the M5 draft head fix (#147), consolidated EXL3
+ tables (#145), block-scaled FP8 (#126) and the NVFP4 expert speedups (#102, #105). Thanks to @gilby,
+ @cwschroeder, @shantanugoel, @jschmied and @tournierjc.
+ - On M5 Macs the 27B fuses the projections of 4-bit group-32 checkpoints such as oQ4e too (#164): four streams went
+ from 315-320 to 332-334 tok/s on an M5 Ultra, with the same bits. Thanks to @gilby.
+ - `--checkpoint-slots` (#125), and the DFlash2 drafter's admitted size (#112, from @jkuepker's ROCm work). Thanks
+ to @nood-co1 and @jkuepker.
+
## 0.5.0 (29 Sep 2026)
- **OpenAI's Responses API on both servers.** `/v1/responses` runs as a chat completion: every response has its chat
diff --git a/src/tensorfold/__init__.py b/src/tensorfold/__init__.py
index a0e41685b..e858797ba 100644
--- a/src/tensorfold/__init__.py
+++ b/src/tensorfold/__init__.py
@@ -1,3 +1,3 @@
"""TensorFold: fast, exact LLM decoding on Apple Silicon and NVIDIA GPUs behind an OpenAI-compatible endpoint."""
-__version__ = "0.5.0"
+__version__ = "0.6.0"
From d61c0acc37615d1f4cf587385f4c73cd3092c00f Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:36:56 +0100
Subject: [PATCH 089/177] fix(cuda): GLM's EXL3 path owns its scratch
---
.../families/glm5_next/cuda/exl3_generic.py | 25 ++------
.../families/glm5_next/cuda/forward.py | 24 ++++----
.../families/glm5_next/cuda/weights.py | 2 +-
tests/cuda/glm_decision_tp.py | 2 +-
tests/cuda/test_exl3_any_width_glm.py | 57 ++++++++-----------
tests/cuda/test_glm_engine.py | 7 ++-
tests/test_cuda_geometry.py | 6 +-
7 files changed, 50 insertions(+), 73 deletions(-)
diff --git a/src/tensorfold/families/glm5_next/cuda/exl3_generic.py b/src/tensorfold/families/glm5_next/cuda/exl3_generic.py
index 07d3346de..992c7f6ee 100644
--- a/src/tensorfold/families/glm5_next/cuda/exl3_generic.py
+++ b/src/tensorfold/families/glm5_next/cuda/exl3_generic.py
@@ -1,26 +1,9 @@
-"""GLM routed experts through the generic mixed-width EXL3 path.
-
-Serves any bit width (k2 2..16) with bit-identical epilogue arithmetic
-(ACT_BF16, GLM tile configs). Uses the exl3_mm.Scratch's rows/slots (passed
-explicitly) to size its own generic Scratch once per key.
-
-The shared expert (pick == E) is excluded by the generic kernels, as in
-exl3_mm; the caller combines its slot from the BF16 shared MLP.
-"""
-from __future__ import annotations
-
-import torch
+"""GLM's universal EXL3 route uses buffer-owned scratch and the existing bf16 expert epilogue."""
from tensorfold.cuda.exl3 import experts as generic
-_scratch: dict = {}
+def routed(x, pick, ex, scratch, rows, limit):
+ """Write the live routed pairs into scratch.y; the caller supplies each row's shared-expert slot."""
-def routed(x, pick, ex, rows, limit):
- """ey[pair] fp32 for routed pairs; the shared slot's row is stale (caller overwrites)."""
- key = (rows, pick.shape[1], ex.count)
- s = _scratch.get(key)
- if s is None:
- s = generic.Scratch(ex, rows, pick.shape[1], device=str(x.device))
- _scratch[key] = s
- return generic.routed(x, pick, None, ex, s, None, rows, limit, act_mode=generic.ACT_BF16)
+ return generic.routed(x, pick, None, ex, scratch, None, rows, limit, act_mode=generic.ACT_BF16)
diff --git a/src/tensorfold/families/glm5_next/cuda/forward.py b/src/tensorfold/families/glm5_next/cuda/forward.py
index f1234d77b..f073b6e4e 100644
--- a/src/tensorfold/families/glm5_next/cuda/forward.py
+++ b/src/tensorfold/families/glm5_next/cuda/forward.py
@@ -4,16 +4,18 @@
from dataclasses import dataclass
from typing import Sequence
+from types import SimpleNamespace
import torch
import triton
import triton.language as tl
from tensorfold.cuda import experts as grouped
+from tensorfold.cuda.exl3.experts import Scratch as Exl3Scratch
from tensorfold.cuda.geometry import MLA_PROMPT_ATT_ROWS as PROMPT_ATT_ROWS # a dense latent call's prompt rows
from tensorfold.cuda.kernels import prefill_attention, qmm as shared
-from . import glue, kda as kda_mod, latent, prof, qmm, sparse
+from . import exl3_generic, glue, kda as kda_mod, latent, prof, qmm, sparse
from .attention import AttnScratch, attention, kv_write
from .weights import LayerW, Weights
@@ -97,15 +99,14 @@ def __init__(self, w: Weights, rows: int, capacity: int = 2560, *, prefill: bool
self.wts = torch.empty((rows, slots), dtype=f32, device=dev)
self.eact = torch.empty((rows * slots, ml), dtype=bf, device=dev)
exl3 = c.quant == "exl3"
- self.ey = torch.empty((rows, slots, D), dtype=bf if prefill and not exl3 else f32, device=dev)
- self.plan = grouped.Plan(rows, slots, c.experts + 1, dev, prefill=prefill and not exl3)
+ self.ey = None if exl3 else torch.empty((rows, slots, D), dtype=bf if prefill else f32, device=dev)
+ self.plan = None if exl3 else grouped.Plan(rows, slots, c.experts + 1, dev, prefill=prefill)
self.exl3 = None
if c.quant == "exl3": # EXL3 routed experts, and the shared expert as a BF16 MLP
- from .exl3_mm import Scratch
-
sl = c.shared_width // w.world
- self.exl3 = Scratch(rows, slots, D, ml, dev)
- self.exl3_rows = rows
+ shape = SimpleNamespace(dims=D, width=ml, count=c.experts)
+ self.exl3 = Exl3Scratch(shape, rows, slots, device=dev)
+ self.ey = self.exl3.y.view(rows, slots, D)
self.sgu = torch.empty((rows, 2 * sl), dtype=bf, device=dev)
self.sact = torch.empty((rows, sl), dtype=bf, device=dev)
self.sxs = torch.empty((rows, sl // 64), dtype=f32, device=dev)
@@ -394,14 +395,11 @@ def moe_block(layer: LayerW, w: Weights, b: Buffers, R: int) -> torch.Tensor:
with prof.timed("moe: route"):
glue.router(b.normed[:R], m.router, b.mlog[:R])
glue.select(b.mlog[:R], m.bias, b.pick[:R], b.wts[:R], c.top_k, c.experts, c.routed_scale, c.norm_topk)
- grouped.route(b.pick[:R], b.plan)
+ if m.shared is None:
+ grouped.route(b.pick[:R], b.plan)
if m.shared is not None:
# EXL3: the routed slots through the trellis kernels, the shared expert (last slot) through BF16 matmuls
- from . import exl3_mm
-
- from . import exl3_generic
- ey_flat = exl3_generic.routed(b.normed[:R], b.pick, m.experts, b.exl3.rows, c.limit)
- b.ey.view(-1, c.hidden)[:ey_flat.shape[0]].copy_(ey_flat)
+ exl3_generic.routed(b.normed[:R], b.pick, m.experts, b.exl3, R, c.limit)
s = m.shared
mm(b, b.normed[:R], s.gu, b.xs[:R], b.sgu[:R])
glue.swiglu(b.sgu[:R], b.sact[:R], b.sxs[:R], c.limit)
diff --git a/src/tensorfold/families/glm5_next/cuda/weights.py b/src/tensorfold/families/glm5_next/cuda/weights.py
index 03771a39e..7146e9a85 100644
--- a/src/tensorfold/families/glm5_next/cuda/weights.py
+++ b/src/tensorfold/families/glm5_next/cuda/weights.py
@@ -11,7 +11,7 @@
from tensorfold.cuda import experts as grouped
-from .exl3_mm import Exl3Experts, words as exl3_words
+from tensorfold.cuda.exl3.experts import Exl3RoutedExperts as Exl3Experts
from . import latent
from .qmm import B16, Q4, as_i32, make_b16, make_q4, quantize4, stack_b16, stack_q4
diff --git a/tests/cuda/glm_decision_tp.py b/tests/cuda/glm_decision_tp.py
index 0ff0c77d6..7f031d614 100644
--- a/tests/cuda/glm_decision_tp.py
+++ b/tests/cuda/glm_decision_tp.py
@@ -48,7 +48,7 @@ def main():
assert e.score_labels(decision, labels) == reference
assert e.live == [] and e.e.st.pos == 0 and e.drafter.context_end == 0
warm, stats = _generate(e, after, None, policy=policy, tokens=16)
- assert stats["cached"] == (len(prompt) if budget and policy == "2" else 0)
+ assert stats["cached"] == (len(prompt) - 1 if budget and policy == "2" else 0)
_generate(e, other, None, policy=policy, tokens=16)
switched, _ = _generate(e, after + [33], None, policy=policy, tokens=16)
cold, _ = _generate(e, after, None, draft=False, tokens=16)
diff --git a/tests/cuda/test_exl3_any_width_glm.py b/tests/cuda/test_exl3_any_width_glm.py
index 30ec5904e..644a33015 100644
--- a/tests/cuda/test_exl3_any_width_glm.py
+++ b/tests/cuda/test_exl3_any_width_glm.py
@@ -1,23 +1,4 @@
-"""GLM-shaped 3-bit (k2 = 6 half-bits) experts through the universal path.
-
-The universal grouped kernel (`tensorfold/cuda/exl3/experts.*`) reads a bit
-width per expert and was tested against a float64 reference at every width
-(`test_mixed_k_rows_are_independent_and_match_the_reference`). GLM's own
-family does not route through it: `families/glm5_next/cuda/exl3_mm.py` pins
-the 4-bit-mcg layout (`words()` requires `int16 [..., 64]`), so a 3-bit
-checkpoint refuses with `only 4-bit EXL3 trellises ... are supported` even
-though the universal kernel serves it. These tests pin the universal path at
-GLM's shapes and bit width, and the two properties an integrator relies on
-before committing to it: determinism across calls, and a row's output being
-the same alone and inside a window.
-
-The layer fixture mirrors a real 3-bit pack's on-disk layout: trellis int16
-`[K/16, N/16, 16 * k2]` words (48 at k2 = 6), fp16 `suh`/`svh` scales, 288
-routed experts at D = 4096, I = 1024 per rank (the pack is rank-split), with
-`GLM_GATEUP`/`GLM_DOWN` tile configs and `ACT_BF16` — the same arithmetic the
-dedicated path runs, verified bit-identical by `test_glm_shaped_bit_identical`
-at 4-bit.
-"""
+"""GLM-shaped 3-bit EXL3 experts: format bits, repeatability, independent rows and buffer ownership."""
import pytest
import torch
@@ -50,21 +31,17 @@ def scale(n, mag):
def test_dequant_identity_against_format_unpack():
- """The kernels' own decode of a 3-bit trellis equals the format's reference
- decoder, bit for bit. If the pack reads wrong here, everything downstream is
- built on a wrong W_q and the rest of the file proves nothing."""
+ """The kernel decodes the same fp16 bits as format.unpack for a 3-bit trellis."""
from tensorfold.cuda.exl3 import experts
from tensorfold.cuda.exl3.format import unpack as fmt_unpack
t = torch.randint(-32768, 32768, (D // 16, I // 16, 8 * K2), dtype=torch.int16).cuda()
ref = torch.from_numpy(fmt_unpack(t.cpu().numpy(), 3.0, "mcg")).cuda()
- assert torch.equal(experts.dequant(t, "mcg"), ref)
+ assert torch.equal(experts.dequant(t, "mcg").view(torch.uint8), ref.view(torch.uint8))
def test_prepared_layer_reports_k2_from_the_checkpoint():
- """`prepare` derives k2 from each tensor's own last dim; a 3-bit pack must come
- out as k2 = 6 half-bits, not the 4-bit assumption of the dedicated wrapper."""
- from tensorfold.cuda.exl3 import experts
+ """Each checkpoint tensor supplies its own bit width."""
ex = _layer()
assert ex.k2_gu == (K2, K2)
@@ -79,14 +56,13 @@ def test_routed_is_deterministic_across_calls():
x = torch.randn(4, D, dtype=torch.bfloat16, device="cuda")
pick = torch.zeros((4, 4), dtype=torch.int32, device="cuda")
s = experts.Scratch(ex, 4, 4, device="cuda")
- y1 = experts.routed(x, pick, None, ex, s, None, 4, act_mode=experts.ACT_BF16)
+ y1 = experts.routed(x, pick, None, ex, s, None, 4, act_mode=experts.ACT_BF16).clone()
y2 = experts.routed(x, pick, None, ex, s, None, 4, act_mode=experts.ACT_BF16)
assert torch.equal(y1, y2), "repeated calls must be bit-identical"
def test_rows_are_independent_of_batch_composition():
- """A row alone must equal its output inside a window, whatever the other rows
- hold: the grouped launch groups by expert and reads rows by index."""
+ """A row gives the same bits alone, duplicated or inside a window."""
from tensorfold.cuda.exl3 import experts
ex = _layer()
@@ -108,9 +84,7 @@ def test_rows_are_independent_of_batch_composition():
def test_shared_expert_pairs_are_skipped():
- """pick == E means the shared expert: the universal kernel leaves those pairs
- unwritten, matching the dedicated path's `slot == slots - 1` guard. The caller
- combines the slot from the BF16 shared MLP before reading it."""
+ """Shared-expert picks are left for the caller to supply."""
from tensorfold.cuda.exl3 import experts
ex = _layer()
@@ -120,3 +94,20 @@ def test_shared_expert_pairs_are_skipped():
y = experts.routed(x, pick, None, ex, s, None, 2, act_mode=experts.ACT_BF16)
torch.cuda.synchronize()
assert torch.all(y == 0)
+
+
+def test_family_adapter_uses_live_rows_and_owned_buffers():
+ from tensorfold.cuda.exl3 import experts
+ from tensorfold.families.glm5_next.cuda import exl3_generic
+
+ ex = _layer()
+ x = torch.randn(4, D, dtype=torch.bfloat16, device="cuda")
+ pick = torch.arange(16, dtype=torch.int32, device="cuda").view(4, 4) % E
+ wide = experts.Scratch(ex, 16, 4, device="cuda")
+ narrow = experts.Scratch(ex, 4, 4, device="cuda")
+ wide.y.fill_(123)
+ got = exl3_generic.routed(x, pick, ex, wide, 4, 7.0).clone()
+ expected = experts.routed(x, pick, None, ex, narrow, None, 4, 7.0, act_mode=experts.ACT_BF16)
+ assert got.shape == (16, D) and torch.equal(got.view(torch.uint8), expected.view(torch.uint8))
+ assert torch.all(wide.y[16:] == 123)
+ assert wide.y.data_ptr() != narrow.y.data_ptr()
diff --git a/tests/cuda/test_glm_engine.py b/tests/cuda/test_glm_engine.py
index 9328f78cc..8ee70b4da 100644
--- a/tests/cuda/test_glm_engine.py
+++ b/tests/cuda/test_glm_engine.py
@@ -377,9 +377,10 @@ def test_exl3_checkpoint_drafted_equals_serial(engine_x, sampling):
"""An EXL3 checkpoint through the same engine: every policy's reply equals serial decoding."""
from tensorfold.families.glm5_next.cuda.engine import EXL3_AUTO, encode_policy
- from tensorfold.families.glm5_next.cuda.exl3_mm import Exl3Experts
+ from tensorfold.cuda.exl3.experts import Exl3RoutedExperts
- assert isinstance(engine_x.w.layers[1].moe.experts, Exl3Experts) and engine_x.w.layers[1].moe.shared is not None
+ assert isinstance(engine_x.w.layers[1].moe.experts, Exl3RoutedExperts)
+ assert engine_x.w.layers[1].moe.shared is not None
assert engine_x._effective(encode_policy("auto")) == encode_policy(EXL3_AUTO) # the default drafts DFlash2
assert engine_x._effective(encode_policy("auto:1:1:0")) == encode_policy("auto:1:1:0")
prompt = list(np.random.default_rng(8).integers(0, 1000, size=45))
@@ -416,7 +417,7 @@ def test_decision_between_chats_preserves_replies(engine_f, cache_bytes, policy)
immediate, stats = _generate(e, after, sampling, policy=policy, tokens=16)
# Saved attention rows retain MTP, but not DFlash2's unsaved draft cache.
- assert stats["cached"] == (len(prompt) if cache_bytes and policy == "2" else 0)
+ assert stats["cached"] == (len(prompt) - 1 if cache_bytes and policy == "2" else 0)
_generate(e, other, sampling, policy=policy, tokens=16)
switched, _ = _generate(e, after + [33], sampling, policy=policy, tokens=16)
_forget(e)
diff --git a/tests/test_cuda_geometry.py b/tests/test_cuda_geometry.py
index 132653a36..372f7a40c 100644
--- a/tests/test_cuda_geometry.py
+++ b/tests/test_cuda_geometry.py
@@ -24,6 +24,9 @@ def __add__(self, other): # arange(rows)[:, None] + arange(k): the [r
return Allocation(self.shape + other.shape, self.dtype, self.device)
def contiguous(self):
return self
+ def view(self, *shape):
+ assert math.prod(shape) == self.numel()
+ return Allocation(shape, self.dtype, self.device)
@pytest.fixture
@@ -239,7 +242,8 @@ def test_mla_exl3_scratch_and_buffers_are_budgeted(monkeypatch, allocations, mtp
mod = importlib.import_module("tensorfold.families.glm5_next.cuda.forward")
names = ("kda", "latent", "attention", "exl3_mm")
mods = [mod] + [importlib.import_module(f"tensorfold.families.glm5_next.cuda.{n}") for n in names]
- for m in mods + [importlib.import_module("tensorfold.cuda.experts")]:
+ for m in mods + [importlib.import_module("tensorfold.cuda.experts"),
+ importlib.import_module("tensorfold.cuda.exl3.experts")]:
monkeypatch.setattr(m, "torch", fake)
monkeypatch.setattr(mods[2], "ENABLED", True)
exl3_mm = mods[-1]
From 3d7a14781acc94a5397d36219d2be84f23465ca7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
<600316+jschmied@users.noreply.github.com>
Date: Thu, 1 Oct 2026 06:45:51 +0200
Subject: [PATCH 090/177] qwen4_exp: refuse quantized bytes where the loader
reads bf16 values
The block-FP8 path's weight_bf16 cast anything that was not block FP8, and dense() returned a weight
without a scale as bf16 whatever its dtype, so an unsupported format (an MXFP8 head, W4A16 NVFP4 DeltaNet
projections) would load with the right shapes and decode fluent garbage. Both now refuse a weight that is
not bf16, fp16 or fp32 by name, and the checkpoint check refuses NVFP4 outside the routed experts before
any weight is read.
---
src/tensorfold/families/qwen4_exp/__init__.py | 6 ++++++
.../families/qwen4_exp/cuda/weights.py | 16 +++++++++++++++-
tests/test_flashnext_plain_weights.py | 18 ++++++++++++++++++
tests/test_hub_and_checks.py | 7 +++++++
4 files changed, 46 insertions(+), 1 deletion(-)
create mode 100644 tests/test_flashnext_plain_weights.py
diff --git a/src/tensorfold/families/qwen4_exp/__init__.py b/src/tensorfold/families/qwen4_exp/__init__.py
index d34620357..81a81ef02 100644
--- a/src/tensorfold/families/qwen4_exp/__init__.py
+++ b/src/tensorfold/families/qwen4_exp/__init__.py
@@ -67,6 +67,12 @@ def check(model_dir: Path) -> None:
raise ValueError(f"TensorFold's Flash Next kernels read NVFP4 (ModelOpt FP4) weights in blocks of 16, the "
f"other linears bf16, MXFP8 or 128x128-block FP8 ({', '.join(NVFP4_MODELS)}); this checkpoint has "
+ describe_quantization(config) + f". {OWN_MODEL_HELP}")
+ # NVFP4 is read in the routed experts only: anywhere else its packed bytes would load as bf16 values
+ outside = sorted(name for name, layer in (found.get("quantized_layers") or {}).items()
+ if "NVFP4" in str(layer.get("quant_algo", "")).upper() and "experts" not in name.split("."))
+ if outside:
+ raise ValueError(f"TensorFold's Flash Next kernels read NVFP4 in the routed experts only; this checkpoint "
+ f"has it on {len(outside)} other layer(s), e.g. {outside[0]}. {OWN_MODEL_HELP}")
if (Path(model_dir) / "model.safetensors.index.json").is_file() and not has_mtp(model_dir):
print("[tensorfold] this NVFP4 checkpoint has no MTP head: decoding without MTP drafts", flush=True)
return
diff --git a/src/tensorfold/families/qwen4_exp/cuda/weights.py b/src/tensorfold/families/qwen4_exp/cuda/weights.py
index 3c33a559b..e96339b7b 100644
--- a/src/tensorfold/families/qwen4_exp/cuda/weights.py
+++ b/src/tensorfold/families/qwen4_exp/cuda/weights.py
@@ -17,6 +17,18 @@
AttnW, Config, GDNW, HC, LayerW, MoEW, MTPW, PLEW, Weights, draft_token_ids, stop_ids) # noqa: F401 (re-exported)
+_PLAIN = (torch.bfloat16, torch.float16, torch.float32)
+
+
+def _plain(name: str, w: torch.Tensor) -> torch.Tensor:
+ """``w`` if it holds real values a bf16 linear can take, else a refusal: quantized bytes cast to bf16 decode garbage."""
+
+ if w.dtype not in _PLAIN:
+ raise ValueError(f"{name}: {str(w.dtype).removeprefix('torch.')} weights without a scale this loader reads; "
+ "Flash Next reads its non-expert linears as bf16, MXFP8 or 128x128-block FP8")
+ return w
+
+
def load(model_dir: str | Path, device: str = "cuda", *, mtp: bool = True, tp: tuple[int, int] | None = None,
draft_vocab: int | str | None = None, ple_on_ssd: bool = False, table_reads: list | None = None) -> Weights:
"""Load rank ``tp``'s head, expert-width and vocabulary shares while replicating other weights; ``draft_vocab`` restricts draft scoring to default/file ids or ids below N, with None using all ids."""
@@ -89,6 +101,8 @@ def dense(name: str, rows=None, cols: slice | None = None):
s = raw(name + ".weight_scale") if w.dtype == torch.float8_e4m3fn else None
if s is not None and s.dtype != torch.uint8:
raise ValueError(f"{name}: FP8 with a per-tensor scale; Flash Next reads MXFP8 (a scale every 32 inputs)")
+ if s is None:
+ _plain(name, w)
if rows is not None:
w, s = w[rows], None if s is None else s[rows]
if cols is not None:
@@ -193,7 +207,7 @@ def weight_bf16(name: str, index: torch.Tensor | None = None) -> torch.Tensor:
full = raw(name + ".weight")
w = full if index is None else full.index_select(0, index)
if w.dtype != torch.float8_e4m3fn or not rd.has(prefix + name + ".weight_scale_inv"):
- return w.to(torch.bfloat16)
+ return _plain(name, w).to(torch.bfloat16)
from tensorfold.cuda.nvfp4.linear import Fp8BlockLinear
cols = Fp8BlockLinear.column_scales(raw(name + ".weight_scale_inv"), *full.shape)
diff --git a/tests/test_flashnext_plain_weights.py b/tests/test_flashnext_plain_weights.py
new file mode 100644
index 000000000..81ee3ac88
--- /dev/null
+++ b/tests/test_flashnext_plain_weights.py
@@ -0,0 +1,18 @@
+"""The Flash Next loader refuses quantized bytes where it reads bf16 values, instead of casting them to garbage."""
+
+import pytest
+
+pytestmark = pytest.mark.torch
+
+
+def test_real_values_pass_and_quantized_bytes_are_refused():
+ import torch
+
+ from tensorfold.families.qwen4_exp.cuda.weights import _plain
+
+ for dtype in (torch.bfloat16, torch.float16, torch.float32):
+ w = torch.ones(2, 3, dtype=dtype)
+ assert _plain("x.weight", w) is w
+ for dtype in (torch.uint8, torch.int8, torch.float8_e4m3fn):
+ with pytest.raises(ValueError, match="x.weight: .* without a scale"):
+ _plain("x.weight", torch.zeros(2, 3, dtype=dtype))
diff --git a/tests/test_hub_and_checks.py b/tests/test_hub_and_checks.py
index ce17da6d8..4a73564ed 100644
--- a/tests/test_hub_and_checks.py
+++ b/tests/test_hub_and_checks.py
@@ -211,6 +211,13 @@ def test_flash_next_reads_the_nvfp4_checkpoint_and_refuses_other_fp4_blocks(tmp_
(tmp_path / "config.json").write_text(json.dumps(other))
with pytest.raises(ValueError, match="blocks of 16"):
qwen4_exp.check(tmp_path)
+ # NVFP4 outside the routed experts (e.g. a W4A16 DeltaNet projection) is refused before any weight is read
+ gdn_fp4 = json.loads(json.dumps(mixed))
+ gdn_fp4["quantization_config"]["quantized_layers"]["model.language_model.layers.0.linear_attn.in_proj_qkv"] = {
+ "quant_algo": "W4A16_NVFP4", "group_size": 16}
+ (tmp_path / "config.json").write_text(json.dumps(gdn_fp4))
+ with pytest.raises(ValueError, match="routed experts only.*in_proj_qkv"):
+ qwen4_exp.check(tmp_path)
def test_models_lists_the_tested_checkpoints(capsys):
From b3c6ece0797f3fc5dcb05f8d2d3fc7679034bc4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
<600316+jschmied@users.noreply.github.com>
Date: Thu, 1 Oct 2026 10:25:41 +0200
Subject: [PATCH 091/177] qwen4_exp cuda: read FP8 experts in the MTP drafter
A ModelOpt export with FP8 MTP experts (per expert: e4m3 with a scale a tensor or a row, or 128x128
blocks) was refused at the check, and the loader would have read them as NVFP4. The drafter's experts are
already re-quantized to NVFP4 at load from bf16 (moe4_from_bf16), and drafts only propose tokens the target
verifies, so FP8 ones are dequantized to bf16 and take the same path. FP8 in the main layers' experts stays
refused: re-quantizing those would change replies. The check accepts FP8 on MTP expert layers only.
Tests: a tiny checkpoint with FP8 MTP experts drafts exactly as one whose stacked bf16 experts hold the same
dequantized values, and its drafts keep the serial tokens; the check accepts FP8 on MTP experts and refuses
it on the main experts.
---
src/tensorfold/families/qwen4_exp/__init__.py | 4 ++-
.../families/qwen4_exp/cuda/weights.py | 22 ++++++++++++-
tests/cuda/nvfp4_tiny.py | 33 ++++++++++++++++---
tests/cuda/test_flashnext_nvfp4_loader.py | 25 ++++++++++++++
tests/test_hub_and_checks.py | 11 +++++++
5 files changed, 89 insertions(+), 6 deletions(-)
diff --git a/src/tensorfold/families/qwen4_exp/__init__.py b/src/tensorfold/families/qwen4_exp/__init__.py
index 81a81ef02..7e3bf8f65 100644
--- a/src/tensorfold/families/qwen4_exp/__init__.py
+++ b/src/tensorfold/families/qwen4_exp/__init__.py
@@ -59,7 +59,9 @@ def check(model_dir: Path) -> None:
# the CUDA engine's NVFP4 route: NVFP4 experts in blocks of 16, other linears bf16, MXFP8, block FP8 or NVFP4
found = config.get("quantization") or config.get("quantization_config") or {}
algo = str(found.get("quant_algo") or "NVFP4").upper()
- layers = {str(v.get("quant_algo", "")).upper() for v in (found.get("quantized_layers") or {}).values()}
+ # FP8 is read in the MTP drafter's experts only (dequantized and re-quantized at load: they only draft)
+ layers = {str(v.get("quant_algo", "")).upper() for k, v in (found.get("quantized_layers") or {}).items()
+ if not (str(v.get("quant_algo", "")).upper() == "FP8" and {"mtp", "experts"} <= set(k.split(".")))}
algos = layers if algo == "MIXED_PRECISION" else {algo}
weights = [g.get("weights") or {} for g in (found.get("config_groups") or {}).values()]
fp4 = {int(w.get("group_size", 16)) for w in weights if int(w.get("num_bits", 4)) == 4}
diff --git a/src/tensorfold/families/qwen4_exp/cuda/weights.py b/src/tensorfold/families/qwen4_exp/cuda/weights.py
index e96339b7b..59a5f6cdc 100644
--- a/src/tensorfold/families/qwen4_exp/cuda/weights.py
+++ b/src/tensorfold/families/qwen4_exp/cuda/weights.py
@@ -263,7 +263,27 @@ def moe_nvfp4(name: str) -> MoEW:
shared = (raw(se + "gate_proj.weight").to(torch.bfloat16)[lo:hi].contiguous(),
raw(se + "up_proj.weight").to(torch.bfloat16)[lo:hi].contiguous(),
raw(se + "down_proj.weight").to(torch.bfloat16)[:, dlo * gs:dhi * gs].contiguous())
- if rd.has(prefix + f"{name}.experts.0.gate_proj.weight"): # the main layers: per-expert FP4
+ fp8_drafter = (rd.has(prefix + f"{name}.experts.0.gate_proj.weight")
+ and raw(f"{name}.experts.0.gate_proj.weight").dtype == torch.float8_e4m3fn)
+ if fp8_drafter and not name.startswith("mtp."):
+ raise ValueError(f"{name}: FP8 routed experts; Flash Next reads the routed experts as NVFP4 (FP8 only in "
+ "the MTP drafter, which is re-quantized at load)")
+ if fp8_drafter: # the MTP layer's per-expert FP8: dequantized, then drafted as below
+ def expert_bf16(base: str) -> torch.Tensor:
+ if rd.has(prefix + base + ".weight_scale_inv"): # 128x128 blocks
+ return weight_bf16(base)
+ w = raw(base + ".weight").float()
+ s = raw(base + ".weight_scale").float() # per tensor, or one scale a row
+ return (w * (s.reshape(-1, 1) if s.numel() > 1 else s)).to(torch.bfloat16)
+
+ def stacked(proj: str) -> torch.Tensor:
+ return torch.stack([expert_bf16(f"{name}.experts.{i}.{proj}") for i in range(e)])
+
+ gate, up, dn = stacked("gate_proj"), stacked("up_proj"), stacked("down_proj")
+ if world > 1:
+ gate, up, dn = gate[:, lo:hi], up[:, lo:hi], dn[:, :, dlo * gs:dhi * gs]
+ moe4 = nvfp4_moe.moe4_from_bf16(torch.cat([gate, up], dim=1), dn, shared)
+ elif rd.has(prefix + f"{name}.experts.0.gate_proj.weight"): # the main layers: per-expert FP4
def stack(proj: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
w = torch.stack([raw(f"{name}.experts.{i}.{proj}.weight") for i in range(e)])
s = torch.stack([raw(f"{name}.experts.{i}.{proj}.weight_scale") for i in range(e)])
diff --git a/tests/cuda/nvfp4_tiny.py b/tests/cuda/nvfp4_tiny.py
index bf8f5341c..3d6ebd978 100644
--- a/tests/cuda/nvfp4_tiny.py
+++ b/tests/cuda/nvfp4_tiny.py
@@ -53,8 +53,11 @@ def write(dir: Path, *, layers: int = 2, experts: int = 2, vocab: int = 256, hid
heads: int = 2, kv_heads: int = 2, hd: int = 64, nk: int = 8, nv: int = 24, dk: int = 128, dv: int = 128,
moe_width: int = 128, shared_width: int = 64, streams: int = 4, low: int = 64,
ple: bool = True, mtp: bool = True, seed: int = 0, prefix: str = "", ple_bf16: bool = False,
- mxfp8: bool = False, ple_nvfp4: bool = False, centred: bool = False, fp8block: bool = False) -> Path:
- """A tiny ModelOpt checkpoint: ``mxfp8``, ``ple_nvfp4``, ``centred`` norms or ``fp8block`` (block FP8 beside bf16)."""
+ mxfp8: bool = False, ple_nvfp4: bool = False, centred: bool = False, fp8block: bool = False,
+ mtp_experts: str = "bf16") -> Path:
+ """A tiny ModelOpt checkpoint: ``mxfp8``, ``ple_nvfp4``, ``centred`` norms or ``fp8block`` (block FP8 beside bf16);
+ ``mtp_experts`` "bf16" (stacked), "fp8" (per-expert e4m3, a scale a tensor) or "fp8_dequant" (stacked bf16 of
+ exactly those e4m3 values times their scales: the same draws)."""
dir.mkdir(parents=True, exist_ok=True)
rng = torch.Generator().manual_seed(seed)
@@ -189,8 +192,30 @@ def linear(name: str, n: int, k: int, *, fp4: bool, mx: bool = False, blk: bool
for proj, n_, k_ in (("gate_proj", moe_width, hidden), ("up_proj", moe_width, hidden),
("down_proj", hidden, moe_width)):
linear(f"mtp.layers.0.mlp.shared_expert.{proj}", n_, k_, fp4=False)
- add("mtp.layers.0.mlp.experts.gate_up_proj", rand(experts, 2 * moe_width, hidden))
- add("mtp.layers.0.mlp.experts.down_proj", rand(experts, hidden, moe_width))
+ if mtp_experts == "bf16":
+ add("mtp.layers.0.mlp.experts.gate_up_proj", rand(experts, 2 * moe_width, hidden))
+ add("mtp.layers.0.mlp.experts.down_proj", rand(experts, hidden, moe_width))
+ else: # e4m3 with one fp32 scale a tensor, per expert
+ def fp8(n_: int, k_: int) -> tuple[torch.Tensor, torch.Tensor]:
+ w = rand(n_, k_).float()
+ scale = w.abs().max() / 448.0
+ return (w / scale).to(torch.float8_e4m3fn), scale.reshape(())
+
+ projs = {p_: [fp8(*shape) for _ in range(experts)] for p_, shape in
+ (("gate_proj", (moe_width, hidden)), ("up_proj", (moe_width, hidden)),
+ ("down_proj", (hidden, moe_width)))}
+ if mtp_experts == "fp8":
+ for p_, items in projs.items():
+ for i, (codes, scale) in enumerate(items):
+ add(f"mtp.layers.0.mlp.experts.{i}.{p_}.weight", codes)
+ add(f"mtp.layers.0.mlp.experts.{i}.{p_}.weight_scale", scale)
+ else:
+ def deq(items):
+ return torch.stack([(c.float() * s_).to(torch.bfloat16) for c, s_ in items])
+
+ add("mtp.layers.0.mlp.experts.gate_up_proj", torch.cat([deq(projs["gate_proj"]),
+ deq(projs["up_proj"])], dim=1))
+ add("mtp.layers.0.mlp.experts.down_proj", deq(projs["down_proj"]))
for proj, n_, k_ in (("q_proj", 2 * heads * hd, hidden), ("k_proj", kv_heads * hd, hidden),
("v_proj", kv_heads * hd, hidden), ("o_proj", hidden, heads * hd),
("indexer.index_qk_proj", (4 + 1) * 128, hidden)):
diff --git a/tests/cuda/test_flashnext_nvfp4_loader.py b/tests/cuda/test_flashnext_nvfp4_loader.py
index dd778dd6e..0f5e33b0f 100644
--- a/tests/cuda/test_flashnext_nvfp4_loader.py
+++ b/tests/cuda/test_flashnext_nvfp4_loader.py
@@ -368,3 +368,28 @@ def test_an_nvfp4_checkpoint_refuses_two_ranks(tiny: Path, monkeypatch) -> None:
monkeypatch.setattr("tensorfold.cuda.comm.NCCL", lambda *a, **k: pytest.fail("the ranks started"))
with pytest.raises(ValueError, match="one GPU"):
FlashNextEngine(tiny, tp=2, rank=0, master="127.0.0.1")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="the loader builds CUDA tensors")
+def test_fp8_drafter_experts_draft_as_their_dequantized_bf16(tmp_path: Path) -> None:
+ """Per-expert FP8 MTP experts (e4m3, a scale a tensor) load: they are dequantized and re-quantized as the bf16
+ drafter is, so they draft exactly as stacked bf16 experts holding the same values, and drafts keep serial tokens."""
+
+ from tensorfold.engine.exact_sampling import Sampling
+ from tensorfold.families.qwen4_exp.cuda.decode import Engine, mtp_decode, prefill, serial_decode
+ from tensorfold.families.qwen4_exp.cuda.weights import load
+
+ prompt = [5, 17, 99, 250, 7, 64, 30, 11, 12, 13]
+ sampling = Sampling(seed=7, top_k=20, top_p=0.95)
+ runs = {}
+ for kind in ("fp8", "fp8_dequant"):
+ w = load(write(tmp_path / kind, hidden=512, mtp_experts=kind), mtp=True, draft_vocab=128)
+ e = Engine(w, capacity=256, max_rows=8, prefill_rows=16, graphs=False)
+ first = prefill(e, prompt, sampling)
+ ref = serial_decode(e, first, 16, sampling).tokens
+ assert prefill(e, prompt, sampling) == first # the prompt's state again
+ out = mtp_decode(e, first, 16, sampling, depth=4, confidence=0.0)
+ assert out.tokens == ref, kind
+ runs[kind] = out
+ same = ("tokens", "rounds", "drafted", "accepted", "keeps", "widths") # the same drafts, accepted the same
+ assert [getattr(runs["fp8"], k) for k in same] == [getattr(runs["fp8_dequant"], k) for k in same]
diff --git a/tests/test_hub_and_checks.py b/tests/test_hub_and_checks.py
index 4a73564ed..83bfa0d65 100644
--- a/tests/test_hub_and_checks.py
+++ b/tests/test_hub_and_checks.py
@@ -218,6 +218,17 @@ def test_flash_next_reads_the_nvfp4_checkpoint_and_refuses_other_fp4_blocks(tmp_
(tmp_path / "config.json").write_text(json.dumps(gdn_fp4))
with pytest.raises(ValueError, match="routed experts only.*in_proj_qkv"):
qwen4_exp.check(tmp_path)
+ # FP8 in the MTP drafter's experts is read (dequantized and re-quantized at load); in the main experts it is not
+ mtp_fp8 = json.loads(json.dumps(mixed))
+ mtp_fp8["quantization_config"]["quantized_layers"]["mtp.layers.0.mlp.experts"] = {"quant_algo": "FP8"}
+ (tmp_path / "config.json").write_text(json.dumps(mtp_fp8))
+ qwen4_exp.check(tmp_path)
+ main_fp8 = json.loads(json.dumps(mixed))
+ main_fp8["quantization_config"]["quantized_layers"]["model.language_model.layers.0.mlp.experts"] = {
+ "quant_algo": "FP8"}
+ (tmp_path / "config.json").write_text(json.dumps(main_fp8))
+ with pytest.raises(ValueError, match="blocks of 16"):
+ qwen4_exp.check(tmp_path)
def test_models_lists_the_tested_checkpoints(capsys):
From 7347e513bb564d29fb2ce756b8d6f4be1b074857 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:36:56 +0100
Subject: [PATCH 092/177] fix(flash-next): FP8 draft experts are validated, and
preflight keeps NVFP4 n-gram tables and skips unused vision weights
---
src/tensorfold/families/qwen4_exp/__init__.py | 22 +++--
.../families/qwen4_exp/cuda/weights.py | 30 ++++--
tests/cuda/nvfp4_tiny.py | 29 ++++--
tests/cuda/test_flashnext_nvfp4_loader.py | 99 +++++++++++++++++--
tests/test_flashnext_modelopt_check.py | 29 ++++++
tests/test_hub_and_checks.py | 8 +-
6 files changed, 183 insertions(+), 34 deletions(-)
create mode 100644 tests/test_flashnext_modelopt_check.py
diff --git a/src/tensorfold/families/qwen4_exp/__init__.py b/src/tensorfold/families/qwen4_exp/__init__.py
index 7e3bf8f65..1f86048ae 100644
--- a/src/tensorfold/families/qwen4_exp/__init__.py
+++ b/src/tensorfold/families/qwen4_exp/__init__.py
@@ -16,7 +16,7 @@
EXL3_VARIANT = "any" # every EXL3 codebook and width (tensorfold.families.EXL3_VARIANT_ANY)
KERNEL_PACKAGE = "tensorfold.kernels.qwen.flash_next.v1"
KERNEL_VERSION = "v1"
-# The CLI sets these before MLX starts, respecting environment overrides, to keep expert bindings from ending each command buffer.
+# The CLI sets defaults before MLX starts so expert bindings do not end each command buffer.
MLX_ENV = {"MLX_MAX_OPS_PER_BUFFER": "200", "MLX_MAX_MB_PER_BUFFER": "100000"}
@@ -56,7 +56,9 @@ def check(model_dir: Path) -> None:
print("[tensorfold] this EXL3 checkpoint has no MTP head: decoding without MTP drafts", flush=True)
return
if quant_method(config) == "modelopt":
- # the CUDA engine's NVFP4 route: NVFP4 experts in blocks of 16, other linears bf16, MXFP8, block FP8 or NVFP4
+ from tensorfold.vision.qwen_checkpoint import vision_key
+
+ # NVFP4 experts and n-gram tables; other linears are bf16, MXFP8 or block FP8
found = config.get("quantization") or config.get("quantization_config") or {}
algo = str(found.get("quant_algo") or "NVFP4").upper()
# FP8 is read in the MTP drafter's experts only (dequantized and re-quantized at load: they only draft)
@@ -67,14 +69,18 @@ def check(model_dir: Path) -> None:
fp4 = {int(w.get("group_size", 16)) for w in weights if int(w.get("num_bits", 4)) == 4}
if not algos <= {"NVFP4", "W4A16_NVFP4", "MXFP8", "FP8_PB_WO"} or fp4 - {16}:
raise ValueError(f"TensorFold's Flash Next kernels read NVFP4 (ModelOpt FP4) weights in blocks of 16, the "
- f"other linears bf16, MXFP8 or 128x128-block FP8 ({', '.join(NVFP4_MODELS)}); this checkpoint has "
+ f"other linears bf16, MXFP8 or 128x128-block FP8 ({', '.join(NVFP4_MODELS)}); "
+ "this checkpoint has "
+ describe_quantization(config) + f". {OWN_MODEL_HELP}")
- # NVFP4 is read in the routed experts only: anywhere else its packed bytes would load as bf16 values
+ # N-gram tables have their own NVFP4 reader; other non-expert layers would cast packed bytes to bf16.
outside = sorted(name for name, layer in (found.get("quantized_layers") or {}).items()
- if "NVFP4" in str(layer.get("quant_algo", "")).upper() and "experts" not in name.split("."))
+ if "NVFP4" in str(layer.get("quant_algo", "")).upper() and "experts" not in name.split(".")
+ and ".ple.ple_embedding.ngram_embedding." not in name + "."
+ and vision_key(name) is None)
if outside:
- raise ValueError(f"TensorFold's Flash Next kernels read NVFP4 in the routed experts only; this checkpoint "
- f"has it on {len(outside)} other layer(s), e.g. {outside[0]}. {OWN_MODEL_HELP}")
+ raise ValueError("TensorFold's Flash Next kernels read NVFP4 in routed experts and n-gram tables only; "
+ f"this checkpoint has it on {len(outside)} other layer(s), e.g. {outside[0]}. "
+ f"{OWN_MODEL_HELP}")
if (Path(model_dir) / "model.safetensors.index.json").is_file() and not has_mtp(model_dir):
print("[tensorfold] this NVFP4 checkpoint has no MTP head: decoding without MTP drafts", flush=True)
return
@@ -161,7 +167,7 @@ def cuda_engine(model_dir: str | Path, *, drafter: str = "", tp: int = 1, rank:
master_port: int = 29551, no_drafts: bool = False, mtp_drafts: int | None = None,
mtp_confidence: float | None = None, context: int | None = None, ple_on_ssd: bool = False,
kv_dtype: str = "bf16", decode_share: float | None = None, **options: Any):
- """The CUDA engine: MTP chains verified exactly on one GPU or two (``tp=2``; start rank 1 first), keys and values bf16, int8 or int4."""
+ """Verify MTP on one or two CUDA GPUs; start rank 1 first for ``tp=2``, with bf16, int8 or int4 KV storage."""
from tensorfold.cuda.exl3.format import is_exl3
diff --git a/src/tensorfold/families/qwen4_exp/cuda/weights.py b/src/tensorfold/families/qwen4_exp/cuda/weights.py
index 59a5f6cdc..5159a73f3 100644
--- a/src/tensorfold/families/qwen4_exp/cuda/weights.py
+++ b/src/tensorfold/families/qwen4_exp/cuda/weights.py
@@ -1,4 +1,4 @@
-"""Pack affine group-32 MLX weights with the shared expert last, concatenate n-gram shards in order, and retain centered-norm gamma itself as fp32 after checking it is around one."""
+"""Load affine MLX or ModelOpt weights, shared experts and n-gram shards with fp32 centered-norm scales."""
from __future__ import annotations
@@ -21,7 +21,7 @@
def _plain(name: str, w: torch.Tensor) -> torch.Tensor:
- """``w`` if it holds real values a bf16 linear can take, else a refusal: quantized bytes cast to bf16 decode garbage."""
+ """Refuse quantized bytes where a bf16 linear needs real values."""
if w.dtype not in _PLAIN:
raise ValueError(f"{name}: {str(w.dtype).removeprefix('torch.')} weights without a scale this loader reads; "
@@ -31,7 +31,7 @@ def _plain(name: str, w: torch.Tensor) -> torch.Tensor:
def load(model_dir: str | Path, device: str = "cuda", *, mtp: bool = True, tp: tuple[int, int] | None = None,
draft_vocab: int | str | None = None, ple_on_ssd: bool = False, table_reads: list | None = None) -> Weights:
- """Load rank ``tp``'s head, expert-width and vocabulary shares while replicating other weights; ``draft_vocab`` restricts draft scoring to default/file ids or ids below N, with None using all ids."""
+ """Load rank ``tp``'s shares; ``draft_vocab`` selects default/file ids or ids below N, None scores all ids."""
import time
from dataclasses import replace
@@ -145,7 +145,7 @@ def hc_nvfp4(name: str, inject: bool) -> HC:
return HC(stack_b16(parts), b16(name + ".input_mix_weight_up"), cscale(name + ".hc_norm.weight"), inject)
def gdn_nvfp4(name: str) -> GDNW:
- """A DeltaNet block from the NVFP4 checkpoint (bf16 or MXFP8 linears): the rank's rows, conv and head vectors."""
+ """A rank's DeltaNet rows, conv and head vectors from bf16 or MXFP8 checkpoint linears."""
kl, vl = full.nk // world, full.nv // world
dk, dv = full.dk, full.dv
@@ -180,7 +180,7 @@ def attention_nvfp4(name: str) -> AttnW:
o)
def ple_nvfp4(name: str, ple_index: int) -> PLEW:
- """A PLE layer from the NVFP4 checkpoint: n-gram rows from bf16, FP8, NVFP4 or MLX 4-bit shards, the rest bf16."""
+ """A PLE layer with bf16, FP8, NVFP4 or MLX 4-bit n-gram shards and bf16 projections."""
if ple_on_ssd:
raise ValueError("--ple-on-ssd reads the MLX checkpoint's n-gram shards from disk; an NVFP4 checkpoint's "
@@ -268,13 +268,20 @@ def moe_nvfp4(name: str) -> MoEW:
if fp8_drafter and not name.startswith("mtp."):
raise ValueError(f"{name}: FP8 routed experts; Flash Next reads the routed experts as NVFP4 (FP8 only in "
"the MTP drafter, which is re-quantized at load)")
- if fp8_drafter: # the MTP layer's per-expert FP8: dequantized, then drafted as below
+ if fp8_drafter: # MTP experts dequantize before draft-only NVFP4 packing
def expert_bf16(base: str) -> torch.Tensor:
+ codes = raw(base + ".weight")
+ if codes.dtype != torch.float8_e4m3fn:
+ raise ValueError(f"{base}: expected FP8 e4m3 MTP expert weights")
if rd.has(prefix + base + ".weight_scale_inv"): # 128x128 blocks
return weight_bf16(base)
- w = raw(base + ".weight").float()
- s = raw(base + ".weight_scale").float() # per tensor, or one scale a row
- return (w * (s.reshape(-1, 1) if s.numel() > 1 else s)).to(torch.bfloat16)
+ if not rd.has(prefix + base + ".weight_scale"):
+ raise ValueError(f"{base}: FP8 MTP experts need a tensor, row or 128x128-block scale")
+ scale = raw(base + ".weight_scale")
+ if scale.dtype not in _PLAIN or scale.numel() not in (1, codes.shape[0]):
+ raise ValueError(f"{base}: FP8 MTP weight_scale must hold one float per tensor or output row")
+ s = scale.float().reshape(-1, 1)
+ return (codes.float() * s).to(torch.bfloat16)
def stacked(proj: str) -> torch.Tensor:
return torch.stack([expert_bf16(f"{name}.experts.{i}.{proj}") for i in range(e)])
@@ -285,7 +292,10 @@ def stacked(proj: str) -> torch.Tensor:
moe4 = nvfp4_moe.moe4_from_bf16(torch.cat([gate, up], dim=1), dn, shared)
elif rd.has(prefix + f"{name}.experts.0.gate_proj.weight"): # the main layers: per-expert FP4
def stack(proj: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- w = torch.stack([raw(f"{name}.experts.{i}.{proj}.weight") for i in range(e)])
+ weights = [raw(f"{name}.experts.{i}.{proj}.weight") for i in range(e)]
+ if any(w.dtype != torch.uint8 for w in weights):
+ raise ValueError(f"{name}.{proj}: expected packed NVFP4 experts; FP8 is read only in MTP")
+ w = torch.stack(weights)
s = torch.stack([raw(f"{name}.experts.{i}.{proj}.weight_scale") for i in range(e)])
s2 = torch.stack([raw(f"{name}.experts.{i}.{proj}.weight_scale_2") for i in range(e)])
return w, s, s2
diff --git a/tests/cuda/nvfp4_tiny.py b/tests/cuda/nvfp4_tiny.py
index 3d6ebd978..5aa5e1c58 100644
--- a/tests/cuda/nvfp4_tiny.py
+++ b/tests/cuda/nvfp4_tiny.py
@@ -54,10 +54,10 @@ def write(dir: Path, *, layers: int = 2, experts: int = 2, vocab: int = 256, hid
moe_width: int = 128, shared_width: int = 64, streams: int = 4, low: int = 64,
ple: bool = True, mtp: bool = True, seed: int = 0, prefix: str = "", ple_bf16: bool = False,
mxfp8: bool = False, ple_nvfp4: bool = False, centred: bool = False, fp8block: bool = False,
- mtp_experts: str = "bf16") -> Path:
- """A tiny ModelOpt checkpoint: ``mxfp8``, ``ple_nvfp4``, ``centred`` norms or ``fp8block`` (block FP8 beside bf16);
- ``mtp_experts`` "bf16" (stacked), "fp8" (per-expert e4m3, a scale a tensor) or "fp8_dequant" (stacked bf16 of
- exactly those e4m3 values times their scales: the same draws)."""
+ mtp_experts: str = "bf16", mtp_scale: str = "tensor") -> Path:
+ """Write ModelOpt weights with bf16 or FP8 MTP experts and tensor, row or block FP8 scales."""
+ if mtp_experts not in ("bf16", "fp8", "fp8_dequant") or mtp_scale not in ("tensor", "row", "block"):
+ raise ValueError("unsupported MTP expert format or scale layout")
dir.mkdir(parents=True, exist_ok=True)
rng = torch.Generator().manual_seed(seed)
@@ -195,11 +195,21 @@ def linear(name: str, n: int, k: int, *, fp4: bool, mx: bool = False, blk: bool
if mtp_experts == "bf16":
add("mtp.layers.0.mlp.experts.gate_up_proj", rand(experts, 2 * moe_width, hidden))
add("mtp.layers.0.mlp.experts.down_proj", rand(experts, hidden, moe_width))
- else: # e4m3 with one fp32 scale a tensor, per expert
+ else: # per-expert e4m3 with fp32 scales
+ def expanded(scale: torch.Tensor) -> torch.Tensor:
+ if mtp_scale == "block":
+ return scale.repeat_interleave(128, 0).repeat_interleave(128, 1)
+ return scale
+
def fp8(n_: int, k_: int) -> tuple[torch.Tensor, torch.Tensor]:
w = rand(n_, k_).float()
- scale = w.abs().max() / 448.0
- return (w / scale).to(torch.float8_e4m3fn), scale.reshape(())
+ if mtp_scale == "block":
+ scale = w.view(n_ // 128, 128, k_ // 128, 128).abs().amax(dim=(1, 3)) / 448.0
+ elif mtp_scale == "row":
+ scale = w.abs().amax(dim=1, keepdim=True) / 448.0
+ else:
+ scale = w.abs().max() / 448.0
+ return (w / expanded(scale)).to(torch.float8_e4m3fn), scale
projs = {p_: [fp8(*shape) for _ in range(experts)] for p_, shape in
(("gate_proj", (moe_width, hidden)), ("up_proj", (moe_width, hidden)),
@@ -208,10 +218,11 @@ def fp8(n_: int, k_: int) -> tuple[torch.Tensor, torch.Tensor]:
for p_, items in projs.items():
for i, (codes, scale) in enumerate(items):
add(f"mtp.layers.0.mlp.experts.{i}.{p_}.weight", codes)
- add(f"mtp.layers.0.mlp.experts.{i}.{p_}.weight_scale", scale)
+ field = "weight_scale_inv" if mtp_scale == "block" else "weight_scale"
+ add(f"mtp.layers.0.mlp.experts.{i}.{p_}.{field}", scale)
else:
def deq(items):
- return torch.stack([(c.float() * s_).to(torch.bfloat16) for c, s_ in items])
+ return torch.stack([(c.float() * expanded(s_)).to(torch.bfloat16) for c, s_ in items])
add("mtp.layers.0.mlp.experts.gate_up_proj", torch.cat([deq(projs["gate_proj"]),
deq(projs["up_proj"])], dim=1))
diff --git a/tests/cuda/test_flashnext_nvfp4_loader.py b/tests/cuda/test_flashnext_nvfp4_loader.py
index 0f5e33b0f..e0652ecf2 100644
--- a/tests/cuda/test_flashnext_nvfp4_loader.py
+++ b/tests/cuda/test_flashnext_nvfp4_loader.py
@@ -199,6 +199,13 @@ def test_the_loader_reads_mxfp8_linears_and_an_nvfp4_table(tmp_path: Path) -> No
from tensorfold.families.qwen4_exp.host_table import NVFP4Table
tiny = write(tmp_path / "mx", mxfp8=True, ple_nvfp4=True, hidden=512) # PLE kernels: 512-wide streams
+ config = json.loads((tiny / "config.json").read_text())
+ config["quantization"]["quantized_layers"] = {
+ "model.layers.1.ple.ple_embedding.ngram_embedding": {"quant_algo": "NVFP4", "group_size": 16}}
+ (tiny / "config.json").write_text(json.dumps(config))
+ from tensorfold.families.qwen4_exp import check
+
+ check(tiny)
if not torch.cuda.is_available(): # the loader builds CUDA tensors
pytest.skip("the loader builds CUDA tensors")
from tensorfold.families.qwen4_exp.cuda.decode import Engine, prefill, serial_decode
@@ -371,19 +378,33 @@ def test_an_nvfp4_checkpoint_refuses_two_ranks(tiny: Path, monkeypatch) -> None:
@pytest.mark.skipif(not torch.cuda.is_available(), reason="the loader builds CUDA tensors")
-def test_fp8_drafter_experts_draft_as_their_dequantized_bf16(tmp_path: Path) -> None:
- """Per-expert FP8 MTP experts (e4m3, a scale a tensor) load: they are dequantized and re-quantized as the bf16
- drafter is, so they draft exactly as stacked bf16 experts holding the same values, and drafts keep serial tokens."""
+@pytest.mark.parametrize("scale", ["tensor", "row", "block"])
+@pytest.mark.parametrize("sampled", [False, True])
+def test_fp8_drafter_experts_draft_as_their_dequantized_bf16(tmp_path: Path, scale: str, sampled: bool,
+ monkeypatch) -> None:
+ """FP8 MTP experts match their dequantized bf16 reference and preserve serial tokens."""
from tensorfold.engine.exact_sampling import Sampling
from tensorfold.families.qwen4_exp.cuda.decode import Engine, mtp_decode, prefill, serial_decode
from tensorfold.families.qwen4_exp.cuda.weights import load
+ from tensorfold.cuda.nvfp4.experts import dense
+
prompt = [5, 17, 99, 250, 7, 64, 30, 11, 12, 13]
- sampling = Sampling(seed=7, top_k=20, top_p=0.95)
- runs = {}
+ sampling = Sampling(seed=7, top_k=20, top_p=0.95) if sampled else Sampling(seed=7, temperature=0, top_k=1)
+ runs, packed, inputs = {}, {}, {}
+ pack = nvfp4_moe.moe4_from_bf16
+
+ def capture(gate_up, down, shared):
+ assert kind not in inputs
+ inputs[kind] = tuple(t.cpu().contiguous().view(torch.uint8) for t in (gate_up, down))
+ return pack(gate_up, down, shared)
+
+ monkeypatch.setattr(nvfp4_moe, "moe4_from_bf16", capture)
for kind in ("fp8", "fp8_dequant"):
- w = load(write(tmp_path / kind, hidden=512, mtp_experts=kind), mtp=True, draft_vocab=128)
+ w = load(write(tmp_path / kind, hidden=512, mtp_experts=kind, mtp_scale=scale), mtp=True, draft_vocab=128)
+ experts = w.mtp.layer.moe.experts.routed_experts
+ packed[kind] = [dense(experts, i, proj).cpu() for i in range(w.cfg.experts) for proj in ("gate", "up", "down")]
e = Engine(w, capacity=256, max_rows=8, prefill_rows=16, graphs=False)
first = prefill(e, prompt, sampling)
ref = serial_decode(e, first, 16, sampling).tokens
@@ -391,5 +412,71 @@ def test_fp8_drafter_experts_draft_as_their_dequantized_bf16(tmp_path: Path) ->
out = mtp_decode(e, first, 16, sampling, depth=4, confidence=0.0)
assert out.tokens == ref, kind
runs[kind] = out
+ assert all(torch.equal(a, b) for a, b in zip(inputs["fp8"], inputs["fp8_dequant"]))
+ assert all(torch.equal(a, b) for a, b in zip(packed["fp8"], packed["fp8_dequant"]))
same = ("tokens", "rounds", "drafted", "accepted", "keeps", "widths") # the same drafts, accepted the same
assert [getattr(runs["fp8"], k) for k in same] == [getattr(runs["fp8_dequant"], k) for k in same]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="the loader builds CUDA tensors")
+@pytest.mark.parametrize("case", ["main_gate", "main_up", "scale_shape", "scale_bytes", "mixed_dtype", "missing_scale"])
+def test_fp8_expert_refusals(tmp_path: Path, monkeypatch, case: str) -> None:
+ """Refuse FP8 target experts and malformed MTP scales before packing or decoding them."""
+
+ from tensorfold.families.qwen4_exp.cuda.reader import _Reader
+ from tensorfold.families.qwen4_exp.cuda.weights import load
+
+ path = write(tmp_path / case, hidden=512, mtp_experts="fp8")
+ get, has = _Reader.get, _Reader.has
+ mtp = "mtp.layers.0.mlp.experts."
+ main = "model.layers.0.mlp.experts."
+ targets = {"main_gate": main + "0.gate_proj.weight", "main_up": main + "1.up_proj.weight",
+ "scale_shape": mtp + "0.gate_proj.weight_scale", "scale_bytes": mtp + "0.gate_proj.weight_scale",
+ "mixed_dtype": mtp + "1.up_proj.weight", "missing_scale": mtp + "0.gate_proj.weight_scale"}
+
+ def altered(self, name):
+ value = get(self, name)
+ if name != targets[case]:
+ return value
+ if case.startswith("main_"):
+ return value.to(torch.float8_e4m3fn)
+ if case == "scale_shape":
+ return torch.ones(512, device=value.device)
+ if case == "scale_bytes":
+ return value.to(torch.uint8)
+ return value.to(torch.float16)
+
+ monkeypatch.setattr(_Reader, "get", altered)
+ if case == "missing_scale":
+ monkeypatch.setattr(_Reader, "has", lambda self, name: False if name == targets[case] else has(self, name))
+ match = {"main_gate": "FP8 routed experts", "main_up": "FP8 is read only in MTP",
+ "scale_shape": "one float per tensor or output row", "scale_bytes": "one float per tensor or output row",
+ "mixed_dtype": "expected FP8 e4m3 MTP expert weights", "missing_scale": "need a tensor, row"}
+ with pytest.raises(ValueError, match=match[case]):
+ load(path, mtp=True, draft_vocab=128)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="the loader builds CUDA tensors")
+@pytest.mark.parametrize("scale", ["tensor", "row", "block"])
+def test_fp8_mtp_real_projection_shapes_dequantize_byte_exactly(tmp_path: Path, monkeypatch, scale: str) -> None:
+ """Capture bf16 inputs before requantization for real 640x2560 gate/up and 2560x640 down projections."""
+
+ from tensorfold.families.qwen4_exp.cuda.weights import load
+
+ actual = write(tmp_path / "fp8", hidden=2560, moe_width=640, mtp_experts="fp8", mtp_scale=scale)
+ reference = write(tmp_path / "bf16", hidden=2560, moe_width=640, mtp_experts="fp8_dequant", mtp_scale=scale)
+ captured = []
+ pack = nvfp4_moe.moe4_from_bf16
+
+ def capture(gate_up, down, shared):
+ assert gate_up.dtype == down.dtype == torch.bfloat16
+ captured.append(tuple(t.cpu().contiguous().view(torch.uint8) for t in (gate_up, down)))
+ return pack(gate_up, down, shared)
+
+ monkeypatch.setattr(nvfp4_moe, "moe4_from_bf16", capture)
+ w = load(actual, mtp=True, draft_vocab=128)
+ assert w.mtp is not None and len(captured) == 1
+ for got, name, shape in zip(captured[0], ("gate_up_proj", "down_proj"), ((2, 1280, 2560), (2, 2560, 640))):
+ want = _tensor(reference, "mtp.layers.0.mlp.experts." + name)
+ assert tuple(want.shape) == shape
+ assert torch.equal(got, want.contiguous().view(torch.uint8)), (scale, name)
diff --git a/tests/test_flashnext_modelopt_check.py b/tests/test_flashnext_modelopt_check.py
new file mode 100644
index 000000000..65e23c46a
--- /dev/null
+++ b/tests/test_flashnext_modelopt_check.py
@@ -0,0 +1,29 @@
+"""Flash Next preflight checks the language weights without claiming support for an unused vision tower."""
+
+import json
+
+import pytest
+
+from tensorfold.families.qwen4_exp import check
+from tensorfold.vision.qwen_checkpoint import PREFIXES
+
+
+def write_config(path, name, algo):
+ config = {"model_type": "qwen4_exp", "quantization_config": {
+ "quant_method": "modelopt", "quant_algo": "MIXED_PRECISION",
+ "quantized_layers": {"model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
+ name: {"quant_algo": algo}},
+ "config_groups": {"experts": {"weights": {"num_bits": 4, "group_size": 16}}}}}
+ (path / "config.json").write_text(json.dumps(config))
+
+
+@pytest.mark.parametrize("prefix", PREFIXES)
+def test_unused_nvfp4_vision_paths_do_not_block_language_preflight(tmp_path, prefix):
+ write_config(tmp_path, prefix + "blocks.0.mlp.linear_fc2", "W4A16_NVFP4")
+ check(tmp_path)
+
+
+def test_nvfp4_language_projection_is_still_refused(tmp_path):
+ write_config(tmp_path, "model.language_model.layers.0.linear_attn.in_proj_qkv", "W4A16_NVFP4")
+ with pytest.raises(ValueError, match="routed experts and n-gram tables only"):
+ check(tmp_path)
diff --git a/tests/test_hub_and_checks.py b/tests/test_hub_and_checks.py
index 83bfa0d65..eb8ff8155 100644
--- a/tests/test_hub_and_checks.py
+++ b/tests/test_hub_and_checks.py
@@ -216,7 +216,13 @@ def test_flash_next_reads_the_nvfp4_checkpoint_and_refuses_other_fp4_blocks(tmp_
gdn_fp4["quantization_config"]["quantized_layers"]["model.language_model.layers.0.linear_attn.in_proj_qkv"] = {
"quant_algo": "W4A16_NVFP4", "group_size": 16}
(tmp_path / "config.json").write_text(json.dumps(gdn_fp4))
- with pytest.raises(ValueError, match="routed experts only.*in_proj_qkv"):
+ with pytest.raises(ValueError, match="routed experts and n-gram tables only.*in_proj_qkv"):
+ qwen4_exp.check(tmp_path)
+ for suffix in ("", ".shard_0"):
+ ple_fp4 = json.loads(json.dumps(mixed))
+ key = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + suffix
+ ple_fp4["quantization_config"]["quantized_layers"][key] = {"quant_algo": "NVFP4", "group_size": 16}
+ (tmp_path / "config.json").write_text(json.dumps(ple_fp4))
qwen4_exp.check(tmp_path)
# FP8 in the MTP drafter's experts is read (dequantized and re-quantized at load); in the main experts it is not
mtp_fp8 = json.loads(json.dumps(mixed))
From 8a836f2d927c19c9296634d454490a045ddf0b06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
<600316+jschmied@users.noreply.github.com>
Date: Thu, 1 Oct 2026 06:42:51 +0200
Subject: [PATCH 093/177] sampling: TENSORFOLD_SEED_SALT moves every
prompt-derived seed
A request without a seed draws one from its prompt, so running an evaluation twice against one server
repeats the samples wherever the conversations agree; an agent benchmark's second pass mostly replays
its first. TENSORFOLD_SEED_SALT (an integer, default 0 = today's seeds) is mixed into every seed drawn
from a prompt, so each run's server can draw independent repeats while staying reproducible per salt.
---
src/tensorfold/engine/exact_sampling.py | 19 ++++++++++++++--
tests/test_seed_salt.py | 29 +++++++++++++++++++++++++
2 files changed, 46 insertions(+), 2 deletions(-)
create mode 100644 tests/test_seed_salt.py
diff --git a/src/tensorfold/engine/exact_sampling.py b/src/tensorfold/engine/exact_sampling.py
index 4c99e6efd..09b1a6ccc 100644
--- a/src/tensorfold/engine/exact_sampling.py
+++ b/src/tensorfold/engine/exact_sampling.py
@@ -5,6 +5,7 @@
from dataclasses import dataclass
import hashlib
import math
+import os
from typing import Any, Sequence
import numpy as np
@@ -32,9 +33,23 @@ def min_log(self) -> float:
return math.log(self.min_p) if self.min_p > 0.0 else -math.inf
-def seed_for(tokens: Sequence[int], salt: int = 0) -> int:
- """A reproducible seed from the prompt: the same conversation samples the same reply."""
+def _salt_from_env() -> int:
+ """``TENSORFOLD_SEED_SALT``: an integer mixed into every prompt-derived seed (0, the default, changes nothing)."""
+ value = os.environ.get("TENSORFOLD_SEED_SALT", "").strip()
+ try:
+ return int(value) if value else 0
+ except ValueError:
+ raise ValueError(f"TENSORFOLD_SEED_SALT={value}: an integer") from None
+
+
+SEED_SALT = _salt_from_env()
+
+
+def seed_for(tokens: Sequence[int], salt: int | None = None) -> int:
+ """A reproducible seed from the prompt: the same conversation samples the same reply (for one salt)."""
+
+ salt = SEED_SALT if salt is None else salt
digest = hashlib.sha256((",".join(str(int(t)) for t in tokens) + f"|{salt}").encode()).digest()
return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1)
diff --git a/tests/test_seed_salt.py b/tests/test_seed_salt.py
new file mode 100644
index 000000000..2567a6e86
--- /dev/null
+++ b/tests/test_seed_salt.py
@@ -0,0 +1,29 @@
+"""TENSORFOLD_SEED_SALT moves every prompt-derived seed; unset or 0 keeps today's seeds."""
+
+import pytest
+
+from tensorfold.engine import exact_sampling
+
+
+def test_the_default_salt_keeps_the_prompt_seed(monkeypatch):
+ monkeypatch.setattr(exact_sampling, "SEED_SALT", 0)
+ assert exact_sampling.seed_for([1, 2, 3]) == exact_sampling.seed_for([1, 2, 3], 0)
+
+
+def test_a_salt_moves_the_prompt_seed_and_stays_reproducible(monkeypatch):
+ base = exact_sampling.seed_for([1, 2, 3], 0)
+ monkeypatch.setattr(exact_sampling, "SEED_SALT", 2)
+ assert exact_sampling.seed_for([1, 2, 3]) != base
+ assert exact_sampling.seed_for([1, 2, 3]) == exact_sampling.seed_for([1, 2, 3], 2)
+
+
+@pytest.mark.parametrize("value,salt", [("", 0), ("0", 0), ("7", 7), (" -3 ", -3)])
+def test_the_salt_reads_the_environment(monkeypatch, value, salt):
+ monkeypatch.setenv("TENSORFOLD_SEED_SALT", value)
+ assert exact_sampling._salt_from_env() == salt
+
+
+def test_a_malformed_salt_is_refused(monkeypatch):
+ monkeypatch.setenv("TENSORFOLD_SEED_SALT", "two")
+ with pytest.raises(ValueError, match="TENSORFOLD_SEED_SALT"):
+ exact_sampling._salt_from_env()
From 303faca549d5ed6e48171cbe4387f8325ba820ff Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:57:16 +0100
Subject: [PATCH 094/177] perf(cuda): the GDN prompt chain loads each stage one
stage ahead (same bits)
---
src/tensorfold/cuda/kernels/gdn_prefill.cu | 86 +++++++++++++++++-----
1 file changed, 66 insertions(+), 20 deletions(-)
diff --git a/src/tensorfold/cuda/kernels/gdn_prefill.cu b/src/tensorfold/cuda/kernels/gdn_prefill.cu
index de354c7c9..d75c4ea01 100644
--- a/src/tensorfold/cuda/kernels/gdn_prefill.cu
+++ b/src/tensorfold/cuda/kernels/gdn_prefill.cu
@@ -4,20 +4,74 @@
#include
#include
#include
+#include
namespace {
constexpr int DK = 128;
constexpr int DV = 128;
-__device__ __forceinline__ float4 load4f(const __nv_bfloat16* p) {
- const uint2 w = *reinterpret_cast(p);
+__device__ __forceinline__ float4 widen(uint2 w) {
const __nv_bfloat162 a = *reinterpret_cast(&w.x);
const __nv_bfloat162 b = *reinterpret_cast(&w.y);
return make_float4(__low2float(a), __high2float(a), __low2float(b), __high2float(b));
}
-__device__ __forceinline__ float4 load4f(const float* p) { return *reinterpret_cast(p); }
+__device__ __forceinline__ float4 widen(float4 w) { return w; }
+
+// A stage's keys, queries, gates and value rows, loaded into registers a stage ahead (in flight while the current
+// stage computes) and then stored as chain_kernel stages them: keys and queries as fp32 float4s, values bf16.
+template
+struct Prefetch {
+ using Raw = typename std::conditional::value, float4, uint2>::type;
+ static constexpr int KQ = STEPS * (DK / 4) / THREADS, VV = STEPS * ROWS / 8 / THREADS;
+ static_assert(KQ * THREADS == STEPS * (DK / 4) && VV * THREADS * 8 == STEPS * ROWS, "whole slots a thread");
+ Raw kr[KQ], qr[KQ];
+ uint4 vr[VV];
+ float gr, br;
+
+ __device__ __forceinline__ void fetch(const QK* q, const QK* k, const __nv_bfloat16* v, const float* g,
+ const float* beta, int t0, int W, int hk, int hv, int head, int key_head,
+ int col0) {
+#pragma unroll
+ for (int it = 0; it < KQ; ++it) {
+ const int i = threadIdx.x + it * THREADS, st = i / (DK / 4), c = i % (DK / 4);
+ const size_t at = (static_cast(min(t0 + st, W - 1)) * hk + key_head) * DK + 4 * c;
+ kr[it] = *reinterpret_cast(k + at);
+ qr[it] = *reinterpret_cast(q + at);
+ }
+#pragma unroll
+ for (int it = 0; it < VV; ++it) {
+ const int i = threadIdx.x + it * THREADS, st = i / (ROWS / 8), c = i % (ROWS / 8);
+ vr[it] = *reinterpret_cast(v + (static_cast(min(t0 + st, W - 1)) * hv + head) * DV
+ + col0 + 8 * c);
+ }
+ if (threadIdx.x < STEPS) {
+ const size_t at = static_cast(min(t0 + static_cast(threadIdx.x), W - 1)) * hv + head;
+ gr = g[at];
+ br = beta[at];
+ }
+ }
+
+ __device__ __forceinline__ void store(float4 (*ks)[DK / 4], float4 (*qs)[DK / 4], float* gs, float* bs,
+ __nv_bfloat16 (*vs)[ROWS]) const {
+#pragma unroll
+ for (int it = 0; it < KQ; ++it) {
+ const int i = threadIdx.x + it * THREADS;
+ ks[i / (DK / 4)][i % (DK / 4)] = widen(kr[it]);
+ qs[i / (DK / 4)][i % (DK / 4)] = widen(qr[it]);
+ }
+#pragma unroll
+ for (int it = 0; it < VV; ++it) {
+ const int i = threadIdx.x + it * THREADS;
+ *reinterpret_cast(&vs[i / (ROWS / 8)][8 * (i % (ROWS / 8))]) = vr[it];
+ }
+ if (threadIdx.x < STEPS) {
+ gs[threadIdx.x] = gr;
+ bs[threadIdx.x] = br;
+ }
+ }
+};
template
__global__ void __launch_bounds__(2 * ROWS) chain_kernel(
@@ -26,7 +80,7 @@ __global__ void __launch_bounds__(2 * ROWS) chain_kernel(
float* __restrict__ last, __nv_bfloat16* __restrict__ y, int W, int hk, int hv) {
__shared__ float4 ks[STEPS][DK / 4], qs[STEPS][DK / 4];
__shared__ float gs[STEPS], bs[STEPS];
- __shared__ __nv_bfloat16 vs[STEPS][ROWS];
+ __shared__ __align__(16) __nv_bfloat16 vs[STEPS][ROWS];
const int head = blockIdx.x, row = blockIdx.y * ROWS + (threadIdx.x >> 1), half = threadIdx.x & 1;
const int key_head = head / (hv / hk);
float s[64];
@@ -36,24 +90,14 @@ __global__ void __launch_bounds__(2 * ROWS) chain_kernel(
const float4 t = *reinterpret_cast(s0 + 8 * j + 4 * half);
s[4 * j] = t.x; s[4 * j + 1] = t.y; s[4 * j + 2] = t.z; s[4 * j + 3] = t.w;
}
+ Prefetch pf; // each stage's loads in flight during the one before
+ pf.fetch(q, k, v, g, beta, 0, W, hk, hv, head, key_head, blockIdx.y * ROWS);
for (int t0 = 0; t0 < W; t0 += STEPS) {
const int n = min(STEPS, W - t0);
__syncthreads();
- for (int i = threadIdx.x; i < n * (DK / 4); i += 2 * ROWS) {
- const int st = i / (DK / 4), c = i % (DK / 4);
- const size_t at = (static_cast(t0 + st) * hk + key_head) * DK + 4 * c;
- ks[st][c] = load4f(k + at);
- qs[st][c] = load4f(q + at);
- }
- for (int i = threadIdx.x; i < n; i += 2 * ROWS) {
- gs[i] = g[static_cast(t0 + i) * hv + head];
- bs[i] = beta[static_cast(t0 + i) * hv + head];
- }
- for (int i = threadIdx.x; i < n * ROWS; i += 2 * ROWS) {
- const int st = i / ROWS, rr = i % ROWS;
- vs[st][rr] = v[(static_cast(t0 + st) * hv + head) * DV + blockIdx.y * ROWS + rr];
- }
+ pf.store(ks, qs, gs, bs, vs);
__syncthreads();
+ if (t0 + STEPS < W) pf.fetch(q, k, v, g, beta, t0 + STEPS, W, hk, hv, head, key_head, blockIdx.y * ROWS);
for (int tt = 0; tt < n; ++tt) {
const float gt = gs[tt], bt = bs[tt];
const size_t vat = (static_cast(t0 + tt) * hv + head) * DV + row;
@@ -95,7 +139,8 @@ __global__ void __launch_bounds__(2 * ROWS) chain_kernel(
float* s1 = last + (static_cast(head) * DV + row) * DK;
#pragma unroll
for (int j = 0; j < 16; ++j)
- *reinterpret_cast(s1 + 8 * j + 4 * half) = make_float4(s[4 * j], s[4 * j + 1], s[4 * j + 2], s[4 * j + 3]);
+ *reinterpret_cast(s1 + 8 * j + 4 * half) = make_float4(s[4 * j], s[4 * j + 1], s[4 * j + 2],
+ s[4 * j + 3]);
}
template
@@ -118,7 +163,8 @@ void gdn_prefill_cuda(const at::Tensor& q, const at::Tensor& k, const at::Tensor
const at::Tensor& beta, const at::Tensor& state, at::Tensor& last, at::Tensor& y, int sms) {
const bool wide = v.size(1) >= sms;
if (q.scalar_type() == at::kFloat) {
- if (wide) launch(q, k, v, g, beta, state, last, y); else launch(q, k, v, g, beta, state, last, y);
+ if (wide) launch(q, k, v, g, beta, state, last, y);
+ else launch(q, k, v, g, beta, state, last, y);
} else {
if (wide) launch<__nv_bfloat16, 128>(q, k, v, g, beta, state, last, y);
else launch<__nv_bfloat16, 64>(q, k, v, g, beta, state, last, y);
From 49b2cb8947d716421cf21aa3e6a4f1ba1d9cb653 Mon Sep 17 00:00:00 2001
From: Isham Rashik <20547074+di37@users.noreply.github.com>
Date: Thu, 1 Oct 2026 17:44:37 +0400
Subject: [PATCH 095/177] fix(mlx): mlx-lm 0.32 support: read cache state the
same under 0.31 and 0.32
mlx-lm 0.32.0 returns cache state with offsets and nested lists (KVCache.state is
(keys, values, offset) over the whole buffer; ArraysCache.state is
(cache, left_padding, lengths)). Under 0.32:
- Qwen dense's load-time stream check (row_forward.check_streams) crashed
- cache_arrays dropped every recurrent layer's arrays from prefill evals
- AlternatingKVCache's state setter raised on the 3-tuple
cache_arrays now walks nested lists; cache_contents gives a layer's rows up to
its offset under either version; AlternatingKVCache defers to KVCache's setter.
Pin widened to mlx-lm <0.33.
Tested on an M4 Max: non-CUDA suite 3592 passed on mlx-lm 0.32.0 (+ mlx-vlm
0.7.4) and 3481 on 0.31.3, the one failure on both being the existing
test_deepseek_v4_family prefill/decode argmax case. Qwen3.8-27B-MLX-4bit with
DFlash2 on both versions: drafted replies equal "draft": false, resends, a
5,058-token cached prefix and concurrent runs (same token_sha).
---
pyproject.toml | 2 +-
src/tensorfold/engine/alternating_kv.py | 7 ++-
src/tensorfold/engine/family_common.py | 29 +++++++----
.../kernels/qwen/dense/v1/row_forward.py | 7 +--
tests/test_alternating_kv.py | 7 +--
tests/test_cache_contents.py | 49 +++++++++++++++++++
tests/test_lane_engine.py | 14 +++---
tests/test_nemotron_pass.py | 5 +-
tests/test_row_forward.py | 17 ++++---
9 files changed, 97 insertions(+), 40 deletions(-)
create mode 100644 tests/test_cache_contents.py
diff --git a/pyproject.toml b/pyproject.toml
index ae9dd3ba3..28be4c1ea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,7 @@ classifiers = [
dependencies = [
# the prompt kernels and SSD streaming are tested on MLX 0.32.2 and 0.32.3 (#88)
"mlx>=0.32.2,<0.32.4; sys_platform == 'darwin'",
- "mlx-lm>=0.31.3,<0.32; sys_platform == 'darwin'",
+ "mlx-lm>=0.31.3,<0.33; sys_platform == 'darwin'",
"numpy>=1.26",
"huggingface-hub>=0.24",
"tokenizers>=0.20; sys_platform != 'darwin'",
diff --git a/src/tensorfold/engine/alternating_kv.py b/src/tensorfold/engine/alternating_kv.py
index c85be9d64..48910c0ad 100644
--- a/src/tensorfold/engine/alternating_kv.py
+++ b/src/tensorfold/engine/alternating_kv.py
@@ -69,14 +69,13 @@ def trim(self, n: int) -> int:
return n
@property
- def state(self) -> tuple[mx.array, mx.array]:
+ def state(self) -> tuple:
return super().state
@state.setter
- def state(self, v: tuple[mx.array, mx.array]) -> None:
+ def state(self, v: tuple) -> None:
self.drop_spare()
- self.keys, self.values = v
- self.offset = self.keys.shape[2]
+ KVCache.state.fset(self, v) # (keys, values) under mlx-lm 0.31, (keys, values, offset) under 0.32
def drop_spare(self) -> None:
"""Forget the spare (a retained or stored cache keeps one buffer)."""
diff --git a/src/tensorfold/engine/family_common.py b/src/tensorfold/engine/family_common.py
index 4ea56b1da..3d403ff40 100644
--- a/src/tensorfold/engine/family_common.py
+++ b/src/tensorfold/engine/family_common.py
@@ -19,16 +19,25 @@ def drop_spares(cache: list[Any]) -> list[Any]:
return drop(cache)
+def _arrays_in(value: Any) -> list[Any]:
+ """The arrays in ``value``, through nested lists and tuples (mlx-lm 0.32 states add offsets and nested lists)."""
+
+ if isinstance(value, (list, tuple)):
+ return [a for v in value for a in _arrays_in(v)]
+ return [value] if value is not None and hasattr(value, "shape") else []
+
+
def cache_arrays(cache: list[Any]) -> list[Any]:
"""Every array a cache list holds (a KV cache nothing was written to yet has none)."""
- arrays: list[Any] = []
- for item in cache:
- if getattr(item, "keys", 0) is None:
- continue
- state = item.state
- if isinstance(state, (list, tuple)):
- arrays.extend(a for a in state if a is not None and hasattr(a, "shape"))
- elif state is not None and hasattr(state, "shape"):
- arrays.append(state)
- return arrays
+ return [a for item in cache if getattr(item, "keys", 0) is not None for a in _arrays_in(item.state)]
+
+
+def cache_contents(item: Any) -> list[Any]:
+ """One layer's cached values, the same under mlx-lm 0.31 and 0.32."""
+
+ # KV rows stop at the offset. Under 0.32, state is the whole buffer plus that offset.
+ if getattr(item, "keys", 0) is None:
+ return []
+ rows = getattr(item, "keys_and_values", None)
+ return _arrays_in(rows() if callable(rows) else item.state)
diff --git a/src/tensorfold/kernels/qwen/dense/v1/row_forward.py b/src/tensorfold/kernels/qwen/dense/v1/row_forward.py
index ed11d2b9a..9e3c3908c 100644
--- a/src/tensorfold/kernels/qwen/dense/v1/row_forward.py
+++ b/src/tensorfold/kernels/qwen/dense/v1/row_forward.py
@@ -7,6 +7,7 @@
import mlx.core as mx
+from tensorfold.engine.family_common import cache_contents
from tensorfold.kernels.inputs import ints
from tensorfold.kernels.qwen.dense.v1 import row_matmul
from tensorfold.kernels.qwen.dense.v1.row_glue import _chain, add_norm, gated_delta, gdn_post, gdn_pre, mlp_act
@@ -375,11 +376,7 @@ def check_streams(core: Any, head: Any, make_cache: Callable[[], list[Any]], cop
"""Return equality and failures for batched versus standalone logits and partial-window cache commits across prompt lengths and window widths."""
def arrays(cache: list[Any]) -> list[mx.array]:
- out = []
- for item in cache:
- state = item.state
- out.extend(a for a in (state if isinstance(state, (list, tuple)) else [state]) if a is not None)
- return out
+ return [a for item in cache for a in cache_contents(item)]
vocab = int(core.embed_tokens["weight"].shape[0]) # MLX's gather reads past the table for larger ids, unchecked
bases = []
diff --git a/tests/test_alternating_kv.py b/tests/test_alternating_kv.py
index 1ae91714d..cbdc602c2 100644
--- a/tests/test_alternating_kv.py
+++ b/tests/test_alternating_kv.py
@@ -6,6 +6,7 @@
from mlx_lm.models.cache import KVCache
from tensorfold.engine.alternating_kv import AlternatingKVCache, drop_spares
+from tensorfold.engine.family_common import cache_contents
from tensorfold.engine.lane_engine import LaneEngine
@@ -29,9 +30,9 @@ def test_matches_kv_cache_under_random_writes_trims_and_copies():
if rng.random() < 0.1:
copies.append((LaneEngine.copy_single_cache([ref])[0], LaneEngine.copy_single_cache([alt])[0]))
assert ref.offset == alt.offset
- assert _same(ref.state, alt.state)
+ assert _same(cache_contents(ref), cache_contents(alt))
for r, a in copies:
- assert _same(r.state, a.state)
+ assert _same(cache_contents(r), cache_contents(a))
def test_decode_writes_alternate_and_spares_drop():
@@ -44,5 +45,5 @@ def test_decode_writes_alternate_and_spares_drop():
assert alt.nbytes > KVCache.nbytes.fget(alt)
drop_spares([alt])
assert alt.spare_keys is None and alt.recent_keys is None
- keys, _ = alt.state
+ keys, _ = cache_contents(alt)
assert keys.shape[2] == 43 and float(keys[0, 0, 42, 0].item()) == 2.0
diff --git a/tests/test_cache_contents.py b/tests/test_cache_contents.py
new file mode 100644
index 000000000..2143ae8cf
--- /dev/null
+++ b/tests/test_cache_contents.py
@@ -0,0 +1,49 @@
+"""cache_arrays and cache_contents read mlx-lm caches the same under 0.31 and 0.32."""
+
+# 0.32 state adds offsets and nests a recurrent layer's arrays in a list.
+
+import pytest
+
+mx = pytest.importorskip("mlx.core")
+pytest.importorskip("mlx_lm")
+
+from mlx_lm.models.cache import ArraysCache, KVCache # noqa: E402
+
+from tensorfold.engine.alternating_kv import AlternatingKVCache # noqa: E402
+from tensorfold.engine.family_common import cache_arrays, cache_contents # noqa: E402
+
+
+def _layers():
+ kv = KVCache()
+ rows = mx.arange(2 * 5 * 4, dtype=mx.float32).reshape(1, 2, 5, 4)
+ kv.update_and_fetch(rows, rows + 1)
+ recurrent = ArraysCache(size=2)
+ recurrent[0], recurrent[1] = mx.zeros((1, 3, 6)), mx.ones((1, 2, 4, 4))
+ return kv, recurrent, rows
+
+
+def test_cache_arrays_holds_every_layers_arrays():
+ kv, recurrent, _ = _layers()
+ held = cache_arrays([kv, recurrent, KVCache()]) # 0.31 gives the KV rows as views, 0.32 the buffers
+ assert len(held) == 4 and all(hasattr(a, "shape") for a in held)
+ assert {id(recurrent[0]), id(recurrent[1])} <= {id(a) for a in held}
+ assert all(a.shape[2] >= kv.offset for a in held[:2])
+
+
+def test_cache_contents_are_the_rows_up_to_the_offset():
+ kv, recurrent, rows = _layers()
+ keys, values = cache_contents(kv)
+ assert keys.shape[2] == kv.offset == 5 # 0.32's state is the whole step-sized buffer
+ assert mx.array_equal(keys, rows).item() and mx.array_equal(values, rows + 1).item()
+ assert [a.shape for a in cache_contents(recurrent)] == [(1, 3, 6), (1, 2, 4, 4)]
+ assert cache_contents(KVCache()) == []
+
+
+def test_alternating_state_round_trips_through_its_setter():
+ kv, _, rows = _layers()
+ alt = AlternatingKVCache()
+ alt.update_and_fetch(rows, rows + 1)
+ alt.update_and_fetch(rows[..., :1, :], rows[..., :1, :]) # a decode write leaves a spare
+ alt.state = kv.state
+ assert alt.spare_keys is None and alt.offset == kv.offset
+ assert all(mx.array_equal(a, b).item() for a, b in zip(cache_contents(alt), cache_contents(kv)))
diff --git a/tests/test_lane_engine.py b/tests/test_lane_engine.py
index 127e8a80e..17f96af7c 100644
--- a/tests/test_lane_engine.py
+++ b/tests/test_lane_engine.py
@@ -83,14 +83,14 @@ def test_copy_single_cache_detaches_kv_and_recurrent_arrays() -> None:
keys = mx.ones((1, 1, 3, 2))
kv.update_and_fetch(keys, keys)
arrays = ArraysCache(size=2)
- arrays.state = [mx.zeros((1, 2)), mx.ones((1, 2))]
+ arrays.cache = [mx.zeros((1, 2)), mx.ones((1, 2))]
clone = LaneEngine.copy_single_cache([kv, arrays])
assert clone[0] is not kv and clone[0].offset == 3
clone[0].keys[..., 0, :] = 9.0
- clone[1].state[0][0, 0] = 5.0
- mx.eval(clone[0].keys, clone[1].state[0], kv.keys, arrays.state[0])
+ clone[1].cache[0][0, 0] = 5.0
+ mx.eval(clone[0].keys, clone[1].cache[0], kv.keys, arrays.cache[0])
assert kv.keys[0, 0, 0, 0].item() == 1.0
- assert arrays.state[0][0, 0].item() == 0.0
+ assert arrays.cache[0][0, 0].item() == 0.0
def test_copy_single_cache_copies_a_view_out_of_its_base() -> None:
@@ -101,10 +101,10 @@ def test_copy_single_cache_copies_a_view_out_of_its_base() -> None:
base = mx.random.normal((64, 256, 256))
mx.eval(base)
arrays = ArraysCache(size=1)
- arrays.state = [base[5:6]]
+ arrays.cache = [base[5:6]]
clone = LaneEngine.copy_single_cache([arrays])
- mx.eval(clone[0].state[0])
- assert mx.array_equal(clone[0].state[0], base[5:6]).item()
+ mx.eval(clone[0].cache[0])
+ assert mx.array_equal(clone[0].cache[0], base[5:6]).item()
before = mx.get_active_memory()
del arrays, base
gc.collect()
diff --git a/tests/test_nemotron_pass.py b/tests/test_nemotron_pass.py
index 3e2d059d9..c068613e3 100644
--- a/tests/test_nemotron_pass.py
+++ b/tests/test_nemotron_pass.py
@@ -7,6 +7,7 @@
nn = pytest.importorskip("mlx.nn")
pytest.importorskip("mlx_lm")
+from tensorfold.engine.family_common import cache_contents # noqa: E402
from tensorfold.families.nemotron_h import prompt_pass # noqa: E402
@@ -47,8 +48,8 @@ def test_a_pass_gives_each_chunk_its_own_forwards_bits(sizes):
both = prompt_pass.hidden(backbone, ids, passed, tuple(sizes))
mx.eval(solo, both)
assert bool(mx.array_equal(solo, both).item())
- a = [x for c in one for x in c.state if x is not None]
- b = [x for c in passed for x in c.state if x is not None]
+ a = [x for c in one for x in cache_contents(c)]
+ b = [x for c in passed for x in cache_contents(c)]
assert len(a) == len(b) and all(bool(mx.array_equal(x, y).item()) for x, y in zip(a, b))
finally:
mx.set_default_device(previous)
diff --git a/tests/test_row_forward.py b/tests/test_row_forward.py
index 741530292..1d4aa6f15 100644
--- a/tests/test_row_forward.py
+++ b/tests/test_row_forward.py
@@ -6,6 +6,7 @@
mx = pytest.importorskip("mlx.core")
nn = pytest.importorskip("mlx.nn")
+from tensorfold.engine.family_common import cache_contents # noqa: E402
from tensorfold.engine.lane_engine import LaneEngine # noqa: E402
from tensorfold.kernels.qwen.dense.v1 import ( # noqa: E402
exact_attention, lane_glue, lane_tree, row_forward, row_glue, row_matmul)
@@ -51,7 +52,7 @@ def _run(model, tokens, cache, start, keep=None):
logits, record = row_forward.forward(core, head, tokens, parents, cache, start, pipeline_layers=2)
keep = len(tokens) if keep is None else keep
row_forward.commit(cache, record, list(range(keep)), len(tokens), start)
- mx.eval(logits, *[a for c in cache for a in c.state if a is not None])
+ mx.eval(logits, *[a for c in cache for a in cache_contents(c)])
return logits
@@ -132,7 +133,7 @@ def test_prefill_chunking_does_not_change_bits(tiny):
_run(model, prompt[begin:begin + size], split, begin)
begin += size
for a, b in zip(whole, split):
- for x, y in zip(a.state, b.state):
+ for x, y in zip(cache_contents(a), cache_contents(b)):
if x is not None:
assert _same(x, y)
@@ -213,7 +214,7 @@ def test_aligned_prefill_resumes_exactly(tiny, monkeypatch):
prompt = [int(t) for t in mx.random.randint(0, 512, (29,)).tolist()]
def arrays(cache):
- return [a for c in cache for a in c.state if a is not None]
+ return [a for c in cache for a in cache_contents(c)]
fresh = engine.prefill_prefix(prompt)
first = LaneStream("first", prompt[:19], max_new_tokens=1)
@@ -259,7 +260,7 @@ def test_tree_window_nodes_equal_serial_paths(tiny, monkeypatch):
for i, r in enumerate(path):
_run(model, [tokens[r]], serial_cache, start + i)
for a, b in zip(cache, serial_cache):
- for x, y in zip(a.state, b.state):
+ for x, y in zip(cache_contents(a), cache_contents(b)):
if x is not None:
assert _same(x, y)
@@ -332,7 +333,7 @@ def test_hidden_rows_and_keep_rows(tiny):
assert _same(ref, lg)
for a, b in zip(caches, mine):
for ia, ib in zip(a, b):
- for u, v in zip(ia.state, ib.state):
+ for u, v in zip(cache_contents(ia), cache_contents(ib)):
if u is not None:
assert _same(u, v)
@@ -365,7 +366,7 @@ def test_streams_with_trees_equal_each_alone(tiny, monkeypatch):
for (ref, own), cache, a, window in zip(alone, caches, offsets, windows):
assert _same(lg[0, a:a + len(window)], ref[0])
for ia, ib in zip(cache, own):
- for u, v in zip(ia.state, ib.state):
+ for u, v in zip(cache_contents(ia), cache_contents(ib)):
if u is not None:
assert _same(u, v)
@@ -403,7 +404,7 @@ def test_one_kernel_signature_for_every_window(monkeypatch):
parents = [-1, 0, 0, 1, 2, 2, 4, 3]
logits, record = row_forward.forward(core, head, tokens[start:start + 8], parents, cache, start)
row_forward.commit(cache, record, [0, 2, 4, 6], len(parents), start)
- mx.eval(logits, *[a for c in cache for a in c.state if a is not None])
+ mx.eval(logits, *[a for c in cache for a in cache_contents(c)])
assert bool(mx.all(mx.isfinite(logits)).item())
monkeypatch.setattr(row_forward, "ROW_ATTENTION", False)
streams = [_prefill(model, tokens[100 + 20 * s:100 + 20 * s + 9 + 4 * s]) for s in range(3)]
@@ -422,7 +423,7 @@ def test_one_kernel_signature_for_every_window(monkeypatch):
keep = max(1, w // 2)
row_forward.commit(streams[s], record, list(range(keep)), w, starts[s])
starts[s] += keep
- mx.eval(logits, *[a for s in ids for c in streams[s] for a in c.state if a is not None])
+ mx.eval(logits, *[a for s in ids for c in streams[s] for a in cache_contents(c)])
assert bool(mx.all(mx.isfinite(logits)).item())
finally:
for c, old in zip(caches, saved):
From b8027439a3affb73afaf217383e837c461183adf Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 1 Oct 2026 19:45:43 +0100
Subject: [PATCH 096/177] 0.6.1 engine work: NVFP4 checkpoints in their own
math, prompts that fill together and wider lane blocks on CUDA, Flash Next on
Macs at long context, and an experimental native Windows host layer
- NVFP4: lane and prompt kernels in the checkpoint's FP4/FP8 math on SM 12.x and W4A16 on SM 8.9-10.x, with --precision full for bf16 activations (quality and speed in docs/recipes/cuda.md); fused SwiGLU into the prompt MLP and one launch for the GDN gates; bulk copies kept inside the stack limit on every toolkit
- CUDA 27B: queued prompts share one prefill forward, each stream's convolution, DeltaNet chain and attention on its own state (same bits); wider lane blocks for NVFP4 from 96 SMs and for MLX 4-bit on SM 12.0 (same bits); GB10 and unmeasured GPUs plan as 0.6.0; engines close their scheduler's worker
- Flash Next on Macs: block scores read each pooled block once a threadgroup, contexts grow in place, a window's tokens stay on the GPU, chained MTP rows sit beside the head's cache, and the sparse decode kernels build at load
- Windows: one GPU a process, memory sized by Windows' API, pinned buffered reads, Ctrl+Break stack dumps; not yet run on Windows hardware
- Docs and packaging: both RoCE devices on two DGX Sparks; wheels ship the NVFP4 headers
---
pyproject.toml | 2 +-
src/tensorfold/cli.py | 17 +-
src/tensorfold/cli_args.py | 7 +
src/tensorfold/cuda/build.py | 25 +-
src/tensorfold/cuda/capacity.py | 9 +-
src/tensorfold/cuda/comm.py | 3 +
src/tensorfold/cuda/direct_read.py | 11 +-
src/tensorfold/cuda/kernels/qmm.py | 14 +-
src/tensorfold/cuda/kernels/qmm_group.cu | 275 +++++++------
src/tensorfold/cuda/kernels/qmm_tiles.py | 11 +
src/tensorfold/cuda/nvfp4/act.cu | 112 +++++
src/tensorfold/cuda/nvfp4/checkpoint.cpp | 171 ++++++++
src/tensorfold/cuda/nvfp4/checkpoint.py | 220 ++++++++++
src/tensorfold/cuda/nvfp4/gemm_ck.cu | 268 ++++++++++++
src/tensorfold/cuda/nvfp4/gemm_ws.cu | 263 ++++++++++++
src/tensorfold/cuda/nvfp4/lane4.cu | 382 ++++++++++++++++++
src/tensorfold/cuda/nvfp4/linear.py | 61 ++-
src/tensorfold/cuda/nvfp4/mma4.cuh | 41 ++
src/tensorfold/cuda/nvfp4/nvfp4q.cuh | 38 ++
src/tensorfold/cuda/nvfp4/qmmf.cu | 25 +-
src/tensorfold/cuda/nvfp4/swiglu4.cuh | 100 +++++
src/tensorfold/cuda/precision.py | 48 +++
src/tensorfold/cuda/scheduler.py | 17 +-
src/tensorfold/engine/lane_engine.py | 16 +-
src/tensorfold/families/qwen3_5/cuda/b16.cpp | 6 +
src/tensorfold/families/qwen3_5/cuda/b16.cu | 262 ++++++++++--
src/tensorfold/families/qwen3_5/cuda/b16.py | 28 +-
.../families/qwen3_5/cuda/engine.py | 7 +
.../families/qwen3_5/cuda/forward.py | 16 +-
src/tensorfold/families/qwen3_5/cuda/multi.py | 110 ++++-
.../families/qwen3_5/cuda/nvfp4_load.py | 44 +-
.../families/qwen3_5/cuda/prefill.py | 216 +++++++++-
.../families/qwen3_5/cuda/weights.py | 11 +-
.../families/qwen3_5_moe/cuda/engine.py | 7 +
.../families/qwen4_exp/cuda/engine.py | 7 +
src/tensorfold/families/qwen4_exp/decode.py | 19 +-
.../families/qwen4_exp/host_table.py | 24 ++
.../families/qwen4_exp/mtp_cache.py | 51 +++
src/tensorfold/families/qwen4_exp/runtime.py | 47 ++-
.../kernels/qwen/flash_next/v1/attention.py | 193 ++++++---
.../kernels/qwen/flash_next/v1/base.py | 17 +-
src/tensorfold/server/cancellation.py | 2 +-
src/tensorfold/server/memory_budget.py | 18 +
src/tensorfold/server/scheduler.py | 6 +-
src/tensorfold/server/stacks.py | 21 +-
tests/cuda/conftest.py | 17 +
tests/cuda/test_nvfp4_checkpoint.py | 294 ++++++++++++++
tests/cuda/test_nvfp4_lane_tiles.py | 82 ++++
tests/cuda/test_nvfp4_linear.py | 26 ++
tests/cuda/test_qmm_group.py | 19 +-
tests/cuda/test_qwen27_b16.py | 56 +++
tests/cuda/test_qwen27_batch_admit.py | 98 +++++
tests/cuda/test_qwen27_checkpoint_slots.py | 1 +
tests/cuda/test_qwen27_nvfp4.py | 27 +-
tests/cuda/test_qwen27_prefill_batch.py | 129 ++++++
tests/test_cuda_27b_ignore_eos.py | 1 +
tests/test_cuda_batch_admit.py | 159 ++++++++
tests/test_cuda_failed_admission.py | 6 +
tests/test_cuda_gpu_floor.py | 15 +-
tests/test_cuda_precision.py | 99 +++++
tests/test_cuda_scheduler_close.py | 35 ++
tests/test_flash_next_chain_side.py | 88 ++++
tests/test_flash_next_index_scores.py | 109 +++++
tests/test_lane_stream_context.py | 41 ++
tests/test_nvfp4_bulk_tile.py | 14 +
tests/test_nvfp4_lane_tile.py | 45 +++
tests/test_qmm_group_tile.py | 16 +
tests/test_qwen27_stream_depth.py | 19 +-
tests/test_windows_platform.py | 144 +++++++
69 files changed, 4471 insertions(+), 317 deletions(-)
create mode 100644 src/tensorfold/cuda/kernels/qmm_tiles.py
create mode 100644 src/tensorfold/cuda/nvfp4/act.cu
create mode 100644 src/tensorfold/cuda/nvfp4/checkpoint.cpp
create mode 100644 src/tensorfold/cuda/nvfp4/checkpoint.py
create mode 100644 src/tensorfold/cuda/nvfp4/gemm_ck.cu
create mode 100644 src/tensorfold/cuda/nvfp4/gemm_ws.cu
create mode 100644 src/tensorfold/cuda/nvfp4/lane4.cu
create mode 100644 src/tensorfold/cuda/nvfp4/mma4.cuh
create mode 100644 src/tensorfold/cuda/nvfp4/nvfp4q.cuh
create mode 100644 src/tensorfold/cuda/nvfp4/swiglu4.cuh
create mode 100644 src/tensorfold/cuda/precision.py
create mode 100644 src/tensorfold/families/qwen4_exp/mtp_cache.py
create mode 100644 tests/cuda/test_nvfp4_checkpoint.py
create mode 100644 tests/cuda/test_nvfp4_lane_tiles.py
create mode 100644 tests/cuda/test_qwen27_b16.py
create mode 100644 tests/cuda/test_qwen27_batch_admit.py
create mode 100644 tests/cuda/test_qwen27_prefill_batch.py
create mode 100644 tests/test_cuda_batch_admit.py
create mode 100644 tests/test_cuda_precision.py
create mode 100644 tests/test_cuda_scheduler_close.py
create mode 100644 tests/test_flash_next_chain_side.py
create mode 100644 tests/test_flash_next_index_scores.py
create mode 100644 tests/test_lane_stream_context.py
create mode 100644 tests/test_nvfp4_bulk_tile.py
create mode 100644 tests/test_nvfp4_lane_tile.py
create mode 100644 tests/test_qmm_group_tile.py
create mode 100644 tests/test_windows_platform.py
diff --git a/pyproject.toml b/pyproject.toml
index 28be4c1ea..787296ed9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ where = ["src"]
"tensorfold.families.glm5_next.cuda" = ["*.cu", "*.cpp"]
"tensorfold.families.nemotron_h.cuda" = ["*.cu", "*.cpp"]
"tensorfold.cuda.exl3" = ["*.cu", "*.cuh", "*.cpp"]
-"tensorfold.cuda.nvfp4" = ["*.cu", "*.cpp"]
+"tensorfold.cuda.nvfp4" = ["*.cu", "*.cuh", "*.cpp"]
"tensorfold.streaming.hostsync" = ["*.cpp", "CMakeLists.txt"]
[tool.pytest.ini_options]
diff --git a/src/tensorfold/cli.py b/src/tensorfold/cli.py
index edf81d6f6..aa23915bb 100644
--- a/src/tensorfold/cli.py
+++ b/src/tensorfold/cli.py
@@ -263,12 +263,18 @@ def _serve_cuda(args: argparse.Namespace, family: Any, model_dir: Path, context:
served = args.name or (args.model.rstrip("/").split("/")[-1] if hub.is_repo_id(args.model) else model_dir.name)
where = f", rank {args.rank} of 2" if args.tp == 2 else ""
print(f"[tensorfold] loading {served}: {family.title} ({family.model_type}) on CUDA{where}", flush=True)
- from tensorfold.cuda import prompt_precision
+ from tensorfold.cuda import precision, prompt_precision
asked = getattr(args, "prefill_fp8", None)
prompt_precision.set_fp8(prompt_precision.FP8_BY_DEFAULT if asked is None else asked) # before any weight loads
+ chosen = getattr(args, "precision", None)
+ precision.set_mode(chosen or precision.CHECKPOINT, asked=chosen is not None)
engine = family.package.cuda_engine(model_dir, **options)
- fp8 = prompt_precision.fp8() and bool(getattr(getattr(engine, "w", None), "fast_prefill", False))
+ weights = getattr(engine, "w", None)
+ fp8 = prompt_precision.fp8() and bool(getattr(weights, "fast_prefill", False))
+ if asked and not fp8 and getattr(weights, "precision", "full") == precision.CHECKPOINT:
+ raise ValueError("--prefill-fp8 is for --precision full: the checkpoint's own math already runs its prompts in "
+ "FP4 and FP8")
if asked and not fp8:
raise ValueError("--prefill-fp8: this checkpoint's prompt matmuls have no FP8 kernel (EXL3 packs, MLX formats "
"other than Qwen's 4-bit g64, Flash Next without MXFP8 layers); drop the flag")
@@ -293,9 +299,10 @@ def _serve_cuda(args: argparse.Namespace, family: Any, model_dir: Path, context:
shown = "greedy" if float(sampling.get("temperature", 1.0)) <= 0 else ", ".join(
f"{k} {v}" for k, v in sampling.items())
effective_context = app.effective_context_window
+ own = getattr(weights, "precision", "") == precision.CHECKPOINT
+ prompts = "FP8 activations" if fp8 else "the checkpoint math" if own else "bf16 activations"
print(f"[tensorfold] serving {served} at http://{args.host}:{args.port}/v1 on CUDA{where} "
- f"(sampling: {shown}; drafts: {'off' if args.no_drafts else 'on'}; "
- f"prompts: {'FP8 activations' if fp8 else 'bf16 activations'}; "
+ f"(sampling: {shown}; drafts: {'off' if args.no_drafts else 'on'}; prompts: {prompts}; "
f"context: {'unlimited' if effective_context is None else effective_context}; "
f"loaded in {time.perf_counter() - started:.1f}s)", flush=True)
serve(app, args.host, int(args.port))
@@ -510,7 +517,7 @@ def _terminate(signum: int, frame: Any) -> None:
signal.signal(signal.SIGTERM, _terminate)
from tensorfold.server import live
- line = live.start(app)
+ line = live.start(app) # connections and decode/prefill tok/s on one line, in a terminal only
try:
server.serve_forever()
except KeyboardInterrupt:
diff --git a/src/tensorfold/cli_args.py b/src/tensorfold/cli_args.py
index 5069da2e8..76260fed2 100644
--- a/src/tensorfold/cli_args.py
+++ b/src/tensorfold/cli_args.py
@@ -131,6 +131,13 @@ def build_parser(handlers: dict[str, Callable[[argparse.Namespace], int]]) -> ar
"docs/recipes/cuda.md#prompt-precision has the measured cost). Default: "
f"{'FP8' if FP8_BY_DEFAULT else 'bf16'} activations. Replies equal this server's own serial "
"decoding either way")
+ cuda.add_argument("--precision", choices=("checkpoint", "full"), default=argparse.SUPPRESS,
+ help="the math for checkpoints that name their activations' formats (NVFP4): checkpoint, the "
+ "default, runs their own math as their runtimes do (FP4 x FP4 in NVFP4 layers on SM 12.x "
+ "GPUs, FP8 x FP8 in FP8 layers from SM 8.9, under the checkpoint's static input scales; "
+ "layers a GPU has no mma for run W4A16, and the startup line says which); full runs bf16 "
+ "activations against the stored weights exactly. The weights never change, only the math; "
+ "MLX checkpoints have one math. Replies equal this server's own serial decoding either way")
serve.set_defaults(func=handlers["serve"])
pull = commands.add_parser("pull", help="download models (or draft models) from Hugging Face")
diff --git a/src/tensorfold/cuda/build.py b/src/tensorfold/cuda/build.py
index 2258f1ed2..63cc5aa41 100644
--- a/src/tensorfold/cuda/build.py
+++ b/src/tensorfold/cuda/build.py
@@ -17,8 +17,8 @@
LOCK_WAIT_SECONDS = 60.0 # a start still waiting on the same lock this long says so again
-def arch_flags(need: tuple[int, int] = MIN_CAPABILITY) -> list[str]:
- """nvcc flags for the current GPU alone; a GPU older than ``need`` is refused by name."""
+def arch_flags(need: tuple[int, int] = MIN_CAPABILITY, arch_specific: bool = False) -> list[str]:
+ """nvcc flags for this GPU alone (``arch_specific``: its ``a`` target); a GPU under ``need`` is refused by name."""
import torch
@@ -28,7 +28,8 @@ def arch_flags(need: tuple[int, int] = MIN_CAPABILITY) -> list[str]:
raise RuntimeError(f"TensorFold's CUDA kernels need compute capability {need[0]}.{need[1]} or newer ({why}"
f"{' for these weights' if need > MIN_CAPABILITY else ''}); this GPU "
f"({torch.cuda.get_device_name()}) is {major}.{minor}")
- return [f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"]
+ a = "a" if arch_specific else ""
+ return [f"-gencode=arch=compute_{major}{minor}{a},code=sm_{major}{minor}{a}"]
def refuse_old_gpu(need: tuple[int, int] = MIN_CAPABILITY) -> None:
@@ -45,12 +46,13 @@ def refuse_old_gpu(need: tuple[int, int] = MIN_CAPABILITY) -> None:
raise ValueError(str(exc)) from None
-def load(name: str, sources: str | list[str], need: tuple[int, int] = MIN_CAPABILITY, **kwargs: Any) -> Any:
+def load(name: str, sources: str | list[str], need: tuple[int, int] = MIN_CAPABILITY, *, arch_specific: bool = False,
+ **kwargs: Any) -> Any:
"""torch's JIT ``load`` for this GPU only (NVIDIA's containers list every architecture back to sm_80), with a line when it compiles or waits on a lock."""
from torch.utils import cpp_extension
- kwargs["extra_cuda_cflags"] = [*kwargs.get("extra_cuda_cflags", []), *arch_flags(need)]
+ kwargs["extra_cuda_cflags"] = [*kwargs.get("extra_cuda_cflags", []), *arch_flags(need, arch_specific)]
links = _toolkit()
if links:
kwargs["extra_ldflags"] = [*kwargs.get("extra_ldflags", []), *links]
@@ -119,7 +121,8 @@ def _toolkit() -> list[str]:
"""The pip route's tools: the venv's ninja on PATH, NVIDIA's pip toolkit, and the link flags it needs."""
venv = Path(sys.executable).parent # an inactive venv: ninja is beside its python, not on PATH
- if shutil.which("ninja") is None and (venv / "ninja").is_file():
+ tool = "ninja.exe" if os.name == "nt" else "ninja" # the venv's ninja, named as Windows names it
+ if shutil.which("ninja") is None and (venv / tool).is_file():
os.environ["PATH"] = f"{venv}{os.pathsep}{os.environ.get('PATH', '')}"
return list(_pip_flags())
@@ -140,6 +143,16 @@ def pip_toolkit(cpp_extension: Any, torch: Any) -> list[str]:
if cpp_extension.CUDA_HOME is not None or not getattr(torch.version, "cuda", None):
return []
home = Path(torch.__file__).resolve().parents[1] / "nvidia" / f"cu{torch.version.cuda.split('.')[0]}"
+ if os.name == "nt": # a Windows wheel: DLLs beside torch and in NVIDIA's cu13 folder
+ search = [folder for folder in (home / "bin", home / "lib", Path(torch.__file__).resolve().parent / "lib")
+ if folder.is_dir()]
+ for folder in search: # a built .pyd loads cudart64_13.dll by name at import
+ try:
+ os.add_dll_directory(str(folder))
+ except (AttributeError, OSError):
+ pass # old CPython or a refused folder: PATH is the fallback
+ os.environ["PATH"] = f"{folder}{os.pathsep}{os.environ.get('PATH', '')}"
+ return [f"/LIBPATH:{folder}" for folder in search if (folder / "cudart.lib").is_file()]
if not (home / "bin" / "nvcc").is_file():
return []
cpp_extension.CUDA_HOME = os.environ["CUDA_HOME"] = str(home) # torch reads its CUDA_HOME at every build
diff --git a/src/tensorfold/cuda/capacity.py b/src/tensorfold/cuda/capacity.py
index cd810d564..333e01f8f 100644
--- a/src/tensorfold/cuda/capacity.py
+++ b/src/tensorfold/cuda/capacity.py
@@ -266,16 +266,11 @@ def choose(plan: Plan, peers: list[list[int]] | None = None) -> int:
def floor(model_dir: str | Path) -> tuple[int, int]:
- """The compute capability a checkpoint's kernels need: NVFP4 and FP8 (ModelOpt, compressed-tensors) use clusters."""
+ """The compute capability a checkpoint's kernels need: 8.9 for every format (clusters are taken where present)."""
from tensorfold.cuda import build
- from tensorfold.cuda.nvfp4.format import is_quantized
- try:
- quantized = is_quantized(model_dir)
- except (OSError, ValueError): # an unreadable config is named by the estimate below
- quantized = False
- return build.CLUSTERS if quantized else build.MIN_CAPABILITY
+ return build.MIN_CAPABILITY
def admit(model_dir: str | Path, requested: int | None, explicit: bool | None, torch,
diff --git a/src/tensorfold/cuda/comm.py b/src/tensorfold/cuda/comm.py
index 2e50cdc90..88b94d948 100644
--- a/src/tensorfold/cuda/comm.py
+++ b/src/tensorfold/cuda/comm.py
@@ -17,6 +17,9 @@ class _UniqueId(ctypes.Structure):
def _library() -> ctypes.CDLL:
+ if os.name == "nt":
+ raise RuntimeError("TensorFold does not run tensor-parallel (NCCL) on Windows: CUDA on Windows has no "
+ "libnccl to wrap; use one GPU per process there")
candidates = [os.environ.get("TF_NCCL_LIB", "")]
found = ctypes.util.find_library("nccl")
if found:
diff --git a/src/tensorfold/cuda/direct_read.py b/src/tensorfold/cuda/direct_read.py
index 616ac1f92..b745725ab 100644
--- a/src/tensorfold/cuda/direct_read.py
+++ b/src/tensorfold/cuda/direct_read.py
@@ -28,6 +28,7 @@ class Reader:
def __init__(self) -> None:
self.direct = hasattr(os, "O_DIRECT")
+ self.staged = os.name == "nt" # Windows has no O_DIRECT: reads there buffer into pinned staging
self.staging: list[list] = [] # [pinned piece, event of its last copy]
self.turn = 0
@@ -43,6 +44,10 @@ def read(self, path: str | Path, offset: int, n: int, device: str | torch.device
if exc.errno != errno.EINVAL:
raise
self.direct = False # the file system refuses O_DIRECT, on the open or on a read
+ if n > 0 and self.staged and (cuda or pinned):
+ # Windows stages here: one page-locked block where that means host RAM, filled by one buffered read
+ raw = self._buffered(path, offset, n, pinned=torch.cuda.is_available())
+ return raw.to(device) if cuda else raw
raw = self._buffered(path, offset, n)
return raw.to(device) if cuda else raw
@@ -56,8 +61,10 @@ def close(self) -> None:
self.staging.clear()
getattr(torch._C, "_host_emptyCache", lambda: None)()
- def _buffered(self, path, offset: int, n: int) -> torch.Tensor:
- raw = torch.empty((n,), dtype=torch.uint8)
+ def _buffered(self, path, offset: int, n: int, pinned: bool = False) -> torch.Tensor:
+ """One buffered read of ``n`` bytes, page-locked when ``pinned`` (how Windows stages, having no O_DIRECT)."""
+
+ raw = torch.empty((n,), dtype=torch.uint8, pin_memory=pinned, device="cpu")
view = memoryview(raw.numpy())
with open(path, "rb", buffering=0) as f:
f.seek(offset)
diff --git a/src/tensorfold/cuda/kernels/qmm.py b/src/tensorfold/cuda/kernels/qmm.py
index 0a95bd3ac..f91fd0b0c 100644
--- a/src/tensorfold/cuda/kernels/qmm.py
+++ b/src/tensorfold/cuda/kernels/qmm.py
@@ -10,18 +10,26 @@
import triton
import triton.language as tl
+from .qmm_tiles import group_tile
+
@lru_cache(maxsize=1)
def _ext():
from tensorfold.cuda.build import load
here = Path(__file__).parent
- return load(name="tensorfold_qmm_v4", sources=[str(here / "qmm.cpp"), str(here / "qmm.cu"),
+ return load(name="tensorfold_qmm_v5", sources=[str(here / "qmm.cpp"), str(here / "qmm.cu"),
str(here / "qmm_group.cu"), str(here / "qmm_prefill.cu"),
str(here / "qmm_prefill8.cu")],
extra_cuda_cflags=["-O3"], verbose=False)
+@lru_cache(maxsize=None)
+def _chip(device: int) -> tuple[int, int, int]:
+ p = torch.cuda.get_device_properties(device)
+ return p.major, p.minor, p.multi_processor_count
+
+
@lru_cache(maxsize=None)
def grouped(device: int) -> bool:
"""sm_12x runs groups of 64 through the grouped kernel: several projections of one input in a launch."""
@@ -161,7 +169,8 @@ def matmul(x: torch.Tensor, q: Q4, xs: torch.Tensor | None = None, *, sk: int |
if out is None:
out = torch.empty((m, q.n), dtype=torch.float32 if f32 else torch.bfloat16, device=x.device)
if q.gs == 64 and reduce and grouped(x.device.index):
- _ext().qmm_group(x, xs, [q.weight], [q.scales], [q.biases], [out], [q.n], [sk], f32, 0, -1)
+ _ext().qmm_group(x, xs, [q.weight], [q.scales], [q.biases], [out], [q.n], [sk], f32,
+ group_tile(m, *_chip(x.device.index)), -1)
return out
if sk > 1 and not reduce and part is None:
part = torch.empty((sk, m, q.n), dtype=torch.float32, device=x.device)
@@ -184,6 +193,7 @@ def matmul_group(x: torch.Tensor, qs: list[Q4], xs: torch.Tensor | None = None,
xs = group_sums(x, 64)
dtype = torch.float32 if f32 else torch.bfloat16
outs = [torch.empty((x.shape[0], q.n), dtype=dtype, device=x.device) for q in qs]
+ tile = tile or group_tile(x.shape[0], *_chip(x.device.index))
_ext().qmm_group(x, xs, [q.weight for q in qs], [q.scales for q in qs], [q.biases for q in qs], outs,
[q.n for q in qs], sks, f32, tile, early)
return outs
diff --git a/src/tensorfold/cuda/kernels/qmm_group.cu b/src/tensorfold/cuda/kernels/qmm_group.cu
index 0196f4e12..e33dbccf3 100644
--- a/src/tensorfold/cuda/kernels/qmm_group.cu
+++ b/src/tensorfold/cuda/kernels/qmm_group.cu
@@ -6,6 +6,7 @@
#include
#include
#include
+#include
#include
#include "qmm_frag.cuh"
@@ -42,9 +43,15 @@ __device__ __forceinline__ uint32_t pairm(uint32_t w, int s, uint32_t mask) {
return r;
}
+template
+using ic = std::integral_constant;
+
// Clusters of C blocks along x, one part each: a cluster covers C / sk column tiles, each split in sk K slices.
// SWAP (8-row tiles): weights are the MMA's A operand (16 columns) and the rows its B (8), half the MMAs of 16 rows.
-template
+// SKIP: a warp's m16 tiles wholly past M skip their fragments, mmas and scaling (a part-filled row tile). SPREAD:
+// every K slice sums its share of the tile's outputs (else slice 0 sums them all); the same adds in the same order.
+template
__global__ void __launch_bounds__(WM * WN * 32) group_kernel(
const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const __grid_constant__ Parts parts,
int M, int K, int ldx, int rows_t, int C) {
@@ -150,101 +157,141 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel(
commit();
}
grid_launch();
- for (int it = 0; it < per; ++it) {
- wait();
- __syncthreads();
- const int next = it + STAGES - 1;
- if (next < per) {
- load_x(next % STAGES, g0 + next);
- load_w(next % STAGES, g0 + next);
- }
- commit();
- const unsigned char* p = stage(it % STAGES);
- const uint32_t* pw = reinterpret_cast(p + T::X);
- const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W);
- const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S);
- uint32_t words[T::NT][GS / 32];
-#pragma unroll
- for (int j = 0; j < T::NT; ++j)
-#pragma unroll
- for (int v = 0; v < GS / 32; ++v)
- words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v];
- float d[I][J][4];
- if constexpr (SWAP) {
-#pragma unroll
- for (int kt = 0; kt < GS / 16; ++kt) {
- uint32_t bx[2]; // rows 0-7 at k lo, k hi: the B fragment
- ldmatrix2(bx, p + (lane & 7) * T::ROW + swz(lane & 7, kt * 2 + ((lane >> 3) & 1)) * 16);
-#pragma unroll
- for (int i = 0; i < I; ++i) { // A rows g, g + 8 are columns of n8 tiles 2i, 2i + 1
- const int s0 = (kt & 1) * 8;
- const uint32_t a[4] = {pairm(words[2 * i][kt / 2], s0, mask),
- pairm(words[2 * i + 1][kt / 2], s0, mask),
- pairm(words[2 * i][kt / 2], s0 + 4, mask),
- pairm(words[2 * i + 1][kt / 2], s0 + 4, mask)};
- if (kt == 0) mma0(d[i][0], a, bx[0], bx[1]);
- else mma(d[i][0], a, bx[0], bx[1]);
- }
+ auto steps = [&](auto live_c) { // LIVE: the warp's m16 tiles holding rows below M
+ [[maybe_unused]] constexpr int LIVE = decltype(live_c)::value;
+ for (int it = 0; it < per; ++it) {
+ wait();
+ __syncthreads();
+ const int next = it + STAGES - 1;
+ if (next < per) {
+ load_x(next % STAGES, g0 + next);
+ load_w(next % STAGES, g0 + next);
}
-#pragma unroll
- for (int i = 0; i < I; ++i)
-#pragma unroll
- for (int e = 0; e < 4; ++e) { // d: column g (+ 8 for e >= 2), rows 2t, 2t + 1
- const int col = wn * (BN / WN) + i * 16 + (lane >> 2) + (e >> 1) * 8;
- const float sv = __bfloat162float(ps[col]), bv = __bfloat162float(ps[BN + col]);
- const float xv = px[(lane & 3) * 2 + (e & 1)];
- acc[i][0][e] = __fmaf_rn(xv, bv, __fmaf_rn(d[i][0][e], sv, acc[i][0][e]));
+ commit();
+ const unsigned char* p = stage(it % STAGES);
+ const uint32_t* pw = reinterpret_cast(p + T::X);
+ const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W);
+ const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S);
+ uint32_t words[T::NT][GS / 32];
+#pragma unroll
+ for (int j = 0; j < T::NT; ++j)
+#pragma unroll
+ for (int v = 0; v < GS / 32; ++v)
+ words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v];
+ float d[I][J][4];
+ if constexpr (SWAP) {
+#pragma unroll
+ for (int kt = 0; kt < GS / 16; ++kt) {
+ uint32_t bx[2]; // rows 0-7 at k lo, k hi: the B fragment
+ const int ch = kt * 2 + ((lane >> 3) & 1);
+ ldmatrix2(bx, p + (lane & 7) * T::ROW + swz(lane & 7, ch) * 16);
+#pragma unroll
+ for (int i = 0; i < I; ++i) { // A rows g, g + 8 are columns of n8 tiles 2i, 2i + 1
+ const int s0 = (kt & 1) * 8;
+ const uint32_t a[4] = {pairm(words[2 * i][kt / 2], s0, mask),
+ pairm(words[2 * i + 1][kt / 2], s0, mask),
+ pairm(words[2 * i][kt / 2], s0 + 4, mask),
+ pairm(words[2 * i + 1][kt / 2], s0 + 4, mask)};
+ if (kt == 0) mma0(d[i][0], a, bx[0], bx[1]);
+ else mma(d[i][0], a, bx[0], bx[1]);
+ }
}
- } else {
#pragma unroll
- for (int kt = 0; kt < GS / 16; ++kt) {
- uint32_t a[T::MT][4];
+ for (int i = 0; i < I; ++i)
#pragma unroll
- for (int i = 0; i < T::MT; ++i) {
- const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8;
- const int ch = kt * 2 + (lane >> 4);
- ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16);
- }
+ for (int e = 0; e < 4; ++e) { // d: column g (+ 8 for e >= 2), rows 2t, 2t + 1
+ const int col = wn * (BN / WN) + i * 16 + (lane >> 2) + (e >> 1) * 8;
+ const float sv = __bfloat162float(ps[col]), bv = __bfloat162float(ps[BN + col]);
+ const float xv = px[(lane & 3) * 2 + (e & 1)];
+ acc[i][0][e] = __fmaf_rn(xv, bv, __fmaf_rn(d[i][0][e], sv, acc[i][0][e]));
+ }
+ } else {
#pragma unroll
- for (int j = 0; j < T::NT; ++j) {
- const uint32_t b0 = pairm(words[j][kt / 2], (kt & 1) * 8, mask);
- const uint32_t b1 = pairm(words[j][kt / 2], (kt & 1) * 8 + 4, mask);
+ for (int kt = 0; kt < GS / 16; ++kt) {
+ uint32_t a[LIVE > 0 ? LIVE : 1][4];
#pragma unroll
- for (int i = 0; i < T::MT; ++i) {
- if (kt == 0) mma0(d[i][j], a[i], b0, b1);
- else mma(d[i][j], a[i], b0, b1);
+ for (int i = 0; i < LIVE; ++i) {
+ const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8;
+ const int ch = kt * 2 + (lane >> 4);
+ ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16);
}
- }
- }
#pragma unroll
- for (int j = 0; j < T::NT; ++j) {
- const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2;
- const __nv_bfloat162 s2 = *reinterpret_cast(ps + col);
- const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col);
- const float sv[2] = {__low2float(s2), __high2float(s2)};
- const float bv[2] = {__low2float(b2), __high2float(b2)};
+ for (int j = 0; j < T::NT; ++j) {
+ const uint32_t b0 = pairm(words[j][kt / 2], (kt & 1) * 8, mask);
+ const uint32_t b1 = pairm(words[j][kt / 2], (kt & 1) * 8 + 4, mask);
#pragma unroll
- for (int i = 0; i < T::MT; ++i) {
- const int row = wm * (BM / WM) + i * 16 + (lane >> 2);
- const float xv[2] = {px[row], px[row + 8]};
+ for (int i = 0; i < LIVE; ++i) {
+ if (kt == 0) mma0(d[i][j], a[i], b0, b1);
+ else mma(d[i][j], a[i], b0, b1);
+ }
+ }
+ }
#pragma unroll
- for (int e = 0; e < 4; ++e) // acc = fma(xs, b, fma(p, s, acc)): qmm.cu's order
- acc[i][j][e] = __fmaf_rn(xv[e >> 1], bv[e & 1],
- __fmaf_rn(d[i][j][e], sv[e & 1], acc[i][j][e]));
+ for (int j = 0; j < T::NT; ++j) {
+ const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2;
+ const __nv_bfloat162 s2 = *reinterpret_cast(ps + col);
+ const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col);
+ const float sv[2] = {__low2float(s2), __high2float(s2)};
+ const float bv[2] = {__low2float(b2), __high2float(b2)};
+#pragma unroll
+ for (int i = 0; i < LIVE; ++i) {
+ const int row = wm * (BM / WM) + i * 16 + (lane >> 2);
+ const float xv[2] = {px[row], px[row + 8]};
+#pragma unroll
+ for (int e = 0; e < 4; ++e) // acc = fma(xs, b, fma(p, s, acc)): qmm.cu's order
+ acc[i][j][e] = __fmaf_rn(xv[e >> 1], bv[e & 1],
+ __fmaf_rn(d[i][j][e], sv[e & 1], acc[i][j][e]));
+ }
}
}
}
+ };
+ if constexpr (!SKIP || SWAP || T::MT == 1) {
+ steps(ic());
+ } else {
+ const int live = min(T::MT, max(0, (M - m0 - wm * (BM / WM) + 15) / 16));
+ if (live == 0) steps(ic<0>());
+ else if (live == 1) steps(ic<1>());
+ else if (T::MT > 2 && live == 2) steps(ic<(T::MT > 2 ? 2 : T::MT)>());
+ else if (T::MT > 3 && live == 3) steps(ic<(T::MT > 3 ? 3 : T::MT)>());
+ else if (T::MT > 4 && live <= 4) steps(ic<(T::MT > 4 ? 4 : T::MT)>());
+ else if (T::MT > 6 && live <= 6) steps(ic<(T::MT > 6 ? 6 : T::MT)>());
+ else steps(ic());
}
wait<0>();
__syncthreads();
} else {
grid_launch();
}
+ const int N = P.n;
+ auto put = [&](int row, int col, float v) {
+ if (row >= M || col >= N) return;
+ if (F32) reinterpret_cast(P.out)[static_cast(row) * N + col] = v;
+ else reinterpret_cast<__nv_bfloat16*>(P.out)[static_cast(row) * N + col] = __float2bfloat16_rn(v);
+ };
+ auto emit = [&](int i, int j, int h, float v0, float v1) { // entries 2h, 2h + 1 of fragment (i, j)
+ if constexpr (SWAP) { // column g (+ 8 for h 1), rows 2t, 2t + 1
+ const int col = n0 + wn * (BN / WN) + i * 16 + (lane >> 2) + h * 8;
+ put(m0 + (lane & 3) * 2, col, v0);
+ put(m0 + (lane & 3) * 2 + 1, col, v1);
+ } else {
+ const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2;
+ const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8;
+ if (row >= M) return;
+ if (!F32 && col + 1 < N && (N & 1) == 0) {
+ auto* dst = reinterpret_cast<__nv_bfloat16*>(P.out) + static_cast(row) * N + col;
+ *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1);
+ } else {
+ put(row, col, v0);
+ put(row, col + 1, v1);
+ }
+ }
+ };
if (sk > 1) { // uniform in a cluster: it holds one part
#if __CUDA_ARCH__ >= 900
auto cluster = cooperative_groups::this_cluster();
float* mine = reinterpret_cast(buf);
- if (slice != 0) {
+ if (SPREAD || slice != 0) {
#pragma unroll
for (int i = 0; i < I; ++i)
#pragma unroll
@@ -253,62 +300,54 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel(
for (int e = 0; e < 4; ++e) mine[((i * J + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e];
}
cluster.sync();
- if (slice == 0 && live) { // slice 0 adds its tile's peers in slice order, as qmm.cu does
- for (int peer = 1; peer < sk; ++peer) {
- const float* theirs = cluster.map_shared_rank(mine, rank + peer);
+ if constexpr (SPREAD) { // slice s adds pairs s, s + sk, .. over the tile's slices in order
+ if (live) {
+ for (int pr = slice; pr < I * J * 2; pr += sk) {
+ const float* first = cluster.map_shared_rank(mine, rank - slice);
+ float v0 = first[2 * pr * T::THREADS + tid], v1 = first[(2 * pr + 1) * T::THREADS + tid];
+#pragma unroll 7
+ for (int peer = 1; peer < sk; ++peer) {
+ const float* theirs = cluster.map_shared_rank(mine, rank - slice + peer);
+ v0 = v0 + theirs[2 * pr * T::THREADS + tid];
+ v1 = v1 + theirs[(2 * pr + 1) * T::THREADS + tid];
+ }
+ emit(pr / (J * 2), pr / 2 % J, pr % 2, v0, v1);
+ }
+ }
+ cluster.sync(); // peers keep their memory until every slice has read it
+ return;
+ } else {
+ if (slice == 0 && live) { // slice 0 adds its tile's peers in slice order, as qmm.cu does
+ for (int peer = 1; peer < sk; ++peer) {
+ const float* theirs = cluster.map_shared_rank(mine, rank + peer);
#pragma unroll
- for (int i = 0; i < I; ++i)
+ for (int i = 0; i < I; ++i)
#pragma unroll
- for (int j = 0; j < J; ++j)
+ for (int j = 0; j < J; ++j)
#pragma unroll
- for (int e = 0; e < 4; ++e)
- acc[i][j][e] = acc[i][j][e] + theirs[((i * J + j) * 4 + e) * T::THREADS + tid];
+ for (int e = 0; e < 4; ++e)
+ acc[i][j][e] = acc[i][j][e] + theirs[((i * J + j) * 4 + e) * T::THREADS + tid];
+ }
}
+ cluster.sync(); // peers keep their memory until slice 0 has read it
+ if (slice != 0) return;
}
- cluster.sync(); // peers keep their memory until slice 0 has read it
- if (slice != 0) return;
#else
__trap();
+ return;
#endif
}
if (!live) return;
- const int N = P.n;
- auto put = [&](int row, int col, float v) {
- if (row >= M || col >= N) return;
- if (F32) reinterpret_cast(P.out)[static_cast(row) * N + col] = v;
- else reinterpret_cast<__nv_bfloat16*>(P.out)[static_cast(row) * N + col] = __float2bfloat16_rn(v);
- };
- if constexpr (SWAP) {
-#pragma unroll
- for (int i = 0; i < I; ++i)
-#pragma unroll
- for (int e = 0; e < 4; ++e)
- put(m0 + (lane & 3) * 2 + (e & 1), n0 + wn * (BN / WN) + i * 16 + (lane >> 2) + (e >> 1) * 8,
- acc[i][0][e]);
- return;
- }
#pragma unroll
for (int i = 0; i < I; ++i)
#pragma unroll
- for (int j = 0; j < J; ++j) {
- const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2;
+ for (int j = 0; j < J; ++j)
#pragma unroll
- for (int h = 0; h < 2; ++h) {
- const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8;
- if (row >= M) continue;
- const float v0 = acc[i][j][2 * h], v1 = acc[i][j][2 * h + 1];
- if (!F32 && col + 1 < N && (N & 1) == 0) {
- auto* dst = reinterpret_cast<__nv_bfloat16*>(P.out) + static_cast(row) * N + col;
- *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1);
- } else {
- put(row, col, v0);
- put(row, col + 1, v1);
- }
- }
- }
+ for (int h = 0; h < 2; ++h) emit(i, j, h, acc[i][j][2 * h], acc[i][j][2 * h + 1]);
}
-template
+template
void launch(const at::Tensor& x, const at::Tensor& xs, Parts& parts, int C, bool pdl) {
using T = LaneTile;
const int M = x.size(0), K = x.size(1), rows_t = (M + BM - 1) / BM;
@@ -319,7 +358,7 @@ void launch(const at::Tensor& x, const at::Tensor& xs, Parts& parts, int C, bool
P.first = clusters;
clusters += rows_t * ((P.tiles + C / P.sk - 1) / (C / P.sk));
}
- auto kernel = group_kernel;
+ auto kernel = group_kernel;
static bool configured = false;
if (!configured) {
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, T::SMEM);
@@ -352,10 +391,12 @@ void launch(const at::Tensor& x, const at::Tensor& xs, Parts& parts, int C, bool
}
// Tiles never change bits; 0 picks by rows and chip: the PRO 6000, power-capped at one row, takes fewer MMAs a weight.
+// 10-12 (``qmm.group_tile`` on SM 12.0 from 96 SMs) spread the slice sums and skip m16 tiles wholly past M.
template
void dispatch(int tile, int M, bool gb10, const at::Tensor& x, const at::Tensor& xs, Parts& parts, int C, bool pdl) {
if (tile == 0 && gb10) tile = M <= 16 ? 2 : M <= 32 ? 3 : M <= 64 ? 4 : 5;
if (tile == 0) tile = M <= 8 ? 7 : M <= 16 ? 8 : M <= 32 ? 9 : M <= 64 ? 4 : 5;
+ auto go = [&](auto full, auto skip, int bm) { M % bm ? skip() : full(); };
switch (tile) {
case 1: launch<64, 16, 64, 1, 4, 4, F32>(x, xs, parts, C, pdl); break;
case 2: launch<64, 16, 64, 1, 4, 8, F32>(x, xs, parts, C, pdl); break;
@@ -366,6 +407,12 @@ void dispatch(int tile, int M, bool gb10, const at::Tensor& x, const at::Tensor&
case 7: launch<64, 8, 128, 1, 4, 4, F32, true>(x, xs, parts, C, pdl); break;
case 8: launch<64, 16, 128, 1, 8, 4, F32>(x, xs, parts, C, pdl); break;
case 9: launch<64, 32, 128, 1, 8, 4, F32>(x, xs, parts, C, pdl); break;
+ case 10: go([&] { launch<64, 128, 128, 2, 4, 2, F32, false, false, true>(x, xs, parts, C, pdl); },
+ [&] { launch<64, 128, 128, 2, 4, 2, F32, false, true, true>(x, xs, parts, C, pdl); }, 128); break;
+ case 11: go([&] { launch<64, 64, 128, 1, 8, 3, F32, false, false, true>(x, xs, parts, C, pdl); },
+ [&] { launch<64, 64, 128, 1, 8, 3, F32, false, true, true>(x, xs, parts, C, pdl); }, 64); break;
+ case 12: go([&] { launch<64, 64, 128, 2, 4, 3, F32, false, false, true>(x, xs, parts, C, pdl); },
+ [&] { launch<64, 64, 128, 2, 4, 3, F32, false, true, true>(x, xs, parts, C, pdl); }, 64); break;
default: TORCH_CHECK(false, "unknown group tile ", tile);
}
}
diff --git a/src/tensorfold/cuda/kernels/qmm_tiles.py b/src/tensorfold/cuda/kernels/qmm_tiles.py
new file mode 100644
index 000000000..79877fe4f
--- /dev/null
+++ b/src/tensorfold/cuda/kernels/qmm_tiles.py
@@ -0,0 +1,11 @@
+"""The grouped 4-bit lane matmul's block by rows and chip (host only; blocks never change bits)."""
+
+WIDE_SMS = 96 # SM 12.0 GPUs from this many SMs (the RTX PRO 6000's 188) take the wide blocks
+
+
+def group_tile(m: int, major: int, minor: int, sms: int) -> int:
+ """0 (the kernel's own pick by rows and chip) or a wide block for ``m`` rows on SM major.minor with ``sms`` SMs."""
+
+ if (major, minor) != (12, 0) or sms < WIDE_SMS or m <= 16:
+ return 0
+ return 12 if m <= 32 else 11 if m <= 64 else 10 if m <= 96 else 11
diff --git a/src/tensorfold/cuda/nvfp4/act.cu b/src/tensorfold/cuda/nvfp4/act.cu
new file mode 100644
index 000000000..cc13bb327
--- /dev/null
+++ b/src/tensorfold/cuda/nvfp4/act.cu
@@ -0,0 +1,112 @@
+// Activations in the checkpoint's own formats, each row alone (a row's codes never depend on another row):
+// NVFP4 (e2m1 codes, an e4m3 scale a 16 inputs under the checkpoint's static global input scale) and FP8 (e4m3 under
+// the static input scale, in the weights' fragment order).
+
+#include
+#include
+#include
+#include
+#include
+
+#include "mma4.cuh"
+#include "nvfp4q.cuh"
+
+namespace {
+
+__device__ __forceinline__ void unpack8(const uint4 u, float (&f)[8]) {
+ const __nv_bfloat162* p = reinterpret_cast(&u);
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ const float2 v = __bfloat1622float2(p[i]);
+ f[2 * i] = v.x;
+ f[2 * i + 1] = v.y;
+ }
+}
+
+// One thread a 16-input block: amax, the block scale e4m3(g * amax / 6), codes e2m1(x * g / scale).
+// codes [M, K/2] (input 2j in byte j's low nibble) and scales [K/64, mpad, 4] (byte b: the step's block b), or with
+// ``tb`` rows a tile the prompt GEMM's tiles, codes [mpad / tb][K/64][tb][32] in its smem swizzle and scales
+// [mpad / tb][K/64][tb][4], rows past M zero.
+__global__ void __launch_bounds__(128) quant4_kernel(const __nv_bfloat16* __restrict__ x, int ldx, int M, int K,
+ float g, uint8_t* __restrict__ codes, uint8_t* __restrict__ scales,
+ int mpad, int tb) {
+ const int row = blockIdx.y, blk = blockIdx.x * blockDim.x + threadIdx.x;
+ if (blk * 16 >= K || row >= (tb ? mpad : M)) return;
+ const int r = tb ? row % tb : 0;
+ const size_t tiled = tb ? ((static_cast(row / tb) * (K / 64) + blk / 4) * tb + r) * 32 : 0;
+ const size_t at = tb == 0 ? static_cast(row) * (K / 2) + blk * 8
+ : tiled + mma4::chunk(r, (blk % 4) / 2) * 16 + (blk % 2) * 8;
+ const size_t sat = (tb ? (static_cast(row / tb) * (K / 64) + blk / 4) * tb + r
+ : static_cast(blk / 4) * mpad + row) * 4 + blk % 4;
+ if (row >= M) {
+ *reinterpret_cast(codes + at) = make_uint2(0u, 0u);
+ scales[sat] = 0;
+ return;
+ }
+ const uint4* src = reinterpret_cast(x + static_cast(row) * ldx + blk * 16);
+ float f[16];
+ unpack8(src[0], *reinterpret_cast(f));
+ unpack8(src[1], *reinterpret_cast(f + 8));
+ float amax = 0.0f;
+#pragma unroll
+ for (int i = 0; i < 16; ++i) amax = fmaxf(amax, fabsf(f[i]));
+ const nvfp4q::Scale sc = nvfp4q::block_scale(amax, g);
+ uint32_t w[2] = {0u, 0u};
+#pragma unroll
+ for (int i = 0; i < 16; ++i) w[i / 8] |= nvfp4q::e2m1(f[i] * sc.mul) << (4 * (i % 8));
+ *reinterpret_cast(codes + at) = make_uint2(w[0], w[1]);
+ scales[sat] = static_cast(sc.sf8);
+}
+
+// One thread 16 inputs: e4m3(x * inv) written in the weights' fragment order (byte 4q + j of a 16-byte group holds
+// input 2q + (j % 2) + 8 (j / 2)), so a lane's four bytes pair with its weight bytes in the e4m3 mma; ``tb``: the
+// prompt GEMM's swizzled tiles [mpad / tb][K/64][tb][64 bytes] (rows past M zero).
+__global__ void __launch_bounds__(128) quant8_kernel(const __nv_bfloat16* __restrict__ x, int ldx, int M, int K,
+ float inv, uint8_t* __restrict__ out, int mpad, int tb) {
+ const int row = blockIdx.y, grp = blockIdx.x * blockDim.x + threadIdx.x;
+ if (grp * 16 >= K || row >= (tb ? mpad : M)) return;
+ const int r = tb ? row % tb : 0;
+ const size_t at = tb == 0 ? static_cast(row) * K + grp * 16
+ : ((static_cast(row / tb) * (K / 64) + grp / 4) * tb + r) * 64 + mma4::chunk(r, grp % 4) * 16;
+ if (row >= M) {
+ *reinterpret_cast(out + at) = make_uint4(0u, 0u, 0u, 0u);
+ return;
+ }
+ const uint4* src = reinterpret_cast(x + static_cast(row) * ldx + grp * 16);
+ float f[16];
+ unpack8(src[0], *reinterpret_cast(f));
+ unpack8(src[1], *reinterpret_cast(f + 8));
+ uint32_t w[4];
+#pragma unroll
+ for (int q = 0; q < 4; ++q) {
+ uint32_t packed = 0;
+#pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ const int k = 2 * q + (j % 2) + 8 * (j / 2);
+ packed |= static_cast(__nv_cvt_float_to_fp8(f[k] * inv, __NV_SATFINITE, __NV_E4M3)) << (8 * j);
+ }
+ w[q] = packed;
+ }
+ *reinterpret_cast(out + at) = make_uint4(w[0], w[1], w[2], w[3]);
+}
+
+} // namespace
+
+void quant4_cuda(const at::Tensor& x, double g, at::Tensor& codes, at::Tensor& scales, int64_t mpad, int64_t tb) {
+ const int M = static_cast(x.size(0)), K = static_cast(x.size(1));
+ const dim3 grid(static_cast((K / 16 + 127) / 128), static_cast(tb ? mpad : M));
+ quant4_kernel<<>>(
+ reinterpret_cast(x.data_ptr()), static_cast(x.stride(0)), M, K,
+ static_cast(g), codes.data_ptr(), scales.data_ptr(), static_cast(mpad),
+ static_cast(tb));
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void quant8_cuda(const at::Tensor& x, double inv, at::Tensor& out, int64_t mpad, int64_t tb) {
+ const int M = static_cast(x.size(0)), K = static_cast(x.size(1));
+ const dim3 grid(static_cast((K / 16 + 127) / 128), static_cast(tb ? mpad : M));
+ quant8_kernel<<>>(
+ reinterpret_cast(x.data_ptr()), static_cast(x.stride(0)), M, K,
+ static_cast(inv), out.data_ptr(), static_cast(mpad), static_cast(tb));
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
diff --git a/src/tensorfold/cuda/nvfp4/checkpoint.cpp b/src/tensorfold/cuda/nvfp4/checkpoint.cpp
new file mode 100644
index 000000000..3c6d53f39
--- /dev/null
+++ b/src/tensorfold/cuda/nvfp4/checkpoint.cpp
@@ -0,0 +1,171 @@
+#include
+#include
+
+void quant4_cuda(const at::Tensor&, double, at::Tensor&, at::Tensor&, int64_t, int64_t);
+void quant8_cuda(const at::Tensor&, double, at::Tensor&, int64_t, int64_t);
+void lane_cuda(int64_t, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, double, at::Tensor&,
+ const at::Tensor&, int64_t, int64_t, int64_t, int64_t, int64_t, bool);
+void pack4_cuda(const at::Tensor&, int64_t, int64_t, at::Tensor&);
+void gemm_cuda(int64_t, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, double, at::Tensor&,
+ int64_t, int64_t, int64_t, int64_t, bool);
+void gemm_gu_ck_cuda(const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&,
+ const at::Tensor&, double, double, int64_t, int64_t, int64_t, double, at::Tensor&, at::Tensor&,
+ bool);
+void gemm_ws_cuda(int64_t, const at::Tensor&, const at::Tensor&, const at::Tensor&, const at::Tensor&, double,
+ at::Tensor&, int64_t, int64_t, int64_t, int64_t, int64_t, bool);
+
+static void rows_in(const at::Tensor& x) {
+ TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1 &&
+ x.stride(0) % 8 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0 && x.size(1) % 64 == 0,
+ "x: (M, K) bf16 rows, 16-byte aligned, K a multiple of 64");
+}
+
+// NVFP4 rows under the static global scale ``g`` (1 / input_scale): codes [M, K/2] and scales [K/64, mpad, 4], or
+// with ``tb`` > 0 the prompt GEMM's tiles: codes [mpad / tb][K/64][tb][32], scales [mpad / tb, K/64, tb, 4].
+void quant4(const at::Tensor& x, double g, at::Tensor codes, at::Tensor scales, int64_t tb) {
+ rows_in(x);
+ const int64_t m = x.size(0), k = x.size(1);
+ TORCH_CHECK(tb == 0 || (tb % 16 == 0 && scales.dim() == 4 && scales.size(1) == k / 64 && scales.size(2) == tb &&
+ scales.size(3) == 4), "tb: 0, or a multiple of 16 with scales (mpad / tb, K/64, tb, 4)");
+ TORCH_CHECK(tb || (scales.dim() == 3 && scales.size(0) == k / 64 && scales.size(2) == 4),
+ "scales: (K/64, mpad, 4)");
+ const int64_t mpad = tb ? scales.size(0) * tb : scales.size(1);
+ TORCH_CHECK(scales.is_contiguous() && scales.scalar_type() == at::kByte && mpad >= m && mpad % 64 == 0,
+ "scales: uint8, mpad a multiple of 64 holding every row");
+ TORCH_CHECK(codes.is_contiguous() && codes.scalar_type() == at::kByte &&
+ codes.numel() == (tb ? mpad : m) * k / 2, "codes: (M, K/2) uint8, (mpad, K/2) when tiled");
+ c10::cuda::CUDAGuard guard(x.device());
+ quant4_cuda(x, g, codes, scales, mpad, tb);
+}
+
+// FP8 rows under the static scale (``inv`` = 1 / input_scale), e4m3 bytes in the weights' fragment order; with
+// ``tb`` > 0 the prompt GEMM's tiles [mpad / tb][K/64][tb][64] for out (mpad, K), rows past M zero.
+void quant8(const at::Tensor& x, double inv, at::Tensor out, int64_t tb) {
+ rows_in(x);
+ const int64_t mpad = tb ? out.size(0) : x.size(0);
+ TORCH_CHECK(tb == 0 || (tb % 16 == 0 && mpad % tb == 0 && mpad >= x.size(0)), "tb: 0 or a multiple of 16 "
+ "dividing the padded rows");
+ TORCH_CHECK(out.is_contiguous() && out.scalar_type() == at::kByte && out.size(0) == mpad &&
+ out.size(1) == x.size(1), "out: (M, K) uint8, (mpad, K) when tiled");
+ c10::cuda::CUDAGuard guard(x.device());
+ quant8_cuda(x, inv, out, mpad, tb);
+}
+
+// out (M, n) = alpha * rows @ weight: mode 0 NVFP4 x NVFP4 (``quant4`` rows, ``pack4`` words, block scales
+// [npad/64, K/64, 64, 4]), mode 1 FP8 x FP8 (``quant8`` rows, fragment-order bytes); part (SK, M, n) past 8 slices.
+// ``tile`` (``checkpoint.lane_tile``) picks the block's shape, never the bits.
+void lane(int64_t mode, const at::Tensor& x, const c10::optional& xs, const at::Tensor& w,
+ const c10::optional& ws, double alpha, at::Tensor out, const c10::optional& part,
+ int64_t n, int64_t k, int64_t sk, int64_t npad, int64_t tile, bool f32) {
+ TORCH_CHECK(mode == 0 || mode == 1, "mode 0 (NVFP4) or 1 (FP8)");
+ const int64_t m = out.size(0);
+ TORCH_CHECK(k % 64 == 0 && (k / 64) % sk == 0, "K in whole steps of 64, split evenly");
+ TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kByte && x.size(0) == m &&
+ x.size(1) == (mode == 0 ? k / 2 : k), "x: quantized rows (M, K/2 | K)");
+ TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.numel() * w.element_size() == npad * k / (mode == 0 ? 2 : 1),
+ "weight bytes do not match npad and K");
+ int64_t mpad = 0;
+ if (mode == 0) {
+ TORCH_CHECK(xs.has_value() && xs->is_contiguous() && xs->size(0) == k / 64 && xs->size(1) >= m &&
+ xs->size(1) % 64 == 0, "xs: (K/64, mpad, 4) row scales");
+ TORCH_CHECK(ws.has_value() && ws->is_contiguous() && ws->numel() == (k / 64) * npad * 4,
+ "ws: [npad/64, K/64, 64, 4] block scales");
+ mpad = xs->size(1);
+ }
+ TORCH_CHECK(out.is_cuda() && out.is_contiguous() && out.size(1) == n &&
+ out.scalar_type() == (f32 ? at::kFloat : at::kBFloat16), "out: (M, n)");
+ TORCH_CHECK(sk <= 8 || (part.has_value() && part->numel() >= sk * m * n), "part: (SK, M, n) fp32");
+ c10::cuda::CUDAGuard guard(x.device());
+ lane_cuda(mode, x, xs.has_value() ? *xs : at::Tensor(), w, ws.has_value() ? *ws : at::Tensor(), alpha, out,
+ part.has_value() ? *part : at::Tensor(), n, k, sk, mpad, tile, f32);
+}
+
+// Prompt rows: out (M, n) = alpha * rows @ weight in ``lane``'s layouts, one K chain a row (``tile`` picks the shape).
+void gemm(int64_t mode, const at::Tensor& x, const c10::optional& xs, const at::Tensor& w,
+ const c10::optional& ws, double alpha, at::Tensor out, int64_t n, int64_t k, int64_t npad,
+ int64_t tile, bool f32) {
+ TORCH_CHECK(mode == 0 || mode == 1, "mode 0 (NVFP4) or 1 (FP8)");
+ const int64_t m = out.size(0);
+ TORCH_CHECK(k % 64 == 0 && npad % 64 == 0 && n <= npad, "K in steps of 64, npad whole 64-column tiles");
+ TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kByte && x.size(0) == m &&
+ x.size(1) == (mode == 0 ? k / 2 : k), "x: quantized rows (M, K/2 | K)");
+ TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.numel() * w.element_size() == npad * k / (mode == 0 ? 2 : 1),
+ "weight bytes do not match npad and K");
+ if (mode == 0)
+ TORCH_CHECK(xs.has_value() && xs->is_contiguous() && xs->size(0) == k / 64 && xs->size(1) >= m &&
+ xs->size(1) % 64 == 0 && ws.has_value() && ws->is_contiguous() &&
+ ws->numel() == (k / 64) * npad * 4, "xs: (K/64, mpad, 4) row scales; ws: [npad/64, K/64, 64, 4]");
+ TORCH_CHECK(out.is_cuda() && out.is_contiguous() && out.size(1) == n &&
+ out.scalar_type() == (f32 ? at::kFloat : at::kBFloat16), "out: (M, n)");
+ c10::cuda::CUDAGuard guard(x.device());
+ gemm_cuda(mode, x, xs.has_value() ? *xs : at::Tensor(), w, ws.has_value() ? *ws : at::Tensor(), alpha, out, n, k,
+ npad, tile, f32);
+}
+
+// NVFP4 checkpoint bytes [N, K/2] -> ``lane``'s words [npad/64, K/64, 8, 32, 2] (int32).
+void pack4(const at::Tensor& src, at::Tensor dst) {
+ TORCH_CHECK(src.is_cuda() && src.is_contiguous() && src.scalar_type() == at::kByte && src.dim() == 2,
+ "src: (N, K/2) uint8");
+ const int64_t n = src.size(0), k = src.size(1) * 2;
+ TORCH_CHECK(k % 64 == 0 && dst.is_contiguous() && (dst.numel() * dst.element_size()) % (k / 2 * 64) == 0 &&
+ dst.numel() * dst.element_size() >= n * k / 2, "dst: [npad/64, K/64, 8, 32, 2] words");
+ c10::cuda::CUDAGuard guard(src.device());
+ pack4_cuda(src, n, k, dst);
+}
+
+// Prompt rows on the warp-specialized GEMM: x the tiled rows (``quant4`` / ``quant8`` with tb 128), else as ``gemm``.
+void gemm_ws(int64_t mode, const at::Tensor& x, const c10::optional& xs, const at::Tensor& w,
+ const c10::optional& ws, double alpha, at::Tensor out, int64_t n, int64_t k, int64_t npad,
+ int64_t tile, bool f32) {
+ TORCH_CHECK(mode == 0 || mode == 1, "mode 0 (NVFP4) or 1 (FP8)");
+ const int64_t m = out.size(0), mpad = x.size(0);
+ TORCH_CHECK(k % 64 == 0 && npad % 64 == 0 && n <= npad, "K in steps of 64, npad whole 64-column tiles");
+ TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kByte && mpad % 128 == 0 && mpad >= m &&
+ x.size(1) == (mode == 0 ? k / 2 : k), "x: tiled rows (mpad, K/2 | K), mpad a multiple of 128");
+ TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.numel() * w.element_size() == npad * k / (mode == 0 ? 2 : 1),
+ "weight bytes do not match npad and K");
+ if (mode == 0)
+ TORCH_CHECK(xs.has_value() && xs->is_contiguous() && xs->dim() == 4 && xs->size(0) * xs->size(2) == mpad &&
+ xs->size(1) == k / 64 && xs->size(2) == 128 && ws.has_value() && ws->is_contiguous() &&
+ ws->numel() == (k / 64) * npad * 4, "xs: (mpad / 128, K/64, 128, 4); ws: [npad/64, K/64, 64, 4]");
+ TORCH_CHECK(out.is_cuda() && out.is_contiguous() && out.size(1) == n &&
+ out.scalar_type() == (f32 ? at::kFloat : at::kBFloat16), "out: (M, n)");
+ c10::cuda::CUDAGuard guard(x.device());
+ gemm_ws_cuda(mode, x, xs.has_value() ? *xs : at::Tensor(), w, ws.has_value() ? *ws : at::Tensor(), alpha, out, n,
+ k, mpad, npad, tile, f32);
+}
+
+// gate|up of the same NVFP4 rows (``quant4``'s row-major layout) -> SiLU(gate) * up -> down's input as NVFP4 rows
+// under down's global scale ``qg`` through the cp.async GEMM: codes (M, npad/2), scales (npad/64, mpad, 4).
+void gemm_gu_ck(const at::Tensor& x, const at::Tensor& xs, const at::Tensor& wg, const at::Tensor& wsg,
+ const at::Tensor& wu, const at::Tensor& wsu, double alpha_g, double alpha_u, int64_t npad, int64_t k,
+ double qg, at::Tensor codes, at::Tensor scales, bool fp32) {
+ const int64_t m = x.size(0), mpad = xs.size(1);
+ TORCH_CHECK(k % 128 == 0 && npad % 64 == 0, "K in steps of 128, npad whole 64-column tiles");
+ TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type() == at::kByte && x.size(1) == k / 2,
+ "x: NVFP4 rows (M, K/2)");
+ TORCH_CHECK(xs.is_contiguous() && xs.dim() == 3 && xs.size(0) == k / 64 && mpad >= m && mpad % 64 == 0,
+ "xs: (K/64, mpad, 4) row scales");
+ for (const auto* w : {&wg, &wu})
+ TORCH_CHECK(w->is_cuda() && w->is_contiguous() && w->numel() * w->element_size() == npad * k / 2,
+ "weight bytes do not match npad and K");
+ for (const auto* ws : {&wsg, &wsu})
+ TORCH_CHECK(ws->is_contiguous() && ws->numel() == (k / 64) * npad * 4, "ws: [npad/64, K/64, 64, 4]");
+ TORCH_CHECK(codes.is_contiguous() && codes.scalar_type() == at::kByte && codes.size(0) == m &&
+ codes.size(1) == npad / 2, "codes: (M, npad/2) uint8");
+ TORCH_CHECK(scales.is_contiguous() && scales.scalar_type() == at::kByte && scales.dim() == 3 &&
+ scales.size(0) == npad / 64 && scales.size(1) == mpad && scales.size(2) == 4,
+ "scales: (npad/64, mpad, 4) uint8");
+ c10::cuda::CUDAGuard guard(x.device());
+ gemm_gu_ck_cuda(x, xs, wg, wsg, wu, wsu, alpha_g, alpha_u, m, npad, k, qg, codes, scales, fp32);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("quant4", &quant4);
+ m.def("quant8", &quant8);
+ m.def("lane", &lane);
+ m.def("pack4", &pack4);
+ m.def("gemm", &gemm);
+ m.def("gemm_ws", &gemm_ws);
+ m.def("gemm_gu_ck", &gemm_gu_ck);
+}
diff --git a/src/tensorfold/cuda/nvfp4/checkpoint.py b/src/tensorfold/cuda/nvfp4/checkpoint.py
new file mode 100644
index 000000000..dcf209f70
--- /dev/null
+++ b/src/tensorfold/cuda/nvfp4/checkpoint.py
@@ -0,0 +1,220 @@
+"""NVFP4 checkpoints in their own math: rows quantized under the static input scales, FP4 x FP4 and FP8 x FP8."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from pathlib import Path
+from typing import NamedTuple
+
+import torch
+
+A4, A8 = 0, 1
+SWIGLU_FP32 = True # the fused gate|up epilogue's SwiGLU in fp32 (False: through bf16, as the unfused path)
+WS = 10 # prompt tiles past WS run the warp-specialized GEMM on rows quantized into its tiles
+TB = 128 # its tile rows
+
+
+def available() -> bool:
+ """Whether this GPU has the block-scaled FP4 mma (compute capability 12.x)."""
+
+ return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12
+
+
+def bulk_tile(tile: int, capability: tuple[int, int]) -> bool:
+ """Whether prompt ``tile`` runs the bulk-copy GEMM (past ``WS``, from sm_90); else gemm_ck's tile, same bits."""
+
+ return tile > WS and tuple(int(v) for v in capability) >= (9, 0)
+
+
+@lru_cache(maxsize=1)
+def _capability() -> tuple[int, int]:
+ return tuple(torch.cuda.get_device_capability())
+
+
+@lru_cache(maxsize=1)
+def _ext():
+ from tensorfold.cuda import precision
+ from tensorfold.cuda.build import MIN_CAPABILITY, load
+
+ here = Path(__file__).parent
+ return load(name="tensorfold_nvfp4_ck_v6",
+ sources=[str(here / "checkpoint.cpp"), str(here / "act.cu"), str(here / "lane4.cu"),
+ str(here / "gemm_ck.cu"), str(here / "gemm_ws.cu")],
+ need=MIN_CAPABILITY, arch_specific=precision.own_math(_capability())["nvfp4"],
+ extra_include_paths=[str(here)], extra_cuda_cflags=["-O3"], verbose=False)
+
+
+class Rows4(NamedTuple):
+ """Rows in NVFP4: e2m1 codes [M, K/2] (input 2j in byte j's low nibble), e4m3 scales [K/64, mpad, 4]."""
+
+ codes: torch.Tensor
+ scales: torch.Tensor
+
+
+def _rows(x: torch.Tensor) -> torch.Tensor:
+ if x.dtype != torch.bfloat16 or x.stride(-1) != 1 or (x.shape[0] > 1 and x.stride(0) % 8) or x.data_ptr() % 16:
+ x = x.to(torch.bfloat16).contiguous()
+ return x
+
+
+def _inv(act: float) -> float:
+ return float(torch.tensor(1.0, dtype=torch.float32) / torch.tensor(act, dtype=torch.float32))
+
+
+def quant4(x: torch.Tensor, act: float, tb: int = 0) -> Rows4:
+ """bf16 rows -> NVFP4 under input scale ``act``: per-16 e4m3 scales, e2m1 codes to nearest even (``tb``: tiled)."""
+
+ x = _rows(x)
+ m, k = x.shape
+ mpad = -(-m // (tb or 64)) * (tb or 64)
+ codes = torch.empty((mpad if tb else m, k // 2), dtype=torch.uint8, device=x.device)
+ shape = (mpad // tb, k // 64, tb, 4) if tb else (k // 64, mpad, 4)
+ scales = (torch.empty if tb else torch.zeros)(shape, dtype=torch.uint8, device=x.device)
+ _ext().quant4(x, _inv(act), codes, scales, tb)
+ return Rows4(codes, scales)
+
+
+def quant8(x: torch.Tensor, act: float, tb: int = 0) -> torch.Tensor:
+ """bf16 rows -> e4m3(x / act), saturating, in the FP8 weights' fragment order (``tb``: the GEMM's tiled rows)."""
+
+ x = _rows(x)
+ m, k = x.shape
+ out = torch.empty((-(-m // tb) * tb if tb else m, k), dtype=torch.uint8, device=x.device)
+ _ext().quant8(x, _inv(act), out, tb)
+ return out
+
+
+def pack4(weight: torch.Tensor, npad: int) -> torch.Tensor:
+ """NVFP4 checkpoint bytes [N, K/2] -> the FP4 lane matmul's words [npad/64, K/64, 8, 32, 2] (int32)."""
+
+ n, k = weight.shape[0], weight.shape[1] * 2
+ words = torch.empty((npad // 64, k // 64, 8, 32, 2), dtype=torch.int32, device=weight.device)
+ _ext().pack4(weight.contiguous().view(torch.uint8), words)
+ return words
+
+
+def alpha(act: float, scale: float) -> float:
+ """The output factor: the input scale times the weight's (fp32, as the checkpoint's runtimes compute it)."""
+
+ return float(torch.tensor(act, dtype=torch.float32) * torch.tensor(scale, dtype=torch.float32))
+
+
+def _rowsq(mode: int, x: torch.Tensor, act: float, tb: int = 0):
+ """Rows quantized for ``mode`` under the static input scale ``act`` (``tb``: the prompt GEMM's tiles)."""
+
+ return quant4(x, act, tb) if mode == A4 else (quant8(x, act, tb), None)
+
+
+def _out(m: int, lin, out: torch.Tensor | None, f32: bool) -> torch.Tensor:
+ want = torch.float32 if f32 else torch.bfloat16
+ return out if out is not None and out.is_contiguous() and out.dtype == want else \
+ torch.empty((m, lin.n), dtype=want, device=lin.words.device if hasattr(lin, "words") else lin.w8.device)
+
+
+FILL = 2 # 128-wide lane blocks while at least SMs / FILL of them fill the GPU, else 64-wide
+WIDE = 96 # SMs from which wide lane blocks pay: a 48-SM GB10 streams faster on 64-wide ones
+
+
+def lane_tile(m: int, n: int, sk: int, sms: int) -> int:
+ """The lane matmul's block (BM * 1000 + BN past 32 rows) for m rows, n columns, sk K slices on sms SMs."""
+
+ if m <= 32:
+ return 16 if m <= 16 else 32 # 16 or 32 rows by 64 columns; blocks never change bits
+ rows = 64 if m <= 64 else 128
+ if sms < WIDE:
+ return 64 # 64 x 64 blocks side by side
+ wide = -(-m // rows) * -(-n // 128) * sk * FILL >= sms
+ return (64128 if wide else 64) if rows == 64 else (128128 if wide else 128064)
+
+
+@lru_cache(maxsize=8)
+def _sms(device: int) -> int:
+ return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+def sm_count() -> int:
+ """SMs on the current GPU (for block shapes only: never a bit of output)."""
+
+ return _sms(torch.cuda.current_device())
+
+
+def _lane(mode: int, rows, lin, y: torch.Tensor, f32: bool) -> None:
+ from tensorfold.cuda.kernels import qmm
+
+ m = y.shape[0]
+ sk = qmm.split_k(lin.n, lin.k)
+ part = torch.empty((sk, m, lin.n), dtype=torch.float32, device=y.device) if sk > 8 else None
+ w, ws = (lin.words, lin.bs) if mode == A4 else (lin.w8, None)
+ _ext().lane(mode, rows[0], rows[1], w, ws, alpha(lin.act, lin.scale), y, part, lin.n, lin.k, sk, lin.npad,
+ lane_tile(m, lin.n, sk, sm_count()), f32)
+
+
+def _tb(tile: int) -> int:
+ """The quantizer's tile rows for a prompt ``tile``: the bulk-copy GEMM's, or 0 for row-major rows."""
+
+ return TB if bulk_tile(tile, _capability()) else 0
+
+
+def _gemm(mode: int, rows, lin, y: torch.Tensor, f32: bool, tile: int = 0) -> None:
+ w, ws = (lin.words, lin.bs) if mode == A4 else (lin.w8, None)
+ run = _ext().gemm_ws if bulk_tile(tile, _capability()) else _ext().gemm
+ run(mode, rows[0], rows[1], w, ws, alpha(lin.act, lin.scale), y, lin.n, lin.k, lin.npad,
+ tile % WS if tile > WS else tile, f32)
+
+
+def matmul(mode: int, x: torch.Tensor, lin, out: torch.Tensor | None = None, f32: bool = False) -> torch.Tensor:
+ """bf16 rows (M, K) @ an NVFP4 (``A4``) or FP8 (``A8``) linear in checkpoint mode -> (M, n); K slices by shape."""
+
+ y = _out(x.shape[0], lin, out, f32)
+ _lane(mode, _rowsq(mode, x, lin.act), lin, y, f32)
+ if out is not None and y is not out:
+ out.copy_(y)
+ return y
+
+
+def prompt(mode: int, x: torch.Tensor, lin, out: torch.Tensor | None = None, f32: bool = False,
+ tile: int = 0) -> torch.Tensor:
+ """Prompt rows (M, K) bf16 @ an NVFP4 or FP8 linear in checkpoint math: one K chain a row (chunk-invariant bits)."""
+
+ y = _out(x.shape[0], lin, out, f32)
+ _gemm(mode, _rowsq(mode, x, lin.act, _tb(tile)), lin, y, f32, tile)
+ if out is not None and y is not out:
+ out.copy_(y)
+ return y
+
+
+def matmul_group(x: torch.Tensor, lins: list, prompt_rows: bool = False, outs: list[torch.Tensor] | None = None,
+ tile: int = 0) -> list[torch.Tensor] | None:
+ """Projections of one input under one input scale: rows quantized once, each its own bits; None if they differ."""
+
+ if len(lins) < 2 or any(getattr(lin, "act", None) is None for lin in lins):
+ return None
+ modes = {A4 if hasattr(lin, "words") else A8 for lin in lins}
+ if len(modes) != 1 or len({float(lin.act) for lin in lins}) != 1:
+ return None
+ mode = modes.pop()
+ rows = _rowsq(mode, x, lins[0].act, _tb(tile) if prompt_rows else 0)
+ outs = [_out(x.shape[0], lin, None if outs is None else outs[j], False) for j, lin in enumerate(lins)]
+ for lin, y in zip(lins, outs):
+ _gemm(mode, rows, lin, y, False, tile) if prompt_rows else _lane(mode, rows, lin, y, False)
+ return outs
+
+
+def mlp_prompt(x: torch.Tensor, gate, up, down, out: torch.Tensor | None = None) -> torch.Tensor | None:
+ """Prompt rows through gate|up -> SiLU(gate) * up -> down, the SwiGLU rows leaving as down's NVFP4 rows; or None."""
+
+ if not all(hasattr(lin, "words") and getattr(lin, "act", None) is not None for lin in (gate, up, down)):
+ return None
+ if float(gate.act) != float(up.act) or gate.n != up.n or gate.npad != gate.n or down.k != gate.n:
+ return None
+ if gate.k % 128 or down.k % 128:
+ return None
+ rows = quant4(x, gate.act)
+ m, mpad = x.shape[0], rows.scales.shape[1]
+ codes = torch.empty((m, down.k // 2), dtype=torch.uint8, device=x.device)
+ scales = torch.empty((down.k // 64, mpad, 4), dtype=torch.uint8, device=x.device)
+ _ext().gemm_gu_ck(rows.codes, rows.scales, gate.words, gate.bs, up.words, up.bs, alpha(gate.act, gate.scale),
+ alpha(up.act, up.scale), gate.npad, gate.k, _inv(down.act), codes, scales, SWIGLU_FP32)
+ y = _out(m, down, out, False)
+ _gemm(A4, Rows4(codes, scales), down, y, False, 0)
+ return y
diff --git a/src/tensorfold/cuda/nvfp4/gemm_ck.cu b/src/tensorfold/cuda/nvfp4/gemm_ck.cu
new file mode 100644
index 000000000..6d2d9a9bf
--- /dev/null
+++ b/src/tensorfold/cuda/nvfp4/gemm_ck.cu
@@ -0,0 +1,268 @@
+// Prompt GEMM in the checkpoint's own math: NVFP4 rows times NVFP4 weights (block-scaled FP4 mma) and FP8 rows times
+// FP8 weights (e4m3 mma), in lane4.cu's layouts. One fp32 chain over K a row: a row's bits never depend on its chunk.
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include "../kernels/qmm_frag.cuh"
+#include "mma4.cuh"
+#include "swiglu4.cuh"
+
+namespace {
+
+using namespace qmm_frag;
+using namespace mma4;
+
+template
+struct Gemm {
+ static constexpr int THREADS = WM * WN * 32;
+ static constexpr int MT = BM / WM / 16, NT = BN / WN / 8; // m16 and n8 tiles a warp
+ static constexpr int ROW = MODE == A4 ? 32 : 64; // bytes a row a step of 64 inputs
+ static constexpr int TILE = MODE == A4 ? 2048 : 4096; // a stored 64-column tile's step
+ static constexpr int X = BM * ROW, SX = MODE == A4 ? BM * 4 : 0; // a step's rows and row scales,
+ static constexpr int W = BN / 64 * TILE, SW = MODE == A4 ? BN * 4 : 0; // weights and their scales
+ static constexpr int STEP = X + SX + W + SW;
+ static constexpr int STAGE = (KS * STEP + 127) / 128 * 128; // KS steps a stage
+ static constexpr int SMEM = STAGES * STAGE;
+};
+
+// The fused gate|up launch's second weight (up), its factor, and down's input it writes: NVFP4 rows [M, kd/2] and
+// scales [kd/64, mpad, 4] under down's global scale qg.
+struct Up {
+ const uint8_t* w;
+ const uint8_t* ws;
+ float alpha;
+ uint8_t* codes;
+ uint8_t* scales;
+ float qg;
+ int kd;
+};
+
+// x: A4 codes [M, K/2] and scales [K/64, mpad, 4]; A8 e4m3 [M, K] in fragment order. w: lane4.cu's words or bytes,
+// ws its block scales [npad/64, K/64, 64, 4]. out (M, N) = alpha * the product, every row one K chain. EPI 1, 2:
+// w is gate and up.w up, a block their same 128 columns, written as SiLU(gate) * up in NVFP4 (swiglu4.cuh).
+template
+__global__ void __launch_bounds__(WM * WN * 32) gemm_kernel(
+ const uint8_t* __restrict__ x, const uint8_t* __restrict__ xs, const uint8_t* __restrict__ w,
+ const uint8_t* __restrict__ ws, float alpha, void* __restrict__ out, int M, int N, int K, int mpad, int npad,
+ int group, Up up) {
+ using G = Gemm;
+ constexpr bool GU = EPI > 0;
+ static_assert(!GU || (MODE == A4 && BM == 128 && BN == 256 && WM == 2 && WN == 4), "128 x 256, 64 x 64 warps");
+ extern __shared__ __align__(128) unsigned char buf[];
+ const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, wm = warp / WN, wn = warp % WN;
+ const int g = lane >> 2, t = lane & 3, q = lane >> 3, rr = lane & 7;
+ const int2 at = tile_of(blockIdx.x, M, npad, BM, GU ? BN / 2 : BN, group);
+ const int m0 = at.x, n0 = at.y, KG = K / 64, tiles = npad / 64;
+ constexpr int CH = G::ROW / 16;
+
+ auto stage = [&](int s) { return buf + s * G::STAGE; };
+ auto load = [&](unsigned char* p, int kg) { // one step of 64 inputs at p
+#pragma unroll
+ for (int c = tid; c < BM * CH; c += G::THREADS) {
+ const int r = c / CH, ch = c % CH;
+ const bool in = m0 + r < M;
+ cp16z(p + r * G::ROW + chunk(r, ch) * 16,
+ x + static_cast(in ? m0 + r : 0) * (K / (MODE == A4 ? 2 : 1)) + kg * G::ROW + ch * 16, in);
+ }
+ if constexpr (MODE == A4) {
+ for (int c = tid; c < BM / 4; c += G::THREADS)
+ cp16z(p + G::X + c * 16, xs + (static_cast(kg) * mpad + m0 + 4 * c) * 4, m0 + 4 * c < mpad);
+ }
+ unsigned char* pw = p + G::X + G::SX;
+#pragma unroll
+ for (int c = tid; c < G::W / 16; c += G::THREADS) {
+ const int tl = c / (G::TILE / 16), off = c % (G::TILE / 16), wt = n0 / 64 + (GU ? tl % 2 : tl);
+ const size_t src = (static_cast(min(wt, tiles - 1)) * KG + kg) * G::TILE + off * 16;
+ cp16z(pw + c * 16, (GU && tl >= 2 ? up.w : w) + src, wt < tiles);
+ }
+ if constexpr (MODE == A4) {
+ for (int c = tid; c < G::SW / 16; c += G::THREADS) {
+ const int tl = c / 16, off = c % 16, wt = n0 / 64 + (GU ? tl % 2 : tl);
+ cp16z(pw + G::W + c * 16, (GU && tl >= 2 ? up.ws : ws) + (static_cast(min(wt, tiles - 1)) *
+ KG + kg) * 256 + off * 16, wt < tiles);
+ }
+ }
+ };
+
+ float acc[G::MT][G::NT][4];
+#pragma unroll
+ for (int i = 0; i < G::MT; ++i)
+#pragma unroll
+ for (int j = 0; j < G::NT; ++j)
+#pragma unroll
+ for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f;
+ const int KT = KG / KS; // stages over K (KG a multiple of KS)
+ auto fill = [&](int s, int kt) {
+#pragma unroll
+ for (int u = 0; u < KS; ++u) load(stage(s) + u * G::STEP, kt * KS + u);
+ };
+#pragma unroll
+ for (int s = 0; s < STAGES - 1; ++s) {
+ if (s < KT) fill(s, s);
+ commit();
+ }
+ for (int kt = 0; kt < KT; ++kt) {
+ wait();
+ __syncthreads();
+ if (kt + STAGES - 1 < KT) fill((kt + STAGES - 1) % STAGES, kt + STAGES - 1);
+ commit();
+#pragma unroll
+ for (int u = 0; u < KS; ++u) {
+ const unsigned char* p = stage(kt % STAGES) + u * G::STEP;
+ const unsigned char* pw = p + G::X + G::SX;
+ if constexpr (MODE == A4) {
+ const uint32_t* sx = reinterpret_cast(p + G::X);
+ const uint32_t* sw = reinterpret_cast(pw + G::W);
+ uint32_t a[G::MT][4], sa[G::MT];
+#pragma unroll
+ for (int i = 0; i < G::MT; ++i) {
+ const int base = wm * (BM / WM) + i * 16, r = base + rr + (q & 1) * 8;
+ ldmatrix4(a[i], p + r * G::ROW + chunk(r, q >> 1) * 16);
+ const uint32_t v = sx[base + g + 8 * (t & 1)]; // row g (t 0), g + 8 (t 1): one load, no branch
+ sa[i] = v; // lanes 2, 3 repeat rows g, g + 8: never read
+ }
+#pragma unroll
+ for (int j = 0; j < G::NT; ++j) {
+ const int jj = wn * G::NT + j;
+ const uint2 b = reinterpret_cast