From 5fc94dd6add113df77dc8ef83904a2ddb901850e Mon Sep 17 00:00:00 2001 From: Clayton Rush Williams Date: Sun, 30 Aug 2026 11:01:14 -0400 Subject: [PATCH] fix(qwen4_exp): fix NVFP4 FTW weight loading, QSA indexer, and PLE tables --- python/freetoken/engine/engine.py | 18 ++++++++++- .../freetoken/models/qwen4_exp/attention.py | 19 ++++++++++++ python/freetoken/models/qwen4_exp/gdn.py | 6 ++-- python/freetoken/models/qwen4_exp/model.py | 9 +++--- python/freetoken/models/qwen4_exp/ple.py | 13 ++++++++ python/freetoken/models/qwen4_exp/weight.py | 30 +++++++++++++++++++ 6 files changed, 86 insertions(+), 9 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 73dc7688d..dedea92c8 100644 --- a/python/freetoken/engine/engine.py +++ b/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). diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index d6ab2867a..633e63110 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/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( diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 69838153f..be57b2fe9 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/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() diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index e5b365aa3..26dc2f2b2 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/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 diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 5100229a9..e2d842ccb 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/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 diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..4497b0c2d 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/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