feat(model): support Qwen3.8 Flash Next

docs/qwen38_adaptation.rst

@@ -0,0 +1,63 @@
+Qwen3.8-Flash-Next
+==================
+
+Qwen3.8-Flash-Next checkpoints use the ``qwen4_exp`` model type;
+text-only checkpoints use ``qwen4_exp_text``. The model
+implementation combines gated residuals, n-gram embeddings, QSA and GDN
+attention, with dedicated multimodal, MTP and cache-state integration.
+Reusable operations live in ``lightllm/common/gated_residual``,
+``lightllm/common/ngram_embedding`` and ``lightllm/common/qsa``.
+
+CPU embedding offload is disabled by default; enable it with
+``--offload_ngram_embedding``. Shared-prefix admission is limited to the
+model's single-node, DP1 normal serving mode with one-page state blocks.
+
+QSA supports BF16 KV storage and ``--llm_kv_type fp8kv_spt``. The latter
+stores main KV and pooled indexer keys as FP8 e4m3 with unit scales; raw
+indexer keys and rotary inputs retain the model dtype. Calibration files
+are not supported for this QSA format. With ``--page_size 4``, pooled
+indexer keys use one entry per compression group and unfinished groups
+use request-local ring storage that includes MTP candidate capacity.
+
+``--enable_ep_moe`` supports unquantized BF16 experts through DeepEP and
+DeepGEMM. It does not require ``--expert_dtype fp8``. The BF16 path supports
+ordinary prefill/decode and CUDA Graphs; split dispatch/GEMM interfaces
+used by microbatch overlap still require quantized experts.
+
+GDN supports ``--ssm_state_mode replay`` with FP32 SSM state and a bounded
+replay history. It uses the upstream GDN state lifecycle: decode advances
+the verification window, and MTP acceptance trims rejected tokens using
+each request's verification width. Cache pages store prefill Conv/SSM state;
+Qwen lexical state participates in the same acceptance call.
+
+FlashInfer 0.7.2 or later provides an optional GDN prefill backend through
+``--llm_prefill_att_backend auto flashinfer``. It accepts the existing KV
+state layout through explicit transposes and preserves zero-length states.
+Prefill CUDA Graphs reuse the captured GDN subgraph when the request count
+matches; other request counts use the current batch state outside the graph.
+The default linear prefill backend selection is unchanged.
+
+MTP projects draft logits only for accepted-tail rows while retaining all
+hidden rows and KV writes. Its graph cache keys include both input and output
+capacity, and all supported capacities are captured during startup before
+target hidden states can occupy the shared graph pool. Ordered TP reduction
+retains rank-ordered FP64 accumulation; the H100 peer path synchronizes
+readers before reusing its shared buffer.
+
+On H100, TP8/DP1 with BF16 weights, FP8 QSA, FP32 ReplaySSM with a
+16-token precomputed history and fixed MTP3 uses a replicated small-batch
+decode path. Up to 64 verification rows retain the complete residual on
+each TP rank while experts remain distributed across eight ranks. Each
+graph shape owns its expert workspace. Shared-expert gating is combined with
+the routed-expert sum while preserving BF16 intermediate rounding. QSA decode
+with up to 128 rows packs the attention and indexer projections into one GEMM;
+prefill and draft execution retain separate projections.
+Larger batches, unsupported expert
+layouts, quantized weights, microbatch overlap and routed-expert collection
+use the existing path. The admitted profile also fixes the QSA split count
+and ReplaySSM kernel configuration across decode batch sizes; this does not
+establish end-to-end deterministic generation.
+
+Kernel microbenchmarks are in ``test/benchmark/qwen38/``. Numerical and
+cache-state regression tests are in ``unit_tests/models/qwen4_exp/`` and
+the corresponding ``unit_tests/common/`` directories.

lightllm/common/attention_state/bounded_att/utils.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from math import prod
+from math import lcm, prod
+from mmap import PAGESIZE
 from weakref import finalize
 
 import torch
@@ -17,6 +18,7 @@
 CUDA_HOST_REGISTER_PORTABLE = 1
 # Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。
 CUDA_HOST_REGISTER_MAPPED = 2
+_HOST_REGISTRATION_BYTES = 1024 * 1024 * 1024
 
 
 # ========================================================================
@@ -43,16 +45,25 @@ def create_page_pool_tensor(
 
     # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。
     if not shared:
-        tensor = torch.empty(
-            size=size,
-            dtype=dtype,
-            device=device,
-            pin_memory=pin_memory,
-        )
+        try:
+            tensor = torch.empty(
+                size=size,
+                dtype=dtype,
+                device=device,
+                pin_memory=pin_memory,
+            )
+        except RuntimeError:
+            if device.type != "cpu" or not pin_memory or not torch.cuda.is_available():
+                raise
+            tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=False)
+            register_page_pool_tensor(tensor=tensor)
+            allocation = "torch.empty + cudaHostRegister"
+        else:
+            allocation = "torch.empty"
         reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled"
         logger.info(
             f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, "
-            f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}"
+            f"device={device}, shared=False, pin_memory={pin_memory}, allocator={allocation}, reason={reason}"
         )
         return tensor
 
@@ -91,24 +102,42 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None:
     """
     if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned():
         return
-    assert tensor.is_shared(), "page pool registration requires shared CPU storage"
     storage = tensor.untyped_storage()
     ptr = storage.data_ptr()
     nbytes = storage.nbytes()
     device_id = torch.cuda.current_device()
     flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED
-    torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags))
+    registration_unit = lcm(prod(tensor.shape[1:]) * tensor.element_size(), PAGESIZE)
+    chunk_bytes = max(registration_unit, _HOST_REGISTRATION_BYTES // registration_unit * registration_unit)
+    if not tensor.is_shared():
+        chunk_bytes = nbytes
+    addresses = []
+    runtime = torch.cuda.cudart()
+    try:
+        for offset in range(0, nbytes, chunk_bytes):
+            address = ptr + offset
+            torch.cuda.check_error(runtime.cudaHostRegister(address, min(chunk_bytes, nbytes - offset), flags))
+            addresses.append(address)
+    except BaseException:
+        logger.error(
+            f"Failed to register bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, "
+            f"chunk_bytes={chunk_bytes}, registered_chunks={len(addresses)}"
+        )
+        if addresses:
+            _unregister_page_pool_storage(tuple(addresses), device_id)
+        raise
     # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。
     # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。
     cleanup = finalize(
         storage,
         _unregister_page_pool_storage,
-        ptr,
+        tuple(addresses),
         device_id,
     )
     cleanup.atexit = False
     logger.info(
-        f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, "
+        f"Registered shared bounded attention page pool: bytes={nbytes}, chunks={len(addresses)}, "
+        f"cuda_device={device_id}, "
         "flags=Portable|Mapped; registration applies to this process's mapping"
     )
 
@@ -118,8 +147,9 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None:
 # ========================================================================
 
 
-def _unregister_page_pool_storage(ptr: int, device_id: int) -> None:
+def _unregister_page_pool_storage(addresses: tuple[int, ...], device_id: int) -> None:
     """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。"""
     torch.cuda.synchronize(device=device_id)
-    torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr))
-    logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}")
+    for ptr in reversed(addresses):
+        torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr))
+    logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, chunks={len(addresses)}")

lightllm/common/attention_state/full_att/gqa/gqa_state.py

@@ -146,14 +146,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor,
         if not len(token_indices):
             return
         assert self.k_storage.is_cuda, "KV writes require CUDA storage"
-        from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv
+        from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv_fused
 
-        # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。
-        destindex_copy_kv(
-            K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index]
-        )
-        destindex_copy_kv(
-            K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index]
+        destindex_copy_kv_fused(
+            KV=kv,
+            DestLoc=token_indices,
+            OutK=self.k_storage[storage_layer_index],
+            OutV=self.v_storage[storage_layer_index],
         )
 
     def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None:
@@ -169,22 +168,16 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten
         assert self.k_storage.is_cuda, "token copies require CUDA storage"
         if not src_indices.numel():
             return
-        from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer
+        from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_pair_to_kv_buffer_pair
 
         # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。
         # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。
         # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。
-        # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射，
-        # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。
-        copy_kv_buffer_to_kv_buffer(
-            src_mem_index=src_indices,
-            dst_mem_index=dst_indices,
-            kv_buffer=self.k_storage,
-        )
-        copy_kv_buffer_to_kv_buffer(
+        copy_kv_buffer_pair_to_kv_buffer_pair(
             src_mem_index=src_indices,
             dst_mem_index=dst_indices,
-            kv_buffer=self.v_storage,
+            k_buffer=self.k_storage,
+            v_buffer=self.v_storage,
         )
 
     # ========================================================================

lightllm/common/basemodel/attention/create_linear_utils.py

@@ -1,41 +1,36 @@
-from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend
-from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend
-from lightllm.common.basemodel.attention.linear.replay_gdn import (
-    ReplayGDNFlashQlaAttBackend,
-    ReplayGDNTritonAttBackend,
-)
-from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend
-from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend
+from .import_utils import load_backend_class
 from lightllm.utils.backend_validator import validate
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
 
 gdn_prefill_backends = {
-    "flashqla": FlashQlaLinearAttBackend,
-    "triton": TritonLinearAttBackend,
+    "flashqla": "linear.flashqla:FlashQlaLinearAttBackend",
+    "flashinfer": "linear.flashinfer:FlashInferLinearAttBackend",
+    "triton": "linear.triton:TritonLinearAttBackend",
 }
 
 gdn_replay_prefill_backends = {
-    "flashqla": ReplayGDNFlashQlaAttBackend,
-    "triton": ReplayGDNTritonAttBackend,
+    "flashqla": "linear.replay_gdn:ReplayGDNFlashQlaAttBackend",
+    "flashinfer": "linear.replay_gdn:ReplayGDNFlashInferAttBackend",
+    "triton": "linear.replay_gdn:ReplayGDNTritonAttBackend",
 }
 
 gdn_decode_backends = {
-    "triton": TritonLinearAttBackend,
+    "triton": "linear.triton:TritonLinearAttBackend",
 }
 
 gdn_replay_decode_backends = {
-    "triton": ReplayGDNTritonAttBackend,
+    "triton": "linear.replay_gdn:ReplayGDNTritonAttBackend",
 }
 
 kda_backends = {
-    "triton": KDALinearAttBackend,
+    "triton": "linear.kda:KDALinearAttBackend",
 }
 
 kda_replay_backends = {
-    "triton": ReplayKDALinearAttBackend,
+    "triton": "linear.replay_kda:ReplayKDALinearAttBackend",
 }
 
 
@@ -54,24 +49,24 @@ def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)):
 def get_kda_linear_att_backend_class():
     args = get_env_start_args()
     classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends
-    return classes["triton"]
+    return load_backend_class(classes["triton"])
 
 
 def _select_backend(classes, backend_args, index, priority_list, stage):
     backend_str = _get_backend_str(backend_args, index)
     if backend_str != "auto":
-        return classes[backend_str]
+        return load_backend_class(classes[backend_str])
 
     for backend_name in priority_list:
         if backend_name == "triton" or validate(backend_name):
             if backend_name == "triton":
                 logger.info(f"Linear {stage} attention backend: triton.")
             else:
                 logger.info(f"Linear {stage} attention backend: {backend_name} (validated).")
-            return classes[backend_name]
+            return load_backend_class(classes[backend_name])
 
     logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.")
-    return classes["triton"]
+    return load_backend_class(classes["triton"])
 
 
 def _get_backend_str(backend_args, index):

lightllm/common/basemodel/attention/linear/flashinfer.py

@@ -0,0 +1,63 @@
+import torch
+
+from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd
+
+from .gdn import LinearAttBackend
+
+
+def chunk_gated_delta_rule(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    initial_state=None,
+    output_final_state=False,
+    cu_seqlens=None,
+    use_qk_l2norm_in_kernel=False,
+    scale=None,
+):
+    from flashinfer.gdn_prefill import chunk_gated_delta_rule as flashinfer_gdn
+
+    if q.shape[0] != 1 or cu_seqlens is None:
+        raise ValueError("FlashInfer GDN requires packed variable-length inputs")
+    q, k, v = (tensor.squeeze(0).contiguous() for tensor in (q, k, v))
+    if use_qk_l2norm_in_kernel:
+        q, k = l2norm_fwd(q), l2norm_fwd(k)
+    state = None
+    if initial_state is not None:
+        state = initial_state.transpose(-1, -2).to(torch.float32).contiguous()
+    final = None
+    if output_final_state:
+        final = (
+            state.clone()
+            if state is not None
+            else torch.zeros(
+                (cu_seqlens.numel() - 1, max(q.shape[1], v.shape[1]), v.shape[2], k.shape[2]),
+                device=q.device,
+                dtype=torch.float32,
+            )
+        )
+    result = flashinfer_gdn(
+        q=q,
+        k=k,
+        v=v,
+        g=g.squeeze(0).float().exp(),
+        beta=beta.squeeze(0).float().contiguous(),
+        scale=scale,
+        initial_state=state,
+        output_final_state=output_final_state,
+        output_state=final,
+        cu_seqlens=cu_seqlens.to(torch.int64),
+        use_cp=False if q.shape[0] <= 1024 * (cu_seqlens.numel() - 1) else "auto",
+        backend="flashinfer",
+    )
+    if output_final_state:
+        output, state = result
+        return output.unsqueeze(0), state.transpose(-1, -2).contiguous()
+    return result.unsqueeze(0), None
+
+
+class FlashInferLinearAttBackend(LinearAttBackend):
+    def get_prefill_kernel(self):
+        return chunk_gated_delta_rule