diff --git a/docs/models.md b/docs/models.md index e4850a12..603759f2 100644 --- a/docs/models.md +++ b/docs/models.md @@ -11,6 +11,7 @@ for them; other checkpoints of the same architectures work too. | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) | | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) | | Qwen3.6 dense | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) | +| Qwen3.8 Flash Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8) (text-only; exact QSA prefix through 2,048 tokens; host-mapped PLE) | | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. | @@ -38,3 +39,13 @@ for them; other checkpoints of the same architectures work too. - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the authoritative model args are read from there. - Multimodal checkpoints are served text-only. +- Qwen3.8 Flash Next natively supports 262,144 tokens, but this integration + intentionally caps total sequence length at 2,048 tokens. Within that range, + dense causal attention is exactly equivalent to the checkpoint's QSA + selection because every visible token remains inside its 2,048-token budget. + Sparse QSA beyond that exact dense prefix is not implemented yet. +- Qwen3.8 Flash Next requires an offload-family MoE backend. FreeToken + automatically disables CUDA graphs and selects the naive cache because PLE + performs host-side gathers and owns per-request convolution state. The FP8 + checkpoint occupies about 173 GiB on disk; routed-expert banks use about + 113 GiB of pinned host RAM while the 48 GiB PLE table remains mmap-backed. diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539c..a654f925 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -125,6 +125,9 @@ class Batch: # flags), built once and shared by all GDN layers. Lazily built by the GDN op if # the scheduler/graph didn't set it. fla_metadata: "FLAMetadata | None" = field(default=None, init=False) + # True only for the static dummy batch while CUDA graphs are captured. Models + # with host-prepared inputs use this to bind stable device buffers into the graph. + cuda_graph_capture: bool = field(default=False, init=False) padded_reqs: List[Req] = field(init=False) # DSV4 paged-KV out-locations for this batch (None for non-DSV4 models). Set by the scheduler. # This decode batch's padded per-row page-table rows. Attention backends that must read diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..33619633 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -323,6 +323,11 @@ def __init__(self, config: EngineConfig): with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) self.model.load_state_dict(self._load_weight_state_dict(config)) + if hasattr(self.model, "load_host_weights"): + self.model.load_host_weights( + config.model_path, + dummy=config.use_dummy_weight, + ) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are @@ -1220,6 +1225,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution model_config = config.model_config single_stream_only = getattr(model_config, "single_stream_only", False) + requires_naive_cache = getattr(model_config, "requires_naive_cache", False) + supports_cuda_graph = getattr(model_config, "supports_cuda_graph", True) is_dsv4 = getattr(model_config, "dsv4_args", None) is not None has_swa_attention = getattr(model_config, "has_swa_attention", False) has_linear_attention = getattr(model_config, "has_linear_attention", False) @@ -1259,6 +1266,14 @@ def override(attr: str, value: Any): # this is dangerous, use with caution override("cuda_graph_bs", [1]) override("cuda_graph_max_bs", 1) + if not supports_cuda_graph: + override("cuda_graph_bs", []) + override("cuda_graph_max_bs", 0) + logger.info_rank0( + f"CUDA graphs disabled for {getattr(model_config, 'model_type', 'model')}: " + "the model requires host-side work during forward" + ) + if config.cuda_graph_max_bs is None: override("cuda_graph_max_bs", config.max_running_req) @@ -1281,6 +1296,14 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("cache_type", "swa_radix") + if requires_naive_cache and getattr(config, "cache_type", "radix") != "naive": + override("cache_type", "naive") + logger.warning_rank0( + f"Cache type overridden to 'naive' for " + f"{getattr(model_config, 'model_type', 'model')}: model-owned runtime state " + "cannot be restored from radix prefixes" + ) + if has_linear_attention: override( "cache_type", diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f202502..7fac4bd3 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -112,6 +112,7 @@ def __init__( free_memory=free_memory, ) self.attn_backend = attn_backend + self.model = model self.max_graph_bs = max(cuda_graph_bs) if cuda_graph_bs else 0 self.graph_bs_list = sorted(cuda_graph_bs) self.dummy_req = dummy_req @@ -164,6 +165,8 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel batch.padded_reqs = batch.reqs self.attn_backend.prepare_for_capture(batch) self.buffer.set_batch(batch) + batch.cuda_graph_capture = True + model.prepare_cuda_graph_capture(batch) # capture on the dummy linear-state slot so GatedDeltaNet gather/scatter # touches scratch (real slot indices are written by copy_from on replay). Hybrid- # radix decouples the GDN slot from table_idx -> use the GDN padding slot. @@ -191,6 +194,7 @@ def can_use_cuda_graph(self, batch: Batch) -> bool: def replay(self, batch: Batch) -> torch.Tensor: assert self.can_use_cuda_graph(batch) + self.model.prepare_cuda_graph_replay(batch) self.buffer.copy_from(batch) g = self.graph_map[batch.padded_size] self.attn_backend.prepare_for_replay(batch) diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index a00154d0..e79810b7 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -174,6 +174,24 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=512, expert_formats=("fp8_block",), ), + AotModel( + name="Qwen/Qwen3.8-Flash-Next-FP8", + architecture="Qwen4ExpForConditionalGeneration", + hidden_size=2560, + kv_groups=((2, 256),), + top_k=10, + moe_intermediate_size=640, + expert_formats=("fp8_block",), + ), + AotModel( + name="RadixArk/Qwen3.8-Flash-Next-NVFP4", + architecture="Qwen4ExpForConditionalGeneration", + hidden_size=2560, + kv_groups=((2, 256),), + top_k=10, + moe_intermediate_size=640, + expert_formats=_NVFP4_FORMATS, + ), AotModel( name="nvidia/Qwen3.6-35B-A3B-NVFP4", architecture="Qwen3_5MoeForConditionalGeneration", diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..8027f690 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -88,9 +88,9 @@ using mem_package_t = decltype(get_mem_package()); template __always_inline __device__ auto load_vec(const void* __restrict__ src) { using Package = mem_package_t; - constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; - constexpr auto kLoopCount = kBytes / kBytesPerLoop; - static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + static_assert(kBytes % sizeof(Package) == 0, "kBytes must contain whole memory packages"); + constexpr auto kPackageCount = kBytes / sizeof(Package); + constexpr auto kLoopCount = (kPackageCount + kThreads - 1) / kThreads; const auto src_packed = static_cast(src); const auto lane_id = threadIdx.x % kThreads; @@ -99,7 +99,9 @@ __always_inline __device__ auto load_vec(const void* __restrict__ src) { #pragma unroll kLoopCount for (std::size_t i = 0; i < kLoopCount; ++i) { const auto j = i * kThreads + lane_id; - vec.data[i] = load_nc(src_packed + j); + if (j < kPackageCount) { + vec.data[i] = load_nc(src_packed + j); + } } return vec; @@ -108,9 +110,9 @@ __always_inline __device__ auto load_vec(const void* __restrict__ src) { template __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) { using Package = mem_package_t; - constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; - constexpr auto kLoopCount = kBytes / kBytesPerLoop; - static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + static_assert(kBytes % sizeof(Package) == 0, "kBytes must contain whole memory packages"); + constexpr auto kPackageCount = kBytes / sizeof(Package); + constexpr auto kLoopCount = (kPackageCount + kThreads - 1) / kThreads; static_assert(std::is_same_v>); const auto dst_packed = static_cast(dst); @@ -119,7 +121,9 @@ __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) #pragma unroll kLoopCount for (std::size_t i = 0; i < kLoopCount; ++i) { const auto j = i * kThreads + lane_id; - details::store_nc(dst_packed + j, vec.data[i]); + if (j < kPackageCount) { + details::store_nc(dst_packed + j, vec.data[i]); + } } } diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 1aaa1303..108cc48e 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -61,7 +61,11 @@ def _jit_fast_index_copy_module( def _default_worker_threads(feature_size: int) -> int: if feature_size <= 1024: - return 8 + if feature_size % 16 == 0: + return 8 + if feature_size % 8 == 0: + return 16 + return 32 if feature_size <= 2048: return 16 return 32 diff --git a/python/freetoken/kernel/fla/layernorm_gated.py b/python/freetoken/kernel/fla/layernorm_gated.py index 55c2f00b..23d31a8c 100644 --- a/python/freetoken/kernel/fla/layernorm_gated.py +++ b/python/freetoken/kernel/fla/layernorm_gated.py @@ -47,6 +47,7 @@ def _layer_norm_fwd_1pass_kernel( NORM_BEFORE_GATE: tl.constexpr, IS_RMS_NORM: tl.constexpr, ACTIVATION: tl.constexpr, + WEIGHT_PLUS_ONE: tl.constexpr, ): # Map the program id to the starting row of X and Y it should compute. row_start = tl.program_id(0) * ROWS_PER_BLOCK @@ -106,6 +107,8 @@ def _layer_norm_fwd_1pass_kernel( w_offsets = cols + group * N w_mask = cols < N w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + if WEIGHT_PLUS_ONE: + w += 1.0 if HAS_BIAS: b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32) @@ -155,6 +158,7 @@ def _layer_norm_fwd( norm_before_gate=True, is_rms_norm=False, activation: str = "swish", + weight_plus_one: bool = False, ): M, N = x.shape if group_size is None: @@ -216,6 +220,7 @@ def _layer_norm_fwd( IS_RMS_NORM=is_rms_norm, num_warps=num_warps, ACTIVATION=activation, + WEIGHT_PLUS_ONE=weight_plus_one, ) return out, mean, rstd @@ -231,8 +236,13 @@ def rms_norm_gated( norm_before_gate=True, is_rms_norm=False, activation: str = "swish", + weight_plus_one: bool = False, ): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" + """Apply grouped (RMS)Norm and an optional gate. + + ``weight_plus_one`` keeps centered RMSNorm checkpoint weights unmodified and + performs the ``1 + weight`` operation in fp32 inside the kernel. + """ x_shape_og = x.shape # reshape input data into 2D tensor @@ -257,6 +267,7 @@ def rms_norm_gated( norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, activation=activation, + weight_plus_one=weight_plus_one, ) return y.reshape(x_shape_og) diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py index ca5b1972..3d106f33 100644 --- a/python/freetoken/models/blocks.py +++ b/python/freetoken/models/blocks.py @@ -16,6 +16,8 @@ if TYPE_CHECKING: import torch + from freetoken.core import Batch + from .config import ModelConfig @@ -23,6 +25,12 @@ class BaseLLMModel(ABC, BaseOP): @abstractmethod def forward(self) -> torch.Tensor: ... + def prepare_cuda_graph_capture(self, batch: Batch) -> None: + """Prepare model-owned stable inputs before capturing a decode graph.""" + + def prepare_cuda_graph_replay(self, batch: Batch) -> None: + """Stage model-owned dynamic inputs before replaying a decode graph.""" + class GatedMLP(BaseOP): def __init__(self, config: ModelConfig): diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f..e786d67a 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -288,9 +288,14 @@ class ModelConfig: # swigluoai/dense-MLP scalars the model module needs. Opaque to model-agnostic engine # code; None for every other model. m3_args: Any | None = None + # Qwen4-Exp payload: hyper-connections, PLE host embedding geometry, and the + # QSA exact-context ceiling. Opaque outside the qwen4_exp model package. + qwen4_args: Any | None = None # Generic execution-path capability flags (set by a model's parse_config) so the engine and # factories stay model-agnostic instead of branching on dsv4_args: single_stream_only: bool = False # model runs one sequence at a time -> force bs=1 + requires_naive_cache: bool = False # model owns host/runtime state radix cannot snapshot + supports_cuda_graph: bool = True # False when forward performs host-side dynamic work @property def is_moe(self) -> bool: diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py new file mode 100644 index 00000000..c4d5c731 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -0,0 +1,19 @@ +from .config import parse_config +from .model import Qwen4ExpForCausalLM +from .weight import ( + iter_weights, + iter_weights_parallel, + load_nvfp4_expert_sources, + load_nvfp4_expert_sources_parallel, + setup_offload_expert_banks, +) + +__all__ = [ + "Qwen4ExpForCausalLM", + "iter_weights", + "iter_weights_parallel", + "load_nvfp4_expert_sources", + "load_nvfp4_expert_sources_parallel", + "parse_config", + "setup_offload_expert_banks", +] diff --git a/python/freetoken/models/qwen4_exp/args.py b/python/freetoken/models/qwen4_exp/args.py new file mode 100644 index 00000000..3122dd48 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/args.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Qwen4ExpArgs: + hc_count: int + hc_lowrank: int + ple_layer_ids: tuple[int, ...] + ple_embed_dim: int + ple_conv_kernel_size: int + ngram_size: int + heads_per_ngram: int + ngram_vocab_size_base: int + split_ngram_parts: int + eos_token_id: int + indexer_budget: int + indexer_compress_ratio: int + output_gate_type: str + + +__all__ = ["Qwen4ExpArgs"] diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py new file mode 100644 index 00000000..9f5a7f7f --- /dev/null +++ b/python/freetoken/models/qwen4_exp/config.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from typing import Any + +from freetoken.models.config import ( + FullAttentionGroupConfig, + LinearGatedDeltaGroupConfig, + ModelConfig, + RotaryConfig, +) + +from .args import Qwen4ExpArgs + + +def parse_config(hf_config: Any) -> ModelConfig: + text = hf_config.text_config + layer_types = list(text.layer_types) + sparse_attention_types = {"full_attention", "qwen_sparse_attention"} + unsupported = sorted(set(layer_types) - {"linear_attention", *sparse_attention_types}) + if unsupported: + raise ValueError(f"Unsupported Qwen4-Exp layer types: {unsupported}") + + head_dim = int(text.head_dim) + rope = text.rope_parameters + rotary_dim = round(head_dim * float(rope.get("partial_rotary_factor", 1.0))) + indexer_budget = int(text.indexer_budget) + rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=rotary_dim, + # QSA selects every visible token through this point. FreeToken currently + # serves that exact dense prefix and rejects longer requests. + max_position=min(int(text.max_position_embeddings), indexer_budget), + base=float(rope["rope_theta"]), + scaling=None, + ) + + full_ids = tuple( + i for i, layer_type in enumerate(layer_types) if layer_type in sparse_attention_types + ) + linear_ids = tuple( + i for i, layer_type in enumerate(layer_types) if layer_type == "linear_attention" + ) + groups = ( + LinearGatedDeltaGroupConfig( + name="linear", + layer_ids=linear_ids, + num_key_heads=int(text.linear_num_key_heads), + num_value_heads=int(text.linear_num_value_heads), + key_head_dim=int(text.linear_key_head_dim), + value_head_dim=int(text.linear_value_head_dim), + conv_kernel_dim=int(text.linear_conv_kernel_dim), + output_gate=True, + ), + FullAttentionGroupConfig( + name="full", + layer_ids=full_ids, + num_kv_heads=int(text.num_key_value_heads), + head_dim=head_dim, + rotary_config=rotary, + ), + ) + + eos_token_id = text.eos_token_id + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] + qwen4_args = Qwen4ExpArgs( + hc_count=int(text.hc_count), + hc_lowrank=int(text.hc_lowrank), + ple_layer_ids=tuple(int(layer_id) - 1 for layer_id in text.ple_layer_ids), + ple_embed_dim=int(text.ple_embed_dim), + ple_conv_kernel_size=int(text.ple_conv_kernel_size), + ngram_size=int(text.ngram_size), + heads_per_ngram=int(text.heads_per_ngram), + ngram_vocab_size_base=int(text.ngram_vocab_size_base), + split_ngram_parts=int(text.split_ngram_parts), + eos_token_id=int(eos_token_id), + indexer_budget=indexer_budget, + indexer_compress_ratio=int(text.indexer_compress_ratio), + output_gate_type=str(text.output_gate_type or text.hidden_act), + ) + + quant = hf_config.quantization_config + if not isinstance(quant, dict): + quant = quant.to_dict() + method = str(quant.get("quant_method") or "").lower() + algo = str(quant.get("quant_algo") or method).lower() + if method == "fp8": + block_size = tuple(int(value) for value in quant["weight_block_size"]) + if block_size != (128, 128): + raise ValueError(f"Qwen4-Exp only supports 128x128 block-FP8, got {block_size}") + expert_quant = "fp8_block" + elif "fp4" in algo: + # ModelOpt NVFP4 checkpoints use packed routed experts while leaving the + # shared expert and other resident text weights in BF16. + block_size = None + expert_quant = "nvfp4" + else: + raise ValueError( + "Qwen4-Exp requires a 128x128 block-FP8 or ModelOpt NVFP4 checkpoint, " + f"got quant_method={method!r}, quant_algo={algo!r}" + ) + + return ModelConfig( + num_layers=int(text.num_hidden_layers), + num_qo_heads=int(text.num_attention_heads), + num_kv_heads=int(text.num_key_value_heads), + head_dim=head_dim, + hidden_size=int(text.hidden_size), + vocab_size=int(text.vocab_size), + intermediate_size=int(getattr(text, "intermediate_size", 0) or 0), + hidden_act=str(text.hidden_act), + rms_norm_eps=float(text.rms_norm_eps), + tie_word_embeddings=bool(getattr(text, "tie_word_embeddings", False)), + rotary_config=rotary, + num_experts=int(text.num_experts), + num_experts_per_tok=int(text.num_experts_per_tok), + moe_intermediate_size=int(text.moe_intermediate_size), + shared_expert_intermediate_size=int(text.shared_expert_intermediate_size), + norm_topk_prob=bool(getattr(text, "norm_topk_prob", True)), + model_type=str(hf_config.model_type), + architectures=list(hf_config.architectures), + moe_enabled=True, + expert_quant=expert_quant, + weight_block_size=block_size, + # Only routed experts and PLE are quantized in the supported checkpoints. + # Attention, hyper-connections, and shared-expert projections stay BF16. + attn_quant="none", + dense_quant="none", + lm_head_quant="none", + use_qk_norm=True, + vision_config=None, + image_token_id=getattr(hf_config, "image_token_id", None), + attention_groups=groups, + qwen4_args=qwen4_args, + # PLE keeps per-request dilated-convolution state outside the generic + # radix cache, so prefix snapshots remain unsupported. Its host embedding + # lookup is staged into a stable device buffer before CUDA-graph replay. + requires_naive_cache=True, + supports_cuda_graph=True, + ) + + +__all__ = ["parse_config"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py new file mode 100644 index 00000000..21ae7019 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/model.py @@ -0,0 +1,629 @@ +from __future__ import annotations + +import json +import math +import os +from dataclasses import replace +from typing import TYPE_CHECKING + +import safetensors +import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.layers import ( + BaseOP, + GemmaPlusOneRMSNorm, + LinearColParallelMerged, + LinearReplicated, + LinearRowParallel, + OPList, + ParallelLMHead, + VocabParallelEmbedding, + make_moe_layer, + silu_and_mul, +) +from freetoken.models.blocks import BaseLLMModel +from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention +from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet +from freetoken.utils import download_hf_weight, init_logger, nvtx_annotate + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from freetoken.core import Batch + from freetoken.models.config import ModelConfig + + from .args import Qwen4ExpArgs + + +class _GroupedRMSNorm(BaseOP): + def __init__(self, size: int, group_size: int, eps: float): + if size % group_size: + raise ValueError(f"RMSNorm size {size} is not divisible by group size {group_size}") + self.weight = torch.empty(size) + self.group_size = group_size + self.eps = eps + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.fla import rms_norm_gated + + return rms_norm_gated( + x=hidden, + weight=self.weight, + bias=None, + eps=self.eps, + group_size=self.group_size, + is_rms_norm=True, + weight_plus_one=True, + ) + + +class _GatedRMSNorm(BaseOP): + def __init__(self, size: int, eps: float, activation: str): + self.weight = torch.empty(size) + self.eps = eps + self.activation = activation + + def forward(self, hidden: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.fla import rms_norm_gated + + return rms_norm_gated( + x=hidden, + weight=self.weight, + bias=None, + z=gate, + eps=self.eps, + is_rms_norm=True, + norm_before_gate=True, + activation=self.activation, + ) + + +class _GatedResidual(BaseOP): + def __init__(self, config: ModelConfig, combine: bool = True): + args: Qwen4ExpArgs = config.qwen4_args + self.hc_count = args.hc_count + self.hidden_size = config.hidden_size + hc_size = self.hc_count * self.hidden_size + self.hc_norm = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.input_mix_weight_down = LinearReplicated(hc_size, args.hc_lowrank, has_bias=False) + self.input_mix_weight_up = LinearReplicated(args.hc_lowrank, hc_size, has_bias=False) + self.block_inject_weight = ( + LinearReplicated(hc_size, self.hc_count, has_bias=False) if combine else None + ) + + def forward(self, hyper_input: torch.Tensor): + normalized = self.hc_norm.forward(hyper_input) + mix = F.silu(self.input_mix_weight_down.forward(normalized) / self.hc_count) + mix = torch.sigmoid(self.input_mix_weight_up.forward(mix)) + mix = mix.view(-1, self.hc_count, self.hidden_size) + mixed = (mix * normalized.view(-1, self.hc_count, self.hidden_size)).mean(dim=1) + if self.block_inject_weight is None: + return mixed + inject = 2 * torch.sigmoid(self.block_inject_weight.forward(normalized) / self.hc_count) + return mixed, hyper_input, inject + + +class _SharedExpert(BaseOP): + def __init__(self, config: ModelConfig): + width = config.shared_expert_intermediate_size + self.gate_up_proj = LinearColParallelMerged( + config.hidden_size, [width, width], has_bias=False + ) + self.down_proj = LinearRowParallel(width, config.hidden_size, has_bias=False) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(hidden))) + + +class _SparseMoE(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.experts = make_moe_layer( + config, + layer_id=layer_id, + renormalize=bool(config.norm_topk_prob), + weight_format="fp8_block", + ) + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.shared_expert = _SharedExpert(config) + self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + router_logits = self.gate.forward(hidden) + shared = self.shared_expert.forward(hidden) + shared *= torch.sigmoid(self.shared_expert_gate.forward(hidden)) + return self.experts.forward(hidden_states=hidden, router_logits=router_logits) + shared + + +def _shift_right_ignore_eos(tokens: torch.Tensor, shift: int, eos_token_id: int) -> torch.Tensor: + if shift == 0: + return tokens + positions = torch.arange(tokens.numel(), dtype=torch.long) + eos_positions = torch.where(tokens == eos_token_id, positions, -1) + previous_eos_inclusive = torch.cummax(eos_positions, dim=0).values + previous_eos = torch.cat([eos_positions.new_full((1,), -1), previous_eos_inclusive[:-1]]) + segment_start = previous_eos + 1 + source_positions = positions - shift + shifted = tokens[source_positions.clamp_min(0)] + valid = (positions - segment_start >= shift) & (source_positions >= 0) + return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) + + +def build_ngram_ids( + tokens: torch.Tensor, + *, + ngram_size: int, + heads_per_ngram: int, + eos_token_id: int, + multipliers: torch.Tensor, + vocab_sizes: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + tokens = tokens.to(dtype=torch.long, device="cpu") + shifted = [ + _shift_right_ignore_eos(tokens, shift, eos_token_id) for shift in range(ngram_size) + ] + blocks = [] + for ngram in range(2, ngram_size + 1): + start = (ngram - 2) * heads_per_ngram + stop = start + heads_per_ngram + mixed = shifted[0] * multipliers[0] + for position in range(1, ngram): + mixed = torch.bitwise_xor(mixed, shifted[position] * multipliers[position]) + sizes = vocab_sizes[start:stop] + heads = torch.remainder(mixed.unsqueeze(-1), sizes) + blocks.append(heads + offsets[start:stop]) + return torch.cat(blocks, dim=-1) + + +def _tokens_for_ngram_forward( + req, + current_ids: torch.Tensor, + *, + start: int = 0, +) -> torch.Tensor: + """Return request history through ``device_len`` without mutating the request. + + Under overlap scheduling the next decode starts before the previous sampled token is + appended to ``req.input_ids``. That token is already present in ``batch.input_ids``; + splice only the missing suffix from there so PLE sees the same history in overlap and + non-overlap modes. ``start`` lets decode retain only the short suffix needed by the + n-gram hash instead of copying and re-hashing the complete request on every token. + """ + if not 0 <= start <= req.device_len: + raise ValueError(f"Qwen4-Exp PLE history start {start} is outside [0, {req.device_len}]") + host_len = min(req.input_ids.numel(), req.device_len) + host_start = min(start, host_len) + tokens = req.input_ids[host_start:host_len].to(dtype=torch.long, device="cpu") + missing = req.device_len - host_len + if missing: + if missing > current_ids.numel(): + raise RuntimeError( + f"Qwen4-Exp PLE history is missing {missing} tokens, " + f"but the forward only carries {current_ids.numel()}" + ) + inflight = current_ids[-missing:] + if start > host_len: + inflight = inflight[start - host_len :] + tokens = torch.cat([tokens, inflight.to(dtype=torch.long, device="cpu")]) + return tokens + + +def _preload_ple_enabled() -> bool: + return os.getenv("FREETOKEN_QWEN4_PLE_PRELOAD", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +class _HostNGramEmbedding(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + args: Qwen4ExpArgs = config.qwen4_args + self.layer_id = layer_id + self.ngram_size = args.ngram_size + self.heads_per_ngram = args.heads_per_ngram + self.eos_token_id = args.eos_token_id + self.embedding_dim = args.ple_embed_dim + self.split_ngram_parts = args.split_ngram_parts + self.ngram_heads = (args.ngram_size - 1) * args.heads_per_ngram + self.head_dim = self.embedding_dim // self.ngram_heads + self.layer_multipliers = torch.empty(args.ngram_size, dtype=torch.long) + self.ngram_heads_vocab_sizes = torch.empty(self.ngram_heads, dtype=torch.long) + self.ngram_heads_offsets = torch.empty(self.ngram_heads, dtype=torch.long) + self._handles = [] + self._shards: list[torch.Tensor] = [] + self._shard_ends = torch.empty(0, dtype=torch.long) + self._scale = torch.tensor(1.0, dtype=torch.bfloat16) + self._host_constants: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + self._dummy = False + self._graph_output: torch.Tensor | None = None + self._device_scale: torch.Tensor | None = None + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + if dummy: + self._dummy = True + return + folder = download_hf_weight(model_path) + index_path = os.path.join(folder, "model.safetensors.index.json") + with open(index_path) as index_file: + weight_map = json.load(index_file)["weight_map"] + prefix = ( + f"model.language_model.layers.{self.layer_id}.ple.ple_embedding." + "ngram_embedding" + ) + shard_count = len([key for key in weight_map if key.startswith(prefix + ".shard_")]) + if shard_count != self.split_ngram_parts: + raise RuntimeError( + f"Qwen4-Exp PLE has {shard_count} shards, expected {self.split_ngram_parts}" + ) + shard_keys = [f"{prefix}.shard_{shard_id}.weight" for shard_id in range(shard_count)] + if not shard_keys or any(key not in weight_map for key in shard_keys): + raise RuntimeError(f"Incomplete Qwen4-Exp PLE shards under {prefix}") + + handles = {} + shards = [] + for key in shard_keys: + filename = weight_map[key] + handle = handles.get(filename) + if handle is None: + handle = safetensors.safe_open( + os.path.join(folder, filename), framework="pt", device="cpu" + ).__enter__() + handles[filename] = handle + shard = handle.get_tensor(key) + if shard.dtype != torch.float8_e4m3fn or shard.shape[1] != self.head_dim: + raise RuntimeError(f"Unexpected PLE shard {key}: {shard.dtype} {tuple(shard.shape)}") + shards.append(shard.view(torch.uint8)) + scale_key = prefix + ".weight_scale" + scale_handle = handles.get(weight_map[scale_key]) + if scale_handle is None: + scale_handle = safetensors.safe_open( + os.path.join(folder, weight_map[scale_key]), framework="pt", device="cpu" + ).__enter__() + handles[weight_map[scale_key]] = scale_handle + + preload = _preload_ple_enabled() + if preload: + total_gib = sum(shard.numel() * shard.element_size() for shard in shards) / (1 << 30) + logger.info_rank0( + f"Preloading {total_gib:.2f} GiB of Qwen4-Exp PLE into host RAM " + "(FREETOKEN_QWEN4_PLE_PRELOAD=1)" + ) + shards = [shard.clone() for shard in shards] + logger.info_rank0("Qwen4-Exp PLE preload complete") + self._handles = [] if preload else list(handles.values()) + self._shards = shards + self._shard_ends = torch.tensor([shard.shape[0] for shard in shards]).cumsum(0) + self._scale = scale_handle.get_tensor(scale_key).reshape(()).clone() + if preload: + for handle in handles.values(): + handle.__exit__(None, None, None) + self._host_constants = ( + self.layer_multipliers.cpu(), + self.ngram_heads_vocab_sizes.cpu(), + self.ngram_heads_offsets.cpu(), + ) + expected_rows = int(self._host_constants[1][-1] + self._host_constants[2][-1]) + if int(self._shard_ends[-1]) < expected_rows: + raise RuntimeError( + f"PLE table has {int(self._shard_ends[-1])} rows, needs {expected_rows}" + ) + + def _current_ngram_ids(self) -> torch.Tensor: + if self._host_constants is None: + raise RuntimeError("Qwen4-Exp PLE host weights are not loaded") + batch = get_global_ctx().batch + reqs = batch.padded_reqs if batch.is_decode else batch.reqs + multipliers, vocab_sizes, offsets = self._host_constants + current_ids = ( + batch.input_ids.to(dtype=torch.long, device="cpu") if batch.is_decode else None + ) + pieces = [] + token_offset = 0 + for req in reqs: + length = req.extend_len + history_start = max(0, req.cached_len - (self.ngram_size - 1)) + if current_ids is not None: + tokens = _tokens_for_ngram_forward( + req, + current_ids[token_offset : token_offset + length], + start=history_start, + ) + token_offset += length + else: + tokens = req.input_ids[history_start : req.device_len] + all_ids = build_ngram_ids( + tokens, + ngram_size=self.ngram_size, + heads_per_ngram=self.heads_per_ngram, + eos_token_id=self.eos_token_id, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + ) + pieces.append( + all_ids[ + req.cached_len - history_start : req.device_len - history_start + ] + ) + if current_ids is not None and token_offset != current_ids.numel(): + raise RuntimeError( + f"Qwen4-Exp PLE consumed {token_offset} decode tokens, " + f"but the batch carries {current_ids.numel()}" + ) + result = torch.cat(pieces, dim=0) + if result.shape[0] != batch.input_ids.numel(): + raise RuntimeError( + f"PLE token count {result.shape[0]} does not match batch {batch.input_ids.numel()}" + ) + return result + + def _lookup(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if self._dummy: + token_count = get_global_ctx().batch.input_ids.numel() + return torch.zeros(token_count, self.embedding_dim, device=device, dtype=dtype) + ngram_ids = self._current_ngram_ids().reshape(-1) + shard_ids = torch.bucketize(ngram_ids, self._shard_ends, right=True) + output = torch.empty( + ngram_ids.numel(), + self.head_dim, + dtype=torch.uint8, + pin_memory=torch.cuda.is_available(), + ) + starts = torch.cat([self._shard_ends.new_zeros(1), self._shard_ends[:-1]]) + for shard_id in shard_ids.unique().tolist(): + positions = torch.nonzero(shard_ids == shard_id, as_tuple=False).flatten() + local_ids = ngram_ids.index_select(0, positions) - starts[shard_id] + rows = self._shards[shard_id].index_select(0, local_ids) + output.index_copy_(0, positions, rows) + fp8 = output.to(device=device, non_blocking=True).view(torch.float8_e4m3fn) + if ( + self._device_scale is None + or self._device_scale.device != device + or self._device_scale.dtype != dtype + ): + self._device_scale = self._scale.to(device=device, dtype=dtype) + embedded = fp8.to(dtype) * self._device_scale + return embedded.view(-1, self.embedding_dim) + + def prepare_cuda_graph_capture( + self, token_count: int, device: torch.device, dtype: torch.dtype + ) -> None: + if self._graph_output is None or self._graph_output.shape[0] < token_count: + self._graph_output = torch.zeros( + token_count, self.embedding_dim, device=device, dtype=dtype + ) + + def prepare_cuda_graph_replay(self, device: torch.device, dtype: torch.dtype) -> None: + assert self._graph_output is not None + embedded = self._lookup(device, dtype) + self._graph_output[: embedded.shape[0]].copy_(embedded) + + def forward(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + batch = get_global_ctx().batch + if batch.cuda_graph_capture: + assert self._graph_output is not None + return self._graph_output[: batch.input_ids.numel()] + return self._lookup(device, dtype) + + +class _DepthwiseConv(BaseOP): + def __init__(self, channels: int, kernel_size: int): + self.weight = torch.empty(channels, 1, kernel_size) + + +class _PLELayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + args: Qwen4ExpArgs = config.qwen4_args + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.hc_count = args.hc_count + hc_size = self.hidden_size * self.hc_count + self.ple_embedding = _HostNGramEmbedding(config, layer_id) + self.key_proj = LinearReplicated(args.ple_embed_dim, hc_size, has_bias=False) + self.value_proj = LinearReplicated(args.ple_embed_dim, self.hidden_size, has_bias=False) + self.norm_key = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.norm_query = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.norm_conv = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.conv1d = _DepthwiseConv(hc_size, args.ple_conv_kernel_size) + self.dilation = args.ngram_size + self.state_len = (args.ple_conv_kernel_size - 1) * self.dilation + self._conv_state_pool: torch.Tensor | None = None + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + self.ple_embedding.load_host_weights(model_path, dummy=dummy) + + def _ensure_conv_state_pool(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + linear_pool = get_global_ctx().linear_state_pool + if linear_pool is None: + raise RuntimeError("Qwen4-Exp PLE requires a linear-state pool") + num_slots = linear_pool.conv_states.shape[1] + expected = (num_slots, self.hc_count * self.hidden_size, self.state_len) + if ( + self._conv_state_pool is None + or self._conv_state_pool.shape != expected + or self._conv_state_pool.device != device + or self._conv_state_pool.dtype != dtype + ): + self._conv_state_pool = torch.zeros(expected, device=device, dtype=dtype) + return self._conv_state_pool + + def prepare_cuda_graph_capture(self, token_count: int) -> None: + device = self.conv1d.weight.device + dtype = self.conv1d.weight.dtype + self.ple_embedding.prepare_cuda_graph_capture(token_count, device, dtype) + self._ensure_conv_state_pool(device, dtype) + + def prepare_cuda_graph_replay(self) -> None: + self.ple_embedding.prepare_cuda_graph_replay( + self.conv1d.weight.device, self.conv1d.weight.dtype + ) + + def _short_conv(self, hidden: torch.Tensor) -> torch.Tensor: + batch = get_global_ctx().batch + reqs = batch.padded_reqs if batch.is_decode else batch.reqs + state_pool = self._ensure_conv_state_pool(hidden.device, hidden.dtype) + if batch.cuda_graph_capture: + assert batch.linear_table_idx is not None + slots = batch.linear_table_idx.long() + state = state_pool.index_select(0, slots) + combined = torch.cat([state, hidden.unsqueeze(-1)], dim=-1) + convolved = F.conv1d( + combined, + self.conv1d.weight, + groups=self.conv1d.weight.shape[0], + dilation=self.dilation, + ) + state_pool.index_copy_(0, slots, combined[..., -self.state_len :]) + return F.silu(convolved).squeeze(-1) + + outputs = [] + offset = 0 + weight = self.conv1d.weight + for req in reqs: + length = req.extend_len + current = hidden[offset : offset + length].transpose(0, 1).unsqueeze(0) + state = state_pool[req.table_idx].unsqueeze(0) + if req.cached_len == 0: + state.zero_() + combined = torch.cat([state, current], dim=-1) + convolved = F.conv1d( + combined, + weight, + groups=weight.shape[0], + dilation=self.dilation, + ) + outputs.append(F.silu(convolved).squeeze(0).transpose(0, 1)) + state_pool[req.table_idx].copy_(combined[0, :, -self.state_len :]) + offset += length + return torch.cat(outputs, dim=0) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + embeddings = self.ple_embedding.forward(hidden.device, hidden.dtype) + key = self.norm_key.forward(self.key_proj.forward(embeddings)) + key = key.view(-1, self.hc_count, self.hidden_size) + value = self.value_proj.forward(embeddings) + query = self.norm_query.forward(hidden).view(-1, self.hc_count, self.hidden_size) + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(1)).flatten(1) + normalized = self.norm_conv.forward(gated) + return gated + self._short_conv(normalized) + + +class Qwen4ExpDecoderLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self._layer_id = layer_id + self._is_linear = config.is_linear_layer(layer_id) + dense_config = replace(config, expert_quant="none", attn_quant="none") + if self._is_linear: + group = config.linear_attention_group() + assert group is not None + self.linear_attn = Qwen3_5GatedDeltaNet( + hidden_size=config.hidden_size, + num_k_heads=group.num_key_heads, + num_v_heads=group.num_value_heads, + head_k_dim=group.key_head_dim, + head_v_dim=group.value_head_dim, + conv_kernel_size=group.conv_kernel_dim, + rms_norm_eps=config.rms_norm_eps, + layer_id=layer_id, + expert_quant="none", + attn_quant="none", + ) + self.linear_attn.norm = _GatedRMSNorm( + group.value_head_dim, + config.rms_norm_eps, + config.qwen4_args.output_gate_type, + ) + else: + self.self_attn = Qwen3_5Attention(dense_config, layer_id) + # Qwen4 stores centered q/k norm weights (effective scale is 1 + w). + # Keep the raw checkpoint values and add one in fp32 inside the kernel. + head_dim = config.head_dim + self.self_attn.q_norm = GemmaPlusOneRMSNorm(head_dim, config.rms_norm_eps) + self.self_attn.k_norm = GemmaPlusOneRMSNorm(head_dim, config.rms_norm_eps) + self.mlp = _SparseMoE(config, layer_id) + self.ple = ( + _PLELayer(config, layer_id) + if layer_id in config.qwen4_args.ple_layer_ids + else None + ) + self.attn_hyper_connection = _GatedResidual(config) + self.mlp_hyper_connection = _GatedResidual(config) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if self.ple is not None: + hidden = hidden + self.ple.forward(hidden) + mixed, residual, weights = self.attn_hyper_connection.forward(hidden) + mixed = ( + self.linear_attn.forward(mixed) + if self._is_linear + else self.self_attn.forward(mixed) + ) + hidden = residual + (mixed.unsqueeze(1) * weights.unsqueeze(-1)).flatten(1) + mixed, residual, weights = self.mlp_hyper_connection.forward(hidden) + mixed = self.mlp.forward(mixed) + return residual + (mixed.unsqueeze(1) * weights.unsqueeze(-1)).flatten(1) + + +class Qwen4ExpModel(BaseOP): + def __init__(self, config: ModelConfig): + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = OPList( + [Qwen4ExpDecoderLayer(config, layer_id) for layer_id in range(config.num_layers)] + ) + self.hyper_connection_mixer = _GatedResidual(config, combine=False) + self.hc_count = config.qwen4_args.hc_count + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + for layer in self.layers.op_list: + if layer.ple is not None: + layer.ple.load_host_weights(model_path, dummy=dummy) + + def prepare_cuda_graph_capture(self, token_count: int) -> None: + for layer in self.layers.op_list: + if layer.ple is not None: + layer.ple.prepare_cuda_graph_capture(token_count) + + def prepare_cuda_graph_replay(self) -> None: + for layer in self.layers.op_list: + if layer.ple is not None: + layer.ple.prepare_cuda_graph_replay() + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + hidden = self.embed_tokens.forward(input_ids).repeat(1, self.hc_count) + for layer in self.layers.op_list: + hidden = layer.forward(hidden) + return self.hyper_connection_mixer.forward(hidden) + + +class Qwen4ExpForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self.model = Qwen4ExpModel(config) + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + ) + super().__init__() + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + self.model.load_host_weights(model_path, dummy=dummy) + + def prepare_cuda_graph_capture(self, batch: Batch) -> None: + self.model.prepare_cuda_graph_capture(batch.input_ids.numel()) + + def prepare_cuda_graph_replay(self, batch: Batch) -> None: + self.model.prepare_cuda_graph_replay() + + def forward(self) -> torch.Tensor: + hidden = self.model.forward(get_global_ctx().batch.input_ids) + return self.lm_head.forward(hidden) + + +__all__ = ["Qwen4ExpForCausalLM", "build_ngram_ids"] diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py new file mode 100644 index 00000000..561c11fb --- /dev/null +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from typing import Iterator + +import safetensors +import torch +from freetoken.distributed import get_tp_info +from freetoken.models.loader import iter_weight_files +from tqdm import tqdm + +from freetoken.models.qwen3_5_moe.weight import ( + iter_weights_parallel, + load_nvfp4_expert_sources, + load_nvfp4_expert_sources_parallel, + setup_offload_expert_banks, +) + + +_FUSIONS = { + ".self_attn.qkv_proj.weight": ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ), + ".linear_attn.in_proj.weight": ( + ".linear_attn.in_proj_qkv.weight", + ".linear_attn.in_proj_z.weight", + ".linear_attn.in_proj_b.weight", + ".linear_attn.in_proj_a.weight", + ), + ".mlp.shared_expert.gate_up_proj.weight": ( + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ), +} + + +def _rename(raw_name: str) -> str | None: + if raw_name.startswith(("mtp.", "model.visual.", "visual.")): + return None + if ".self_attn.indexer." in raw_name: + return None + if ".ple.ple_embedding.ngram_embedding." in raw_name: + return None + if raw_name.startswith("model.language_model."): + return "model." + raw_name[len("model.language_model.") :] + if raw_name.startswith("language_model."): + return "model." + raw_name[len("language_model.") :] + return raw_name + + +def _try_fuse(name: str, tensor: torch.Tensor, buffers: dict): + for fused_suffix, parts in _FUSIONS.items(): + for index, part in enumerate(parts): + if name.endswith(part): + fused_name = name[: -len(part)] + fused_suffix + slots = buffers.setdefault(fused_name, {}) + slots[index] = tensor + if len(slots) == len(parts): + del buffers[fused_name] + return fused_name, torch.cat([slots[i] for i in range(len(parts))], dim=0) + return () + return None + + +def iter_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + if get_tp_info().size > 1: + raise NotImplementedError("Qwen4-Exp currently supports TP=1 only") + if include_moe_experts: + raise ValueError("Qwen4-Exp requires --moe-backend offload, cpu, or hybrid") + if not include_non_moe: + return + + buffers = {} + for filename in tqdm( + iter_weight_files(model_path), + desc="Loading Qwen4-Exp resident weights", + disable=not get_tp_info().is_primary(), + ): + with safetensors.safe_open(filename, framework="pt", device=str(device)) as handle: + for raw_name in handle.keys(): + name = _rename(raw_name) + if ( + name is None + or ".mlp.experts." in name + or raw_name.endswith(".weight_scale_inv") + ): + continue + tensor = handle.get_tensor(raw_name) + fused = _try_fuse(name, tensor, buffers) + if fused is not None: + if fused: + yield fused + continue + yield name, tensor + if buffers: + raise RuntimeError(f"Incomplete Qwen4-Exp projection fusions: {sorted(buffers)}") + + +__all__ = [ + "iter_weights", + "iter_weights_parallel", + "load_nvfp4_expert_sources", + "load_nvfp4_expert_sources_parallel", + "setup_offload_expert_banks", +] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..9626261a 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -65,6 +65,12 @@ class ModelSpec: "freetoken.models.qwen3_5_moe", "Qwen3_5MoEForCausalLM", ), + # Qwen3.8-Flash-Next / Qwen4-Exp text tower. Routed FP8 experts use the + # standard offload banks; its 51 GB PLE table is mmap-backed on the host. + "Qwen4ExpForConditionalGeneration": ModelSpec( + "freetoken.models.qwen4_exp", + "Qwen4ExpForCausalLM", + ), # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. Dense gated GQA # with a [SWA x3, full] pattern -- full layers are NoPE -- weightless qk norms, centered diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index fe7e417d..736bc91f 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -46,19 +46,30 @@ def fused_topk( from freetoken.kernel.backend import is_triton_kernels_installed + fallback_reason = None + kernel_topk = topk + if not is_triton_kernels_installed(): + fallback_reason = "triton_kernels is not installed" + elif topk & (topk - 1): + # triton_kernels.topk uses tl.arange(0, k), which requires a power-of-two + # range. Select the next power of two and discard the extra tail instead + # of routing Qwen's top_k=10 through the much slower multi-kernel fallback. + kernel_topk = 1 << (topk - 1).bit_length() + if kernel_topk > gating_output.shape[-1]: + fallback_reason = ( + f"triton_kernels padded topk={kernel_topk} exceeds " + f"num_experts={gating_output.shape[-1]}" + ) + # triton_kernels ships no Windows wheel, and unlike flashinfer/sgl_kernel it is not one # of the six ops the in-repo triton kernels cover -- so this router needs its own fallback. - if not is_triton_kernels_installed(): + if fallback_reason is not None: global _warned_torch_topk if not _warned_torch_topk: _warned_torch_topk = True - # Once, not per call: this runs every MoE forward. On Linux a missing - # triton_kernels used to fail fast with ImportError; keep the misconfiguration - # visible without giving up the fallback that Windows needs. logger.warning_rank0( - "fused_topk: triton_kernels is not installed -> pure-torch router fallback " - "(numerically equivalent, slower). Expected on Windows (no wheel); on Linux " - "install triton_kernels to restore the fused router." + f"fused_topk: {fallback_reason} -> pure-torch router fallback " + "(numerically equivalent, slower)." ) return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded) @@ -70,7 +81,7 @@ def fused_topk( logits = torch.softmax(logits, dim=-1) sparse_topk = triton_kernels_topk( logits, - topk, + kernel_topk, apply_softmax=not softmax_first, ) if hasattr(sparse_topk, "vals"): @@ -78,6 +89,12 @@ def fused_topk( topk_ids = sparse_topk.indx else: topk_weights, topk_ids = sparse_topk[:2] + if kernel_topk != topk: + topk_weights = topk_weights[:, :topk] + topk_ids = topk_ids[:, :topk] + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + topk_weights = topk_weights.contiguous() topk_ids = topk_ids.to(torch.int32) if num_token_non_padded is not None: indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..98562ca0 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -150,6 +150,8 @@ def _infer_tool_call_parser(model_path: str) -> str: if ( "qwen3_5" in marker or "qwen3.5" in marker + or "qwen4_exp" in marker + or "qwen4exp" in marker or ("qwen3" in marker and "coder" in marker) ): return "qwen3_coder" @@ -188,7 +190,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None: tag in marker for tag in ("v4", "deepseek_v4", "v3.2", "v32") ): return "deepseekv32" - if "qwen3" in marker or "qwen3.5" in marker or "qwen3_5" in marker: + if any( + tag in marker + for tag in ("qwen3", "qwen3.5", "qwen3_5", "qwen4_exp", "qwen4exp") + ): return "qwen3" if "glm" in marker: return "glm" diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py index f27640b9..0280874f 100644 --- a/tests/engine/test_attention_backend_matrix.py +++ b/tests/engine/test_attention_backend_matrix.py @@ -245,6 +245,22 @@ def _info(name): _adjust_config(config) +def test_model_runtime_capabilities_force_safe_cache_and_graph(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch) + config = _config("linear_hybrid", attention_backend="auto") + object.__setattr__(config, "cache_type", "radix") + config.model_config.requires_naive_cache = True + config.model_config.supports_cuda_graph = False + + _adjust_config(config) + + assert config.cache_type == "naive" + assert config.cuda_graph_bs == [] + assert config.cuda_graph_max_bs == 0 + + def test_trtllm_page_size_coercion_is_part_aware(monkeypatch): from freetoken.engine.engine import _adjust_config diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd..51873267 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -4,6 +4,15 @@ import torch +def test_fast_index_copy_worker_geometry_supports_qwen4_scale_rows(): + from freetoken.kernel.aot_models import aggregate_fast_index_copy_feature_sizes + from freetoken.kernel.fast_index_copy import default_worker_args + + assert default_worker_args(400)[:2] == (8, 400) + assert default_worker_args(200)[:2] == (16, 200) + assert {200, 400} <= set(aggregate_fast_index_copy_feature_sizes()) + + def test_pinned_extension_uses_packaged_module_not_runtime_jit(monkeypatch): import torch.utils.cpp_extension as cpp_extension @@ -67,6 +76,25 @@ def test_fast_index_copy_accepts_exact_pinned_cpu_source(): torch.testing.assert_close(output.cpu(), source[[5, 2, 0]]) +@pytest.mark.parametrize("row_bytes", [200, 400]) +def test_fast_index_copy_accepts_non_128_byte_rows(row_bytes): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for fast index copy") + + from freetoken.kernel import copy_to_pinned_tensor, fast_index_copy_jit + + values = (torch.arange(6 * row_bytes, dtype=torch.int32) % 251).to(torch.uint8) + source = copy_to_pinned_tensor(values.reshape(6, row_bytes)) + output = torch.empty((3, row_bytes), dtype=torch.uint8, device="cuda") + dst_indices = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") + src_indices = torch.tensor([5, 2, 0], dtype=torch.int32, device="cuda") + + fast_index_copy_jit(output, dst_indices, source, src_indices) + torch.cuda.synchronize() + + assert torch.equal(output.cpu(), source[[5, 2, 0]]) + + def test_fast_index_copy_skip_env_noops_without_jit(monkeypatch): import freetoken.kernel.fast_index_copy as fast_index_copy diff --git a/tests/models/test_qwen4_exp.py b/tests/models/test_qwen4_exp.py new file mode 100644 index 00000000..2ce0e519 --- /dev/null +++ b/tests/models/test_qwen4_exp.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.model import ( + _PLELayer, + _preload_ple_enabled, + _tokens_for_ngram_forward, + build_ngram_ids, +) +from freetoken.models.qwen4_exp.weight import _rename, _try_fuse +from freetoken.models.register import get_model_spec + + +def _config(): + text = SimpleNamespace( + layer_types=[ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + head_dim=256, + rope_parameters={"partial_rotary_factor": 0.25, "rope_theta": 10_000_000}, + indexer_budget=2048, + max_position_embeddings=262_144, + num_key_value_heads=2, + 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, + eos_token_id=248044, + hc_count=4, + hc_lowrank=320, + ple_layer_ids=[2], + ple_embed_dim=2560, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=8, + ngram_vocab_size_base=20_000_000, + split_ngram_parts=128, + indexer_compress_ratio=4, + output_gate_type="sigmoid", + hidden_act="silu", + num_hidden_layers=4, + num_attention_heads=24, + hidden_size=2560, + vocab_size=248320, + rms_norm_eps=1e-6, + num_experts=512, + num_experts_per_tok=10, + moe_intermediate_size=640, + shared_expert_intermediate_size=640, + norm_topk_prob=None, + tie_word_embeddings=False, + ) + return SimpleNamespace( + text_config=text, + quantization_config={"quant_method": "fp8", "weight_block_size": [128, 128]}, + model_type="qwen4_exp", + architectures=["Qwen4ExpForConditionalGeneration"], + image_token_id=248056, + ) + + +def test_qwen4_config_uses_exact_qsa_prefix(): + config = parse_config(_config()) + assert config.rotary_config.max_position == 2048 + assert config.expert_quant == "fp8_block" + assert config.attn_quant == "none" + assert config.qwen4_args.ple_layer_ids == (1,) + assert config.qwen4_args.output_gate_type == "sigmoid" + assert config.requires_naive_cache + assert config.supports_cuda_graph + assert config.is_linear_layer(0) + assert not config.is_linear_layer(3) + + +def test_qwen4_config_accepts_transformers_sparse_attention_alias(): + hf_config = _config() + hf_config.text_config.layer_types[-1] = "qwen_sparse_attention" + config = parse_config(hf_config) + assert not config.is_linear_layer(3) + + +def test_qwen4_config_accepts_modelopt_nvfp4_experts(): + hf_config = _config() + hf_config.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + } + + config = parse_config(hf_config) + + assert config.expert_quant == "nvfp4" + assert config.weight_block_size is None + assert config.attn_quant == "none" + assert config.dense_quant == "none" + + +def test_qwen4_registry_entry(): + spec = get_model_spec("Qwen4ExpForConditionalGeneration") + assert spec.module == "freetoken.models.qwen4_exp" + assert spec.model_cls == "Qwen4ExpForCausalLM" + + +def test_qwen4_weight_names(): + assert _rename("model.language_model.layers.1.ple.key_proj.weight") == ( + "model.layers.1.ple.key_proj.weight" + ) + assert _rename("model.visual.blocks.0.attn.qkv.weight") is None + assert _rename("model.language_model.layers.3.self_attn.indexer.q_layernorm.weight") is None + + +def test_qwen4_projection_fusion_order(): + buffers = {} + base = "model.layers.3.self_attn." + parts = [ + ("q_proj.weight", torch.full((2, 3), 1.0)), + ("k_proj.weight", torch.full((1, 3), 2.0)), + ("v_proj.weight", torch.full((1, 3), 3.0)), + ] + assert _try_fuse(base + parts[0][0], parts[0][1], buffers) == () + assert _try_fuse(base + parts[1][0], parts[1][1], buffers) == () + name, fused = _try_fuse(base + parts[2][0], parts[2][1], buffers) + assert name == base + "qkv_proj.weight" + assert fused[:, 0].tolist() == [1.0, 1.0, 2.0, 3.0] + + +def test_ngram_hash_resets_at_eos(): + tokens = torch.tensor([4, 5, 99, 6, 7]) + multipliers = torch.tensor([3, 5, 7]) + sizes = torch.tensor([101, 103]) + offsets = torch.tensor([0, 101]) + ids = build_ngram_ids( + tokens, + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=sizes, + offsets=offsets, + ) + assert ids.shape == (5, 2) + expected_bigram_after_eos = (6 * 3) ^ (99 * 5) + assert ids[3, 0].item() == expected_bigram_after_eos % 101 + + +def test_ngram_history_includes_inflight_overlap_token(): + req = SimpleNamespace(input_ids=torch.tensor([4, 5, 6]), device_len=4) + + tokens = _tokens_for_ngram_forward(req, torch.tensor([7], device="cpu")) + + assert tokens.tolist() == [4, 5, 6, 7] + assert req.input_ids.tolist() == [4, 5, 6] + + +def test_ngram_history_does_not_duplicate_drained_token(): + req = SimpleNamespace(input_ids=torch.tensor([4, 5, 6, 7]), device_len=4) + + tokens = _tokens_for_ngram_forward(req, torch.tensor([7], device="cpu")) + + assert tokens.tolist() == [4, 5, 6, 7] + + +def test_ngram_history_can_select_only_required_suffix(): + req = SimpleNamespace(input_ids=torch.tensor([4, 5, 6]), device_len=5) + + tokens = _tokens_for_ngram_forward( + req, + torch.tensor([7, 8], device="cpu"), + start=2, + ) + + assert tokens.tolist() == [6, 7, 8] + assert _tokens_for_ngram_forward(req, torch.tensor([7, 8]), start=4).tolist() == [8] + + +def test_incremental_ngram_hash_matches_full_history(): + tokens = torch.tensor([10, 99, 4, 5, 6, 7]) + kwargs = { + "ngram_size": 3, + "heads_per_ngram": 1, + "eos_token_id": 99, + "multipliers": torch.tensor([3, 5, 7]), + "vocab_sizes": torch.tensor([101, 103]), + "offsets": torch.tensor([0, 101]), + } + + full = build_ngram_ids(tokens, **kwargs) + history_start = 3 + incremental = build_ngram_ids(tokens[history_start:], **kwargs) + + assert torch.equal(incremental[2], full[5]) + + +def test_ple_graph_convolution_matches_eager(monkeypatch): + def make_layer(): + layer = object.__new__(_PLELayer) + layer.hidden_size = 2 + layer.hc_count = 2 + layer.state_len = 1 + layer.dilation = 1 + layer.conv1d = SimpleNamespace(weight=torch.randn(4, 1, 2)) + layer._conv_state_pool = None + return layer + + req = SimpleNamespace(table_idx=1, extend_len=1, cached_len=0) + linear_pool = SimpleNamespace(conv_states=torch.empty(1, 3, 1, 1)) + hidden = torch.randn(1, 4) + + eager = make_layer() + eager_batch = SimpleNamespace( + is_decode=True, + reqs=[req], + padded_reqs=[req], + cuda_graph_capture=False, + linear_table_idx=torch.tensor([1], dtype=torch.int32), + ) + context = SimpleNamespace(batch=eager_batch, linear_state_pool=linear_pool) + monkeypatch.setattr("freetoken.models.qwen4_exp.model.get_global_ctx", lambda: context) + eager_output = eager._short_conv(hidden) + + graph = make_layer() + graph.conv1d.weight.copy_(eager.conv1d.weight) + context.batch = SimpleNamespace( + is_decode=True, + reqs=[req], + padded_reqs=[req], + cuda_graph_capture=True, + linear_table_idx=torch.tensor([1], dtype=torch.int32), + ) + graph_output = graph._short_conv(hidden) + + torch.testing.assert_close(graph_output, eager_output) + torch.testing.assert_close(graph._conv_state_pool, eager._conv_state_pool) + + +def test_qwen4_ple_preload_is_opt_in(monkeypatch): + monkeypatch.delenv("FREETOKEN_QWEN4_PLE_PRELOAD", raising=False) + assert not _preload_ple_enabled() + + monkeypatch.setenv("FREETOKEN_QWEN4_PLE_PRELOAD", "true") + assert _preload_ple_enabled() diff --git a/tests/models/test_qwen4_exp_raw_config.py b/tests/models/test_qwen4_exp_raw_config.py new file mode 100644 index 00000000..131e5d9c --- /dev/null +++ b/tests/models/test_qwen4_exp_raw_config.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.utils.hf import RawConfigShim + + +def _raw_checkpoint_config() -> RawConfigShim: + """Raw config shape used when installed Transformers predates Qwen4-Exp.""" + return RawConfigShim( + { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "model_type": "qwen4_exp", + "image_token_id": 248056, + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + "text_config": { + "model_type": "qwen4_exp_text", + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "qwen_sparse_attention", + ], + "head_dim": 256, + "rope_parameters": { + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "rope_type": "default", + }, + "indexer_budget": 2048, + "max_position_embeddings": 262_144, + "num_key_value_heads": 2, + "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, + "eos_token_id": 248044, + "hc_count": 4, + "hc_lowrank": 320, + "ple_layer_ids": [2], + "ple_embed_dim": 2560, + "ple_conv_kernel_size": 4, + "ngram_size": 3, + "heads_per_ngram": 8, + "ngram_vocab_size_base": 20_000_000, + "split_ngram_parts": 128, + "indexer_compress_ratio": 4, + "output_gate_type": "sigmoid", + "hidden_act": "silu", + "num_hidden_layers": 4, + "num_attention_heads": 24, + "hidden_size": 2560, + "vocab_size": 248320, + "rms_norm_eps": 1e-6, + "num_experts": 512, + "num_experts_per_tok": 10, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "tie_word_embeddings": False, + }, + } + ) + + +def test_qwen4_raw_config_uses_official_topk_normalization_default(): + config = parse_config(_raw_checkpoint_config()) + + assert config.norm_topk_prob is True + assert config.rotary_config.max_position == 2048 diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index 1fd0f2e5..068d2e23 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -1,3 +1,6 @@ +import sys +from types import SimpleNamespace + import pytest import torch @@ -58,6 +61,42 @@ def test_fused_topk_accepts_triton_kernel_tuple_output(): torch.testing.assert_close(weights, ref_weights, rtol=2e-4, atol=2e-4) +def test_fused_topk_pads_non_power_of_two_for_triton(monkeypatch): + import freetoken.moe.fused as fused + + monkeypatch.setattr( + "freetoken.kernel.backend.is_triton_kernels_installed", + lambda: True, + ) + called = {} + + def fake_topk(logits, topk, *, apply_softmax): + called["topk"] = topk + values, ids = torch.topk(logits, topk, dim=-1) + if apply_softmax: + values = torch.softmax(values, dim=-1) + return SimpleNamespace(vals=values, indx=ids) + + monkeypatch.setitem(sys.modules, "triton_kernels", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "triton_kernels.topk", SimpleNamespace(topk=fake_topk)) + logits = torch.arange(32, dtype=torch.float32).reshape(2, 16) + hidden_states = torch.zeros((2, 8)) + + weights, ids = fused.fused_topk( + hidden_states, + logits, + topk=10, + renormalize=True, + ) + + assert called["topk"] == 16 + ref_weights, ref_ids = torch.topk(torch.softmax(logits, dim=-1), 10, dim=-1) + ref_weights = ref_weights / ref_weights.sum(dim=-1, keepdim=True) + torch.testing.assert_close(weights, ref_weights) + torch.testing.assert_close(ids, ref_ids.to(torch.int32)) + assert weights.is_contiguous() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 24]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e..1fd464e2 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -87,6 +87,10 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25" +def test_qwen4_exp_uses_qwen3_coder_tool_parser(): + assert _inferred("Qwen4ExpForConditionalGeneration")[0] == "qwen3_coder" + + def test_an_explicit_choice_beats_inference(): config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"}) with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):