fix(qwen4_exp): fix nvfp4 ftw weight loading, qsa indexer, and ple tables

python/freetoken/engine/engine.py

@@ -459,7 +459,7 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso
         # _materialize casts each loaded tensor to its model-param dtype (model_state), so
         # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved;
         # offload models exclude experts (served from the offload cache, not dense weights).
-        return _materialize_loaded_weight_state_dict(
+        state_dict = _materialize_loaded_weight_state_dict(
             model_state,
             load_weight(
                 config.model_path,
@@ -468,6 +468,22 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso
             ),
             device=self.device,
         )
+        # Fuse any separate Hyper-Connection mix + block injection weights
+        keys = list(state_dict.keys())
+        for k in keys:
+            if k.endswith(".input_mix_weight_down.weight"):
+                prefix = k[:-len(".input_mix_weight_down.weight")]
+                inject_k = f"{prefix}.block_inject_weight.weight"
+                if inject_k in state_dict:
+                    w_down = state_dict.pop(k)
+                    w_inject = state_dict.pop(inject_k)
+                    pad = (-(w_down.shape[0] + w_inject.shape[0])) % 16
+                    parts = [w_down, w_inject]
+                    if pad:
+                        parts.append(torch.zeros(pad, *w_down.shape[1:], dtype=w_down.dtype, device=w_down.device))
+                    fused_name = f"{prefix}.input_mix_weight_down_block_inject.weight"
+                    state_dict[fused_name] = torch.cat(parts, dim=0)
+        return state_dict
 
     def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int, int, bool]:
         """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap).

python/freetoken/models/qwen4_exp/attention.py

@@ -88,6 +88,25 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None:
         self.q_layernorm = GemmaPlusOneRMSNorm(self.head_dim, eps=self.eps)
         self.k_layernorm = GemmaPlusOneRMSNorm(self.head_dim, eps=self.eps)
 
+    def load_state_dict(self, state_dict: dict, *, prefix: str = "", _internal: bool = False) -> None:
+        # Handles both self_attn.indexer.* and flat self_attn.index_* checkpoint layouts
+        parent_prefix = prefix[:-len(".indexer")] if prefix.endswith(".indexer") else prefix
+
+        for k in (f"{prefix}.index_qk_proj.weight", f"{parent_prefix}.index_qk_proj.weight"):
+            if k in state_dict:
+                self.index_qk_proj.weight = state_dict.pop(k)
+                break
+
+        for k in (f"{prefix}.q_layernorm.weight", f"{parent_prefix}.index_q_norm.weight", f"{parent_prefix}.q_layernorm.weight"):
+            if k in state_dict:
+                self.q_layernorm.weight = state_dict.pop(k)
+                break
+
+        for k in (f"{prefix}.k_layernorm.weight", f"{parent_prefix}.index_k_norm.weight", f"{parent_prefix}.k_layernorm.weight"):
+            if k in state_dict:
+                self.k_layernorm.weight = state_dict.pop(k)
+                break
+
     def forward(self, x: torch.Tensor) -> QSAIndexerInputs:
         q, k = self.index_qk_proj.forward(x).split(self._split, dim=-1)
         return QSAIndexerInputs(

python/freetoken/models/qwen4_exp/gdn.py

@@ -4,7 +4,7 @@
 import torch.nn.functional as F
 from freetoken.core import get_global_ctx
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
-from freetoken.layers import BaseOP, LinearColParallelMerged
+from freetoken.layers import BaseOP, LinearColParallelMerged, LinearReplicated
 
 from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged
 from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged
@@ -110,9 +110,7 @@ def __init__(
         # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors
         # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors
         # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4.
-        self.out_proj = make_replicated_quant(
-            expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False
-        )
+        self.out_proj = LinearReplicated(self.value_dim, hidden_size, has_bias=False)
 
     def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
         beta = b.sigmoid()

python/freetoken/models/qwen4_exp/model.py

@@ -14,6 +14,7 @@
 """
 
 from __future__ import annotations
+import os
 
 from typing import TYPE_CHECKING, List
 
@@ -150,7 +151,7 @@ def load_host_tables(self, engine_config) -> int:
             return 0
         from .ple import PinnedUVATable, ZeroTable, derive_ngram_hash_constants
 
-        if getattr(engine_config, "use_dummy_weight", False):
+        if getattr(engine_config, "use_dummy_weight", False) or os.environ.get("FREETOKEN_ZERO_PLE", "0") == "1":
             # Dummy fill leaves the int64 hash buffers garbage (a zero vocab size divides by
             # zero in the hash), so re-derive the real constants and read a zero table.
             for ple in ple_layers:
@@ -163,9 +164,9 @@ def load_host_tables(self, engine_config) -> int:
                     ple_layer_index=ple.ple_index,
                 )
                 emb = ple.ple_embedding
-                emb.layer_multipliers.copy_(torch.tensor(mult, dtype=torch.int64))
-                emb.ngram_heads_vocab_sizes.copy_(torch.tensor(sizes, dtype=torch.int64))
-                emb.ngram_heads_offsets.copy_(torch.tensor(offsets, dtype=torch.int64))
+                emb.layer_multipliers = torch.tensor(mult, dtype=torch.int64, device=torch.cuda.current_device())
+                emb.ngram_heads_vocab_sizes = torch.tensor(sizes, dtype=torch.int64, device=torch.cuda.current_device())
+                emb.ngram_heads_offsets = torch.tensor(offsets, dtype=torch.int64, device=torch.cuda.current_device())
                 emb.attach_table(ZeroTable(offsets[-1] + sizes[-1], args.ngram_head_dim))
             return 0
 

python/freetoken/models/qwen4_exp/ple.py

@@ -417,6 +417,14 @@ def __init__(self, args: Qwen4ExpArgs, table: PLETableBackend | None = None) ->
         self.ngram_heads_offsets = torch.empty(self.num_heads, dtype=torch.int64)
         self._table = table
 
+    def load_state_dict(self, state_dict: dict, *, prefix: str = "", _internal: bool = False) -> None:
+        # Pop metadata tensors if present in state_dict, otherwise keep computed buffers
+        for name in ("layer_multipliers", "ngram_heads_vocab_sizes", "ngram_heads_offsets"):
+            k = f"{prefix}.{name}" if prefix else name
+            if k in state_dict:
+                val = state_dict.pop(k)
+                getattr(self, name).copy_(val)
+
     def attach_table(self, table: PLETableBackend) -> None:
         self._table = table
 
@@ -466,6 +474,11 @@ def row_ids(self, meta: PLEMetadata) -> torch.Tensor:
         """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64."""
         packed, select = self._window(meta)
         tokens = [select(s) for s in self._shift_ignore_eos(packed)]
+        device = tokens[0].device
+        if not self.layer_multipliers.is_meta and self.layer_multipliers.device != device:
+            self.layer_multipliers = self.layer_multipliers.to(device)
+            self.ngram_heads_vocab_sizes = self.ngram_heads_vocab_sizes.to(device)
+            self.ngram_heads_offsets = self.ngram_heads_offsets.to(device)
         blocks = []
         for ngram in range(2, self.ngram_size + 1):
             start = (ngram - 2) * self.heads_per_ngram

python/freetoken/models/qwen4_exp/weight.py

@@ -232,6 +232,36 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True,
     table is ~47.7 GiB and must not also sit in the page cache while the bank holds the same bytes.
     """
     folder = download_hf_weight(model_path)
+    ngram_json_path = os.path.join(folder, "qwen4_ngram.json")
+    if os.path.exists(ngram_json_path):
+        with open(ngram_json_path) as f:
+            ngram_meta = json.load(f)
+        bin_path = os.path.join(folder, ngram_meta["file"])
+        rows = int(ngram_meta["rows"])
+        cols = int(ngram_meta["dim"])
+        nbytes = int(ngram_meta["nbytes"])
+        dtype_str = ngram_meta.get("dtype", "bfloat16")
+        dtype_map = {
+            "bfloat16": torch.bfloat16,
+            "float8_e4m3fn": torch.float8_e4m3fn,
+            "fp8": torch.float8_e4m3fn,
+            "float16": torch.float16,
+        }
+        dtype = dtype_map.get(dtype_str, torch.bfloat16)
+        scale = torch.tensor(1.0, dtype=torch.float32)
+
+        bank = HostBank((rows, cols), dtype)
+        bar = byte_bar(nbytes, "Loading PLE table (binary)")
+        try:
+            buf = bank.memoryview()
+            read_range_into(buf, bin_path, file_offset=0, nbytes=nbytes,
+                            dest_offset=0, workers=workers, chunk=chunk)
+            bar.update(nbytes)
+        finally:
+            bar.close()
+        if pin and torch.cuda.is_available():
+            bank.pin()
+        return PleTable(bank=bank, weight_scale=scale)
     parts: dict[int, tuple[str, int, int]] = {}  # shard index -> (path, file offset, bytes)
     scale: torch.Tensor | None = None
     rows = cols = 0