perf(gguf): sm_120 dispatch thresholds + upstream int8-MMA MMQ port (Ornith)

[lucaspirola]

What

Two related optimizations for Ornith-1.5-35B-Q4_K_M on RTX 5080 (sm_120,
Blackwell), on top of the existing Ornith INT4 262K-context serving path:

1. Arch-aware GGUF dispatch thresholds — the dense/MoE dequant-vs-MMQ
   crossover points measured on Ada don't hold on sm_120; make them
   architecture-gated instead of hardcoded.
2. int8-tensor-core MMQ port — the vendored GGUF kernels (llama.cpp
   b2899 via vLLM/sgl-kernel) are DP4A-only with zero tensor-core use.
   Upstream llama.cpp master has since rewritten Q4_K/Q6_K MMQ around
   int8 tensor-core MMA tiles (turing_mma, sm_75+). This PR vendors
   that upstream kernel verbatim (kernel/csrc/gguf_mmq/, llama.cpp
   eab8ee41) and wires it in for both dense projections and grouped
   MoE experts, gated to sm_120 for now.

FREETOKEN_GGUF_DISABLE_MMA=1 forces the old DP4A/dequant path (debugging,
or a toolchain that can't build the extension) — used below as the control
in every comparison.

Gains

Synthetic kernel benchmarks (real Ornith-1.5-35B-Q4_K_M tensors, sm_120, median of CUDA-event timings):

shape	rows/tokens	int8-MMA MMQ	DP4A MMQ (old)	dequant+cuBLAS	speedup vs DP4A
Q4_K attn_q [8192x2048] (dense)	8192 rows	1.79 ms	22.9 ms	2.40 ms	12.8x
Q6_K lm_head [248320x2048] (dense)	2048 tokens	17.5 ms	262 ms	19.2 ms	15.0x
Q4_K gate_up (MoE, E=256 top-8)	8192 tokens	4.16 ms	23.2 ms	n/a	5.6x
Q6_K down (MoE, E=256 top-8)	8192 tokens	4.90 ms	15.3 ms	n/a	3.1x

MMA beats dequant+cuBLAS too (the prior best option at large batch), not just
the DP4A kernel it replaces — e.g. attn_q 1.79ms vs 2.40ms, lm_head 17.5ms vs
19.2ms.

Live end-to-end A/B, full 262K-context serve config, same box, same prompt, minutes apart:

Config: --attention-backend triton --kv-cache-dtype q4_0 --num-tokens 262144 --kv-reserve-tokens 262144 --max-seq-len-override 262144 --max-running-requests 1 --moe-backend offload --moe-cache-auto --max-prefill-length 8192. Prompt: 28K words of seeded non-repetitive text
(~50K tokens, 6+ chunked prefill batches) with three distinct needles at
10/50/90% depth, greedy decode.

	int8-MMA (this PR)	DP4A/dequant (FREETOKEN_GGUF_DISABLE_MMA=1)
wall time	13.88 s	21.94 s
effective prefill rate	~4,700 tok/s	~2,700 tok/s
needles recovered	3/3 exact	3/3 exact
tracebacks	0	0

~1.75x prefill, ~1.58x end-to-end wall time, identical answers on the
same hardware and prompt.

Dispatch thresholds (sm_120 vs the existing Ada-tuned defaults, unchanged
off sm_120):

• dense dequant-vs-MMQ crossover: 24 rows (was 32) — Q4_K attention shapes
  cross at 24 (0.0645ms dequant vs 0.0778ms MMQ); 16 would regress the Q6_K
  lm_head, where MMQ still wins at 16.
• MoE grouped-MMQ-vs-vec crossover: 16 tokens (was 32) — 0.314ms MMQ vs
  0.324ms vec at 16 tokens, widening to 0.382 vs 0.475 at 24.
• Live-reverified on real Ornith tensors: 24 rows now 0.069ms via dequant vs
  0.081ms with the old threshold.

How

• python/freetoken/kernel/csrc/gguf_mmq/ — llama.cpp master's CUDA MMQ
  vendored verbatim (mmq/mma/load-tiles/vec-dot/configs/quantize/mmid + the
  ggml headers it needs). mmq_ext.cu is the only hand-written file: backend
  shims (device info, a torch-allocator-backed pool, error/abort plumbing)
  and the torch bindings for the dense (ggml_mul_mat_a8_mma) and grouped-MoE
  (ggml_moe_a8_mma) entry points. Only Q4_K/Q6_K mul_mat_q cases are
  instantiated to keep JIT compile time down.
• layers/gguf.py — dequant_gemm_min_rows(cc) (arch-gated threshold),
  _use_mma_mmq dispatch gate (sm_120 + supported type + successful JIT
  build), wired into fused_mul_mat_gguf.
• moe/fused_gguf.py — mmq_min_tokens(cc) (arch-gated threshold),
  _use_mma_moe gate (adds a block-alignment check on the padded expert
  slot stride), wired into _moe_matmul for the 320–16384 token band
  (below 320, DP4A wins — per-expert tiles waste work at small batches;
  MMVQ still owns decode).
• kernel/gguf.py — _mma_module() lazy JIT loader, mma_mmq_supported().

Testing

• Numerics: MMA output verified against the dequantized reference on real
  Ornith tensors (rel error <= 0.013, on par with the existing DP4A kernel)
  and against gguf-py on random-but-safe packed bytes
  (tests/kernels/test_gguf_mma.py). MoE broadcast (gate/up) and gather
  (down) forms verified against the per-expert dense reference and the
  existing vec kernel, end-to-end through fused_experts_gguf (rel ~1e-3).
• tests/kernels/test_gguf_dispatch.py — pure threshold-function tests for
  both archs, plus CUDA dispatch-branch tests with faked device capability.
• Kernel + benchmark suites: 336 passed, 1 skipped.
• Full non-slow suite (-m "not slow", excluding e2e/server): failure set
  unchanged from clean main (same 7 pre-existing failures).
• Live 262K-context A/B above, including a hostile (non-repetitive,
  multi-needle) prompt specifically to catch any tile/row-range corruption a
  compressible-filler prompt could hide.
• ruff check clean on all changed files.

Falls back cleanly to the existing DP4A/dequant path off sm_120, on a build
failure, or with FREETOKEN_GGUF_DISABLE_MMA=1; existing sm_89/Ada behavior
is unchanged (verified via the A/B, not just by inspection).

[zyy0212time-del]

I tested revision 242c37a4ecd7cccaed24d0ffdb5989e24fd81334 from PR #196 on a native Windows host with an RTX 5060 Laptop GPU (compute capability 12.0), driver 610.47, CUDA 13.0.48, Python 3.12.10, PyTorch 2.11.0+cu130, and FreeToken 0.1.2. The worktree was clean.

Model identity: PocketAiHub/Ornith-1.5-35B-A3B-Abliterated-GGUF, file Ornith-1.5-35B-A3B-Abliterated-Q4_K_M.gguf, GGUF Q4_K_M, 21,166,757,664 bytes, SHA-256 a07f299e83a398b5078c1cb8ab4ec96333c8ea9d10d0cb479cb073056fedd3d0 (GGUF metadata: qwen35moe, 256 experts, top-k 8). The model card attributes the conversion and validation to PocketAI Model Lab. The local download used the public resolve/main URL and the filename matches the model page exactly. The GGUF repository revision was not recorded in the local download metadata: revision not recorded / unavailable. The model card's pinned upstream base revision e4dfb35a93d4b6822a811a7676f3488514abe7e2 is recorded separately and is not being presented as the GGUF repository revision.

There is a separate Windows portability issue before model load: the current source uses ipc:///tmp/... endpoints, which the Windows pyzmq runtime rejects with Protocol not supported. For this validation only, I enabled a venv-local, opt-in TCP shim without modifying the worktree. It maps the scheduler endpoints to tcp://127.0.0.1:<server_port+2/+3/+4>, the frontend to +5, and the tokenizer to +6 when enabled.

With that shim, I ran the same command/configuration twice on the same machine (base ports 1926 and 1940): ft.exe serve --model-path <MODEL> --served-model-name ornith-196 --attention-backend triton --moe-backend offload --moe-cache-auto --num-tokens 16384 --max-running-requests 1 --max-prefill-length 8192 --max-seq-len-override 16384.

On this native Windows host and configuration, both valid launches hit the same load-blocking failure during host-bank registration. Both entered the Qwen3.5 MoE GGUF loader, reported expert banks: slow path (serial build), and then the backend terminated with the identical message: RuntimeError: cudaHostRegister failed for 0.3 GiB. The 0.3 GiB value is the failed registration buffer reported by the wrapper, not the total pinned-memory requirement. The service remained in loading; no request completed.

This observation is limited to the tested native Windows host/configuration. It is not evidence about PR #196's MMA kernel correctness, prefill/decode throughput, or a causal WDDM/pageable-locking failure. The inner CUDA error code, physical/system memory snapshot, WDDM quota, cumulative registered bytes, and other GPU-process state were not captured. I also cannot equate the result with issue #55 without that instrumentation.

I noticed the related native-Windows pin-budget discussion in #120 and #55; this report does not establish causality with either issue. Could maintainers advise whether native Windows is expected to use a pin-budget/residency strategy analogous to the WSL path, and whether the current IPC endpoints are intentionally Linux-only? Raw local paths and usernames are omitted; the exact local logs can be provided privately if useful.

[zyy0212time-del]

Follow-up: --moe-pageable-gpu skips the pin-budget path for this mixed-GGUF checkpoint

Follow-up to my earlier Windows validation comment on this PR. While investigating why --moe-pageable-gpu had no effect for this checkpoint, I traced the residency-planning path at runtime. All instrumentation below was venv-only; the FreeToken worktree remained clean.

Environment

• FreeToken: 242c37a4ecd7cccaed24d0ffdb5989e24fd81334 (this PR branch)
• Model: PocketAiHub/Ornith-1.5-35B-A3B-Abliterated-GGUF
• File: Ornith-1.5-35B-A3B-Abliterated-Q4_K_M.gguf
• SHA-256: A07F299E83A398B5078C1CB8AB4EC96333C8EA9D10D0CB479CB073056FEDD3D0
• Host: native Windows, RTX 5060 Laptop GPU (8151 MiB), 31.36 GiB RAM
• Relevant config: --moe-backend offload --moe-cache-auto --moe-pageable-gpu, with FREETOKEN_PIN_BUDGET_GB=15

Observed control flow

For this checkpoint, _auto_pageable_gpu_layers() evaluates:

ftw_bank_bytes(model_path) or bank_bytes_estimate(model_config)

Runtime instrumentation showed:

ftw_bank_bytes()       -> None
bank_bytes_estimate() -> None

The checkpoint has fmt="gguf" and gguf_expert_types metadata, but the current bank_bytes_estimate() path does not produce a value for this case: "gguf" has no _BANK_BYTES_PER_EXPERT entry, and the precise laguna_int4 branch is not taken.

Because bank_bytes is falsy, _auto_pageable_gpu_layers() returns an empty set before reaching _pin_budget_bytes(). The configured 15 GiB budget is therefore never consulted, no pageable layers are selected, and all 40 MoE layers are planned for pinned residency.

In the baseline run, execution reached layer 35 before registration failed:

• 70 successful cudaHostRegister calls
• next registration failed
• cumulative successful registration: 17.0215 GiB

That 17.0215 GiB value is specific to this host/run set; I am not treating it as a general Windows threshold.

Metadata-only sizing audit + causal check

Using the checkpoint's GGUF metadata and the in-tree geometry logic (row_bytes, BLOCK_SHAPE, _expert_bank_geometry), I get:

• 40 MoE layers
• per-layer HostBank allocation:
  256 × (1,179,648 + 860,160) = 522,190,848 B
• full routed-expert HostBank allocation:
  20,887,633,920 B = 19.453125 GiB

The two terms above are the per-expert gate/up and down strides selected by the uniform bank geometry. Twenty layers have Q4_K down experts, but the HostBank allocation uses the Q6_K-sized down stride uniformly.

I then overrode bank_bytes_estimate() through venv-only instrumentation gated by FT55_DEBUG_BANK_BYTES; no FreeToken source files were changed. This was a debug-only, checkpoint-specific causal test, not a proposed production fix.

	Baseline	Corrected sizing
bank_bytes seen by selector	None	19.453125 GiB
pin budget consulted	no	yes → 15 GiB
pageable layers	0	10 → [0,1,2,3,4,35,36,37,38,39]
planned pinned layers	40	30
layers reaching cudaHostRegister	36 before failure	30 complete
cudaHostRegister	70 OK + 1 FAIL	60 OK + 0 FAIL
cumulative registered	17.0215 GiB at failure	14.589844 GiB

The engine's own log independently reported the same 19.45 GiB vs 15.00 GiB budget decision and the same 10 pageable layers.

So, for this checkpoint and this commit, the observed causal chain is:

no bank-size estimate
→ pageable selector returns early
→ pin budget is never consulted
→ all MoE layers remain on the pinned path

Providing the correct checkpoint-specific sizing makes the existing budget-selection path behave as predicted in this test: 10 layers become pageable, the pinned HostBank footprint falls to 14.589844 GiB, and all 60 cudaHostRegister calls succeed.

Scope / stop point

I am not claiming that all mixed-GGUF checkpoints are affected, that there is a universal Windows pin threshold, or that the model now loads end-to-end.

After the host-registration blocker was removed, execution reached a subsequent distinct failure point while building freetoken_gguf_kernels through the JIT/ninja path. I have not investigated that failure, so I cannot attribute it to this PR, Windows, PyTorch, or my local toolchain.

This host-registration failure behavior is related context to #55, but the 17.0215 GiB threshold here is specific to this host, and I am not claiming this pageable-planner issue is the root cause of #55.

Question

Could the existing GGUF expert-type metadata and bank-geometry logic (gguf_expert_types, row_bytes, and _expert_bank_geometry) be used to provide bank_bytes for the fmt="gguf" path, so _auto_pageable_gpu_layers() can reach the configured pin-budget logic?

Happy to share the runtime traces and metadata-audit script if useful.

[zyy0212time-del]

Follow-up on the mixed-GGUF bank-sizing gap from my earlier validation on #196: I opened the isolated fix as a stacked PR on top of 242c37a4 here: lucaspirola#1.

bank_bytes_estimate() now covers expert_quant="gguf" using a shared expert_bank_geometry() helper, so the residency planner uses the same mixed-quant geometry as the qwen35moe/Laguna GGUF loaders rather than having no usable pre-load estimate. The important detail is that these banks use aligned flat [num_experts, stride] slots with the maximum gate/up and down geometry across participating quant types; raw GGUF tensor payload bytes therefore are not the quantity the planner needs.

The patch touches 5 files (+163/-17) and adds 8 CPU/synthetic repository regression tests covering mixed and homogeneous GGUF geometry, exact sizing, malformed/unsupported metadata, the existing non-GGUF path, and the shared geometry helper.

On Ornith-1.5-35B-A3B-Abliterated-Q4_K_M, the metadata-derived logical bank size is exactly 20,887,633,920 bytes. In my native Windows / RTX 5060 Laptop E2E run, with a 15 GiB pin budget and no FT55_DEBUG_BANK_BYTES override, the planner selected 10 pageable + 30 pinned MoE layers, completed 60/60 host registrations (14.589844 GiB cumulative), reached /health = ready, and returned HTTP 200 with the expected "OK" smoke response.

I stacked it here because the prerequisite qwen35moe/Laguna GGUF support is not currently present on main in a form that allows this to be submitted as a clean standalone upstream diff. If you'd prefer this to be rebased/retargeted once the prerequisite work lands on main, I'm happy to do that.