feat(models): add Qwen3.8 Flash Next support

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.

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

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",

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)

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",

python/freetoken/kernel/csrc/jit/fast_index_copy.cuh

@@ -88,9 +88,9 @@ using mem_package_t = decltype(get_mem_package<kUnit>());
 template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads>
 __always_inline __device__ auto load_vec(const void* __restrict__ src) {
     using Package = mem_package_t<kUnit>;
-    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<const Package*>(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 <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads, typename Tp>
 __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
     using Package = mem_package_t<kUnit>;
-    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<Tp, device::device_vec<Package, kLoopCount>>);
 
     const auto dst_packed = static_cast<Package*>(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]);
+        }
     }
 }
 

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

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)
 

python/freetoken/models/blocks.py

@@ -16,13 +16,21 @@
 if TYPE_CHECKING:
     import torch
 
+    from freetoken.core import Batch
+
     from .config import ModelConfig
 
 
 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):

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:

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",
+]

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"]