(WIP) falcon-ocr, for discussion

convert_hf_to_gguf.py

@@ -2630,6 +2630,147 @@
         yield from super().modify_tensors(data_torch, name, bid)
 
 
+@ModelBase.register("FalconOCRForCausalLM")
+class FalconOCRModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.FALCON_OCR
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+        # this model does not actually use the chat template, but we need to make sure to avoid any additional formatting
+        self.gguf_writer.add_chat_template("{% for m in messages %}{{ m['content'] + '\\n' }}{% endfor %}")
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        self.gguf_writer.add_context_length(hparams["max_seq_len"])
+        self.gguf_writer.add_feed_forward_length(hparams["ffn_dim"])
+
+        # head_dim (64) differs from n_embd/n_heads (768/16=48)
+        self.gguf_writer.add_key_length(hparams["head_dim"])
+        self.gguf_writer.add_value_length(hparams["head_dim"])
+
+        self.gguf_writer.add_layer_norm_rms_eps(hparams.get("norm_eps", 1e-5))
+        self.gguf_writer.add_rope_freq_base(hparams.get("rope_theta", 10000))
+        self.gguf_writer.add_rope_dimension_count(hparams["head_dim"] // 2)
+        self.gguf_writer.add_add_bos_token(False)
+
+        # important: because "golden" rope must be applied to fit Q shape,
+        # we must force number of KV heads to be the same as number of Q heads
+        self.gguf_writer.add_head_count_kv(hparams["n_heads"]) # not n_kv_heads
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if "freqs" in name:
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if "img_projector" in name:
+            return
+
+        if name == "freqs_cis_golden":
+            # original shape: [n_heads, rope_dim // 2, 2]
+            # permute to [2, n_heads, rope_dim//2] so h-freqs and w-freqs are contiguous,
+            # then flatten to [2, n_heads * rope_dim//2]
+            # ggml loads this as ne[0]=n_heads*rope_dim//2, ne[1]=2
+            data_torch = data_torch.permute(2, 0, 1).contiguous().reshape(2, -1)
+            # ggml_rope_ext computes theta = pos_int / freq_factor (freq_base=1.0)
+            # pos_int is fixed-point: pos_int = actual_pos * 1e6
+            # golden rope needs theta = freqs_actual * actual_pos = freqs_actual * pos_int / 1e6
+            # => freq_factor = 1e6 / freqs_actual
+            data_torch = 1e6 / data_torch
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), data_torch)
+            return
+
+        # Deinterleave fused w13 into separate gate (even rows) and up (odd rows)
+        if "feed_forward.w13" in name:
+            gate = data_torch[0::2, :]
+            up = data_torch[1::2, :]
+            yield from super().modify_tensors(gate, name.replace("w13", "w1"), bid)
+            yield from super().modify_tensors(up, name.replace("w13", "w3"), bid)
+            return
+
+        # Unfused w1 needs sqrt(2) scaling to match reference numerics
+        if "feed_forward.w1" in name:
+            data_torch = data_torch * math.sqrt(2.0)
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("FalconOCRForCausalLM")
+class FalconOCRMmprojModel(MmprojModel):
+    has_vision_encoder = True
+
+    # Important: Falcon OCR model does not actually have a vision encoder,
+    # the image patches are projected directly to the text embedding space and fed to the text encoder.
+    # we only use the mmproj to store the image projector weights and preprocessor config
+
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        # Inject synthetic vision_config / text_config for MmprojModel base class
+        hparams = ModelBase.load_hparams(dir_model, False)
+        hparams["text_config"] = {"hidden_size": hparams["dim"]}
+        hparams["vision_config"] = {
+            "hidden_size": hparams["dim"],
+            "patch_size": hparams["spatial_patch_size"],
+            "image_size": 1024,
+            "intermediate_size": hparams["dim"],
+            "num_attention_heads": 1,
+            "num_hidden_layers": 0, # no actual vision encoder layers
+        }
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+    def set_gguf_parameters(self):
+        self.gguf_writer.add_file_type(self.ftype)
+        self.gguf_writer.add_clip_has_vision_encoder(True)
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.FALCON_OCR)
+        self.gguf_writer.add_vision_projection_dim(self.n_embd_text)
+        self.gguf_writer.add_vision_patch_size(self.global_config["spatial_patch_size"])
+        self.gguf_writer.add_vision_image_size(1024)
+        self.gguf_writer.add_vision_embedding_length(self.global_config["dim"])
+        self.gguf_writer.add_vision_image_mean([0.5, 0.5, 0.5])
+        self.gguf_writer.add_vision_image_std([0.5, 0.5, 0.5])
+        self.gguf_writer.add_vision_min_pixels(64 * 64)
+        self.gguf_writer.add_vision_max_pixels(1024 * 1024)
+        self.gguf_writer.add_vision_head_count(1)
+        self.gguf_writer.add_vision_feed_forward_length(1)
+        self.gguf_writer.add_vision_block_count(0)
+        self.gguf_writer.add_vision_attention_layernorm_eps(1e-5)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if "tok_embeddings" in name:
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "img_projector.weight":
+            # The HF linear weight [n_embd, patch_dim] has patch_dim = H*W*C with C
+            # fastest (from einops rearrange). Rearrange to PyTorch conv2d format
+            # [C_out, C_in, kH, kW] so ggml_conv_2d can use it directly.  The planar
+            # image data fed to ggml matches this convention.
+            ps = self.global_config["spatial_patch_size"]   # 16
+            ch = self.global_config["channel_size"]         # 3
+            n_embd = data_torch.shape[0]
+            w = data_torch.reshape(n_embd, ps, ps, ch)     # [n_embd, H, W, C]
+            w = w.permute(0, 3, 1, 2).contiguous()         # [n_embd, C, H, W]
+            w = w.reshape(data_torch.shape)                 # flatten back to 2-D
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, bid=0), w)
+            return
+
+        if name == "tok_embeddings.weight":
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
+            prefix_str = "<|image_cls|><|image_reg_1|><|image_reg_2|><|image_reg_3|><|image_reg_4|>"
+            ids = tokenizer.encode(prefix_str, add_special_tokens=False)
+            prefix_embd = data_torch[ids].contiguous()
+            logger.info(f"Extracted {len(ids)} prefix embeddings (token IDs: {ids})")
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_TOK_IMG_BEGIN, suffix=""), prefix_embd)
+            return
+
+        return
+
+
 @ModelBase.register("GPTBigCodeForCausalLM")
 class StarCoderModel(TextModel):
     model_arch = gguf.MODEL_ARCH.STARCODER

gguf-py/gguf/constants.py

@@ -494,6 +494,7 @@ class MODEL_ARCH(IntEnum):
     LLAMA_EMBED      = auto()
     MAINCODER        = auto()
     KIMI_LINEAR      = auto()
+    FALCON_OCR       = auto()
 
 
 class VISION_PROJECTOR_TYPE(IntEnum):
@@ -980,6 +981,7 @@ class MODEL_TENSOR(IntEnum):
     MODEL_ARCH.LLAMA_EMBED:      "llama-embed",
     MODEL_ARCH.MAINCODER:        "maincoder",
     MODEL_ARCH.KIMI_LINEAR:      "kimi-linear",
+    MODEL_ARCH.FALCON_OCR:       "falcon-ocr",
 }
 
 VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
@@ -3877,6 +3879,18 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
     ],
+    MODEL_ARCH.FALCON_OCR: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
     # TODO
 }
 
@@ -4137,7 +4151,8 @@ class VisionProjectorType:
     GLM4V = "glm4v"
     YOUTUVL = "youtuvl"
     NEMOTRON_V2_VL = "nemotron_v2_vl"
-    HUNYUANOCR     = "hunyuanocr"
+    HUNYUANOCR = "hunyuanocr"
+    FALCON_OCR = "falcon-ocr"
 
 
 # Items here are (block size, type size)

gguf-py/gguf/tensor_mapping.py

@@ -230,6 +230,7 @@ class TensorNameMap:
             "layers.{bid}.attn.Wqkv",                                              # modern-bert
             "model.layers.{bid}.self_attn.language_expert_query_key_value",        # cogvlm
             "model.layers.{bid}.linear_attn.in_proj_qkv",                          # qwen3.5
+            "layers.{bid}.attention.wqkv",                                         # falcon-ocr
         ),
 
         # Attention query
@@ -357,6 +358,7 @@ class TensorNameMap:
         MODEL_TENSOR.ATTN_SINKS: (
             "model.layers.{bid}.self_attn.sinks", # openai-moe
             "model.layers.{bid}.self_attn.attention_sink_bias", # mimov2
+            "layers.{bid}.attention.sinks", # falcon_ocr
         ),
 
         MODEL_TENSOR.ATTN_GATE: (

src/CMakeLists.txt

@@ -67,6 +67,7 @@ add_library(llama
             models/exaone.cpp
             models/exaone4.cpp
             models/falcon-h1.cpp
+            models/falcon-ocr.cpp
             models/falcon.cpp
             models/gemma-embedding.cpp
             models/gemma.cpp

src/llama-arch.cpp

@@ -63,6 +63,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_MAMBA2,           "mamba2"           },
     { LLM_ARCH_JAMBA,            "jamba"            },
     { LLM_ARCH_FALCON_H1,        "falcon-h1"        },
+    { LLM_ARCH_FALCON_OCR,       "falcon-ocr"       },
     { LLM_ARCH_XVERSE,           "xverse"           },
     { LLM_ARCH_COMMAND_R,        "command-r"        },
     { LLM_ARCH_COHERE2,          "cohere2"          },

src/llama-arch.h

@@ -67,6 +67,7 @@ enum llm_arch {
     LLM_ARCH_MAMBA2,
     LLM_ARCH_JAMBA,
     LLM_ARCH_FALCON_H1,
+    LLM_ARCH_FALCON_OCR,
     LLM_ARCH_XVERSE,
     LLM_ARCH_COMMAND_R,
     LLM_ARCH_COHERE2,

src/llama-context.cpp

@@ -2070,7 +2070,7 @@ void llama_context::output_reorder() {
 //
 
 uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
-    if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) {
+    if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_FALCON_OCR) {
         return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
     }
     uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());

src/llama-graph.cpp

@@ -1201,8 +1201,9 @@ ggml_tensor * llm_graph_context::build_ffn(
 
     if (down) {
         cur = build_lora_mm(down, cur);
-        if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) {
+        if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2 || arch == LLM_ARCH_FALCON_OCR) {
             // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators
+            // Falcon-OCR's ReLU^2 activation produces values > 65504 (FP16 max), causing overflow.
             ggml_mul_mat_set_prec(cur, GGML_PREC_F32);
         }
     }

src/llama-kv-cache.cpp

@@ -1088,7 +1088,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
 
 bool llama_kv_cache::get_can_shift() const {
     // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.
-    if (model.arch == LLM_ARCH_STEP35) {
+    // falcon-ocr uses custom version of M-RoPE
+    if (model.arch == LLM_ARCH_STEP35 || model.arch == LLM_ARCH_FALCON_OCR) {
         return false;
     }
     if (hparams.n_pos_per_embd() > 1) {

src/llama-model.cpp

@@ -2936,6 +2936,15 @@ void llama_model::load_hparams(llama_model_loader & ml) {
                     default: type = LLM_TYPE_UNKNOWN;
                 }
             } break;
+        case LLM_ARCH_FALCON_OCR:
+            {
+                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
+                switch (hparams.n_layer) {
+                    case 22: type = LLM_TYPE_SMALL; break;
+                    default: type = LLM_TYPE_UNKNOWN;
+                }
+            } break;
         default: throw std::runtime_error("unsupported model architecture: " + arch_name());
     }
 
@@ -7969,6 +7978,34 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                         layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
 
                         layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+                        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
+                        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
+                        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+                    }
+                } break;
+            case LLM_ARCH_FALCON_OCR:
+                {
+                    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+                    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+                    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);
+
+                    // note: the model doesn't actually use GQA due to "golden" rope enforcing Q dimension
+                    const int64_t n_head_kv_ratio = 2;
+                    const int64_t n_embd_qkv = (n_embd_head_k * n_head)
+                                                + n_embd_k_gqa / n_head_kv_ratio
+                                                + n_embd_v_gqa / n_head_kv_ratio;
+
+                    for (int i = 0; i < n_layer; ++i) {
+                        auto & layer = layers[i];
+
+                        layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_head * n_rot/2, 2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+
+                        layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_qkv}, 0);
+                        layer.wo   = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+                        layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, i), {n_head}, TENSOR_NOT_REQUIRED);
+
                         layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
                         layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
                         layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
@@ -9262,6 +9299,10 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
             {
                 llm = std::make_unique<llm_build_step35_iswa>(*this, params);
             } break;
+        case LLM_ARCH_FALCON_OCR:
+            {
+                llm = std::make_unique<llm_build_falcon_ocr>(*this, params);
+            } break;
         default:
             GGML_ABORT("fatal error");
     }
@@ -9518,6 +9559,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
 
         case LLM_ARCH_QWEN2VL:
         case LLM_ARCH_PADDLEOCR:
+        case LLM_ARCH_FALCON_OCR: // note: falcon-ocr uses a variant of m-rope
             return LLAMA_ROPE_TYPE_MROPE;
         case LLM_ARCH_QWEN3VL:
         case LLM_ARCH_QWEN3VLMOE: