diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py
index 7ba00e5911..0e24f439dd 100644
--- a/lightllm/common/basemodel/attention/fa3/fp.py
+++ b/lightllm/common/basemodel/attention/fa3/fp.py
@@ -10,6 +10,20 @@
from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor
from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers
+try:
+ from flash_attn_interface import flash_attn_with_kvcache as flash_attn_with_kvcache_neo
+ import inspect
+
+ # Verify this is the neo-patched FA3 build (with image_token_tag support),
+ _sig = inspect.signature(flash_attn_with_kvcache_neo)
+ if "image_token_tag" not in _sig.parameters:
+ raise ImportError("flash_attn_interface found but missing image_token_tag support (need neo build)")
+
+ HAS_FLASH_ATTN_INTERFACE = True
+except ImportError:
+ flash_attn_with_kvcache_neo = None
+ HAS_FLASH_ATTN_INTERFACE = False
+
class Fa3AttBackend(BaseAttBackend):
"""Common fixed page-table storage for FA3 attention backends."""
@@ -129,6 +143,37 @@ def _nomarl_prefill_att(
k_descale, v_descale = None, None # disable quantization
Lq = q.shape[-1]
sm_scale = 1.0 / (Lq ** 0.5)
+
+ # neo_chat*: image-token bidirectional attention requires flash_attn_interface
+ # (sgl_kernel's flash_attn_with_kvcache does not support image_token_tag).
+ if att_control.image_token_tag is not None:
+ if not HAS_FLASH_ATTN_INTERFACE:
+ raise ImportError(
+ "flash_attn_interface (fa3-neo) is required for image_token_tag bidirectional "
+ "attention. Install it or set LIGHTLLM_NEO_PREFILL_TRITON_BACKEND=1 to use the "
+ "triton fallback."
+ )
+ extra_kwargs = {"image_token_tag": att_control.image_token_tag}
+ o = flash_attn_with_kvcache_neo(
+ q=q,
+ k_cache=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
+ v_cache=v.view(v.shape[0], 1, v.shape[1], v.shape[2]),
+ page_table=self.page_table,
+ cache_seqlens=self.infer_state.b_seq_len,
+ cu_seqlens_q=self.cu_seqlens_q,
+ cu_seqlens_k_new=self.cu_seqlens_k,
+ max_seqlen_q=self.infer_state.max_q_seq_len,
+ softmax_scale=sm_scale,
+ causal=True,
+ window_size=window_size,
+ softcap=0.0,
+ k_descale=k_descale,
+ v_descale=v_descale,
+ return_softmax_lse=False,
+ **extra_kwargs,
+ )
+ return o
+
o = flash_attn_with_kvcache(
q=q,
k_cache=k.view(k.shape[0], 1, k.shape[1], k.shape[2]),
diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py
index c7e9a59aad..18baf58207 100644
--- a/lightllm/models/__init__.py
+++ b/lightllm/models/__init__.py
@@ -41,6 +41,8 @@
)
from lightllm.models.gpt_oss.model import GptOssTpPartModel
from lightllm.models.qwen3_omni_moe_thinker.model import Qwen3OmniMOETpPartModel
+from lightllm.models.neo_chat_moe.model import NeoTpMOEPartModel
+from lightllm.models.neo_chat.model import NeoTpPartModel
from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel
from lightllm.models.qwen3_5_moe.model import Qwen3_5MOETpPartModel
from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel
diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py
index c104ebccc9..b6ab734c2a 100644
--- a/lightllm/models/llama/model.py
+++ b/lightllm/models/llama/model.py
@@ -74,14 +74,20 @@ def _init_custom(self):
rope_scaling = self.config.get("rope_scaling", None)
if rope_scaling is None:
self._init_to_get_rotary()
- return
-
- if "rope_type" in rope_scaling:
+ elif "rope_type" in rope_scaling:
scaling_type = rope_scaling["rope_type"]
+ self._init_rotary_by_scaling_type(scaling_type, rope_scaling)
elif "type" in rope_scaling:
scaling_type = rope_scaling["type"]
+ self._init_rotary_by_scaling_type(scaling_type, rope_scaling)
else:
raise ValueError(f"Unknown RoPE scaling format {rope_scaling}")
+
+ if "rope_theta_hw" in self.config:
+ self._init_to_get_hw_rotary()
+ super()._init_custom()
+
+ def _init_rotary_by_scaling_type(self, scaling_type, rope_scaling):
if scaling_type == "default" or "mrope_section" in rope_scaling:
self._init_to_get_rotary()
elif scaling_type == "yarn":
@@ -126,9 +132,10 @@ def _init_to_get_rotary(self, default_base=10000):
except:
pass
- inv_freq = 1.0 / (
+ full_inv_freq = 1.0 / (
base ** (torch.arange(0, partial_head_dim, 2, device="cpu", dtype=torch.float32) / partial_head_dim)
)
+ inv_freq = full_inv_freq[::2] # for neo
t = (
torch.arange(max(max_seq_len + 1024 * 128, self.max_seq_length), device="cpu", dtype=torch.float32)
/ rope_scaling_factor
@@ -139,6 +146,46 @@ def _init_to_get_rotary(self, default_base=10000):
self._sin_cached = torch.sin(freqs).to(self.data_type).cuda()
return
+ def _init_to_get_hw_rotary(self, default_base=10000):
+ partial_head_dim = int(self.config.get("partial_rotary_factor", 1) * self.head_dim_ // 2)
+ if self.config.get("rope_scaling", {}) is None:
+ rope_scaling_factor = 1.0
+ else:
+ rope_scaling_factor = self.config.get("rope_scaling", {}).get("factor", 1.0)
+
+ base = self.config.get("rope_theta_hw", float(default_base))
+ if "max_sequence_length" in self.config:
+ max_seq_len = self.config["max_sequence_length"]
+ else:
+ max_position_embeddings = self.config.get(
+ "max_position_embeddings_hw", 2048 if base <= 10000.0 + 1e-5 else 16384
+ )
+ max_seq_len = max_position_embeddings * rope_scaling_factor
+
+ # NTK
+ try:
+ ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1))
+ assert ntk_alpha >= 1
+ if ntk_alpha > 1:
+ logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}")
+ max_seq_len *= ntk_alpha
+ base = base * (ntk_alpha ** (partial_head_dim / (partial_head_dim - 2))) # Base change formula
+ except:
+ pass
+ full_inv_freq = 1.0 / (
+ base ** (torch.arange(0, partial_head_dim, 2, device="cpu", dtype=torch.float32) / partial_head_dim)
+ )
+ inv_freq = full_inv_freq[::2]
+ t = (
+ torch.arange(max(max_seq_len + 1024 * 128, self.max_seq_length), device="cpu", dtype=torch.float32)
+ / rope_scaling_factor
+ )
+ freqs = torch.outer(t, inv_freq)
+
+ self._hw_cos_cached = torch.cos(freqs).to(self.data_type).cuda()
+ self._hw_sin_cached = torch.sin(freqs).to(self.data_type).cuda()
+ return
+
def _init_to_get_dynamic_ntk_rotary(self):
partial_head_dim = int(self.config.get("partial_rotary_factor", 1) * self.head_dim_)
max_position_embeddings = self.config.get("max_position_embeddings", 2048)
diff --git a/lightllm/models/neo_chat/__init__.py b/lightllm/models/neo_chat/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat/layer_infer/__init__.py b/lightllm/models/neo_chat/layer_infer/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py
new file mode 100644
index 0000000000..cf88868a8a
--- /dev/null
+++ b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py
@@ -0,0 +1,148 @@
+import os
+import torch
+from functools import partial
+from typing import Tuple
+from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
+from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo
+from lightllm.models.neo_chat_moe.triton_kernel.context_attention_fwd_neo import context_attention_fwd_neo
+from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd
+from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer
+from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight
+from lightllm.distributed import all_reduce
+import torch.distributed as dist
+from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
+from lightllm.common.basemodel.attention.base_att import AttControl
+from lightllm.common.basemodel.attention.fa3.fp import HAS_FLASH_ATTN_INTERFACE
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+_USE_TRITON_PREFILL = os.environ.get("LIGHTLLM_NEO_PREFILL_TRITON_BACKEND", "0").strip().lower() in ("1", "true")
+if not _USE_TRITON_PREFILL and not HAS_FLASH_ATTN_INTERFACE:
+ logger.warning(
+ "flash_attn_interface (fa3-neo) is not installed; falling back to triton prefill backend "
+ "for neo_chat. Install fa3-neo or set LIGHTLLM_NEO_PREFILL_TRITON_BACKEND=1 to silence "
+ "this warning."
+ )
+ _USE_TRITON_PREFILL = True
+
+
+class NeoChatTransformerLayerInfer(Qwen3TransformerLayerInfer):
+ def __init__(self, data_type, network_config):
+ super().__init__(data_type, network_config)
+ return
+
+ def _bind_attention(self):
+ self._context_attention_kernel = self._context_attention_kernel
+ self._token_attention_kernel = self._token_attention_kernel
+ return
+
+ def _get_qkv(self, input, infer_state: NeoChatInferStateInfo, layer_weight: NeoChatTransformerLayerWeight):
+ input = input.view(-1, self.embed_dim_)
+
+ qkv = layer_weight.qkv_proj.mm(input)
+ q, cache_kv = qkv.split(
+ [self.tp_q_head_num_ * self.head_dim_, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], dim=-1
+ )
+ q = q.view(q.shape[0], self.tp_q_head_num_, self.head_dim_)
+ q_t, q_hw = q.chunk(2, dim=-1)
+
+ cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+ k = cache_kv[:, : self.tp_k_head_num_, :]
+ v = cache_kv[:, self.tp_k_head_num_ :, :]
+ k_t, k_hw = k.chunk(2, dim=-1)
+
+ q_t_2d = q_t.reshape(q.shape[0], -1)
+ q_hw_2d = q_hw.reshape(q.shape[0], -1)
+ k_t_2d = k_t.reshape(k.shape[0], -1)
+ k_hw_2d = k_hw.reshape(k.shape[0], -1)
+
+ layer_weight.qk_norm_weight_(q_t_2d, k_t_2d, eps=self.eps_)
+ layer_weight.qk_hw_norm_weight_(q_hw_2d, k_hw_2d, eps=self.eps_)
+
+ q_t = q_t_2d.view(q.shape[0], self.tp_q_head_num_, self.head_dim_ // 2)
+ q_hw = q_hw_2d.view(q.shape[0], self.tp_q_head_num_, self.head_dim_ // 2)
+ q_h, q_w = q_hw.chunk(2, dim=-1)
+
+ k_t = k_t_2d.view(k.shape[0], self.tp_k_head_num_, self.head_dim_ // 2)
+ k_hw = k_hw_2d.view(k.shape[0], self.tp_k_head_num_, self.head_dim_ // 2)
+ k_h, k_w = k_hw.chunk(2, dim=-1)
+
+ rotary_emb_fwd(
+ q_t,
+ k_t,
+ infer_state.position_cos,
+ infer_state.position_sin,
+ )
+ rotary_emb_fwd(
+ q_h,
+ k_h,
+ infer_state.position_cos_h,
+ infer_state.position_sin_h,
+ )
+ rotary_emb_fwd(
+ q_w,
+ k_w,
+ infer_state.position_cos_w,
+ infer_state.position_sin_w,
+ )
+
+ q = torch.cat([q_t, q_h, q_w], dim=-1)
+ q = q.reshape(q.shape[0], -1)
+
+ k = torch.cat([k_t, k_h, k_w], dim=-1)
+ cache_kv = torch.cat([k, v], dim=1)
+ return q, cache_kv
+
+ def _context_attention_kernel(
+ self, q, kv, infer_state: NeoChatInferStateInfo, layer_weight, out=None
+ ) -> torch.Tensor:
+
+ if _USE_TRITON_PREFILL:
+ o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
+ kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+ context_attention_fwd_neo(
+ q.view(-1, self.tp_q_head_num_, self.head_dim_),
+ kv[:, 0 : self.tp_k_head_num_, :],
+ kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
+ o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
+ infer_state.b_req_idx,
+ infer_state.b_q_start_loc,
+ infer_state.b_seq_len,
+ infer_state.b_ready_cache_len,
+ infer_state.max_q_seq_len,
+ infer_state.req_manager.req_to_token_indexs,
+ infer_state.b_image_token_end,
+ )
+ return o_tensor
+
+ _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+ _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+
+ att_control = AttControl()
+ att_control.image_token_tag = getattr(infer_state, "b_image_token_tag", None)
+
+ o_tensor = infer_state.prefill_att_state.prefill_att(
+ q=_q,
+ k=_k,
+ v=_v,
+ att_control=att_control,
+ alloc_func=self.alloc_tensor,
+ )
+ return o_tensor.view(q.shape)
+
+ def _token_attention_kernel(
+ self,
+ q: torch.Tensor,
+ infer_state: NeoChatInferStateInfo,
+ layer_weight: NeoChatTransformerLayerWeight,
+ ) -> torch.Tensor:
+ _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+ _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+ att_control = AttControl()
+ # att_control.mla_decode_dict["softmax_scale"] = 1.0 / (self.head_dim_ ** 0.5)
+ o_tensor = infer_state.decode_att_state.decode_att(
+ q=_q, k=_k, v=_v, att_control=att_control, alloc_func=self.alloc_tensor
+ )
+ o_tensor = o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_)[:, :, : self.head_dim_].contiguous()
+ return o_tensor
diff --git a/lightllm/models/neo_chat/layer_weights/__init__.py b/lightllm/models/neo_chat/layer_weights/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/neo_chat/layer_weights/pre_and_post_layer_weight.py
new file mode 100644
index 0000000000..e6489f39af
--- /dev/null
+++ b/lightllm/models/neo_chat/layer_weights/pre_and_post_layer_weight.py
@@ -0,0 +1,23 @@
+import torch
+import numpy as np
+from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight
+
+# add key: language_model.xxx -> xxx
+# only change keys at PreAndPostLayerWeight load, TransformLayerWeight is correct now
+def rename_weight_keys(weights):
+ prefix = "language_model."
+ keys = list(weights.keys())
+ for k in keys:
+ if prefix in k:
+ weights[k.replace(prefix, "")] = weights.pop(k)
+
+
+class NeoChatPreAndPostLayerWeight(Qwen2PreAndPostLayerWeight):
+ def __init__(self, data_type, network_config):
+ super().__init__(data_type, network_config)
+ return
+
+ def load_hf_weights(self, weights):
+ rename_weight_keys(weights)
+ super().load_hf_weights(weights)
+ return
diff --git a/lightllm/models/neo_chat/layer_weights/transformer_layer_weight.py b/lightllm/models/neo_chat/layer_weights/transformer_layer_weight.py
new file mode 100644
index 0000000000..8351369fd8
--- /dev/null
+++ b/lightllm/models/neo_chat/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,57 @@
+from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+ QKRMSNORMWeight,
+ RMSNormWeight,
+ QKVROWNMMWeight,
+)
+
+
+class NeoChatTransformerLayerWeight(Qwen3TransformerLayerWeight):
+ def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+ super().__init__(layer_num, data_type, network_config, quant_cfg)
+ return
+
+ def _init_weight_names(self):
+ super()._init_weight_names()
+ self._q_norm_hw_name = f"model.layers.{self.layer_num_}.self_attn.q_norm_hw.weight"
+ self._k_norm_hw_name = f"model.layers.{self.layer_num_}.self_attn.k_norm_hw.weight"
+
+ def _init_qkv(self):
+ in_dim = self.n_embed
+ self.qkv_proj = QKVROWNMMWeight(
+ in_dim=in_dim,
+ q_head_num=self.q_head_num_,
+ kv_head_num=self.k_head_num_,
+ head_dim=self.head_dim,
+ weight_names=[self._q_weight_name, self._k_weight_name, self._v_weight_name],
+ data_type=self.data_type_,
+ bias_names=[self._q_bias_name, self._k_bias_name, self._v_bias_name],
+ quant_method=self.get_quant_method("qkv_proj"),
+ )
+
+ def _init_norm(self):
+ hidden_size = self.network_config_["hidden_size"]
+ self.att_norm_weight_ = RMSNormWeight(
+ dim=hidden_size,
+ weight_name=self._att_norm_weight_name,
+ data_type=self.data_type_,
+ )
+ self.ffn_norm_weight_ = RMSNormWeight(
+ dim=hidden_size,
+ weight_name=self._ffn_norm_weight_name,
+ data_type=self.data_type_,
+ )
+
+ self.qk_norm_weight_ = QKRMSNORMWeight(
+ dim=self.head_dim // 2,
+ q_weight_name=self._q_norm_name,
+ k_weight_name=self._k_norm_name,
+ data_type=self.data_type_,
+ )
+
+ self.qk_hw_norm_weight_ = QKRMSNORMWeight(
+ dim=self.head_dim // 2,
+ q_weight_name=self._q_norm_hw_name,
+ k_weight_name=self._k_norm_hw_name,
+ data_type=self.data_type_,
+ )
diff --git a/lightllm/models/neo_chat/model.py b/lightllm/models/neo_chat/model.py
new file mode 100644
index 0000000000..80621034ae
--- /dev/null
+++ b/lightllm/models/neo_chat/model.py
@@ -0,0 +1,66 @@
+import os
+import json
+from lightllm.common.build_utils import repair_config
+from lightllm.models.registry import ModelRegistry, llm_model_type_is
+from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo
+from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer
+from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer
+from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight
+from lightllm.models.qwen2_vl.model import QWen2VLTokenizer
+from lightllm.models.qwen3.model import Qwen3TpPartModel
+from lightllm.server.core.objs import SamplingParams
+from lightllm.models.qwen3_moe.model import Qwen3MOEModel
+from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem
+from lightllm.models.neo_chat_moe.vision_process import smart_resize
+from lightllm.models.internvl.model import InternvlTokenizer
+from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
+from lightllm.models.neo_chat.layer_infer.transformer_layer_infer import NeoChatTransformerLayerInfer
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight
+from lightllm.models.neo_chat.layer_weights.pre_and_post_layer_weight import NeoChatPreAndPostLayerWeight
+from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer
+from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo
+from lightllm.common.basemodel.attention import (
+ get_prefill_att_backend_class,
+ get_decode_att_backend_class,
+ BaseAttBackend,
+)
+
+
+@ModelRegistry(["neo_chat"], is_multimodal=True, condition=llm_model_type_is("qwen3"))
+class NeoTpPartModel(Qwen3TpPartModel):
+
+ pre_layer_infer_class = LlamaMultimodalPreLayerInfer
+ transformer_layer_infer_class = NeoChatTransformerLayerInfer
+
+ pre_and_post_weight_class = NeoChatPreAndPostLayerWeight
+ transformer_weight_class = NeoChatTransformerLayerWeight
+
+ infer_state_class = NeoChatInferStateInfo
+
+ def __init__(self, kvargs):
+ super().__init__(kvargs)
+ return
+
+ def _init_inferstate_cls(self):
+ pass
+
+ def _init_att_backend(self):
+ self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0, priority_list=["fa3"])(
+ model=self
+ )
+ self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0, priority_list=["fa3"])(
+ model=self
+ )
+
+ def _init_config(self):
+ with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
+ all_config = json.load(json_file)
+ self.config = all_config["llm_config"]
+ # rename keys
+ repair_config(self.config, same_names=["num_attention_heads", "n_head"])
+ repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"])
+ repair_config(self.config, same_names=["num_hidden_layers", "n_layer"])
+ if self.finetune_config:
+ self.config["vocab_size"] = self.finetune_config.vocab_size
+ return
diff --git a/lightllm/models/neo_chat_moe/__init__.py b/lightllm/models/neo_chat_moe/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat_moe/infer_struct.py b/lightllm/models/neo_chat_moe/infer_struct.py
new file mode 100644
index 0000000000..c52147cd9b
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/infer_struct.py
@@ -0,0 +1,101 @@
+from typing import Optional, List
+import torch
+import numpy as np
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.common.req_manager import ReqManager
+from lightllm.models.neo_chat_moe.triton_kernel.get_neo_position import get_neo_position_triton
+from lightllm.models.llama.model import LlamaTpPartModel
+
+
+class NeoChatInferStateInfo(LlamaInferStateInfo):
+ def __init__(self):
+ super().__init__()
+ self.position_cos = None
+ self.position_sin = None
+ self.position_cos_h = None
+ self.position_sin_h = None
+ self.position_cos_w = None
+ self.position_sin_w = None
+
+ def init_some_extra_state(self, model: LlamaTpPartModel):
+ LlamaInferStateInfo.init_some_extra_state(self, model)
+ if self.is_prefill:
+ self.b_image_token_end = torch.zeros([self.position_ids.size(0)], dtype=torch.int32, device="cpu").cuda(
+ non_blocking=True
+ )
+ self.position_ids = self.get_neo_position(self.multimodal_params)
+ else:
+ b_position_delta = [0 for _ in range(self.b_seq_len.shape[0])]
+ for batch_idx, p in enumerate(self.multimodal_params):
+ position_delta = 0
+ for image in p["images"]:
+ position_delta += image["grid_thwd"][3]
+ b_position_delta[batch_idx] = position_delta
+ position_ids = self.position_ids + torch.tensor(b_position_delta, device=self.position_ids.device)
+ self.position_ids = position_ids.unsqueeze(0).expand(3, -1).clone()
+ self.position_ids[1:].zero_()
+
+ self.position_ids = self.position_ids.contiguous()
+ self.position_cos = model._cos_cached[self.position_ids[0]]
+ self.position_sin = model._sin_cached[self.position_ids[0]]
+ self.position_cos_h = model._hw_cos_cached[self.position_ids[1]]
+ self.position_sin_h = model._hw_sin_cached[self.position_ids[1]]
+ self.position_cos_w = model._hw_cos_cached[self.position_ids[2]]
+ self.position_sin_w = model._hw_sin_cached[self.position_ids[2]]
+ return
+
+ def get_neo_position(self, multimodal_params: List[dict]) -> torch.Tensor:
+ if len(multimodal_params) == 0:
+ position_ids = self.position_ids.new_zeros((3, self.position_ids.size(0)))
+ position_ids[0].copy_(self.position_ids)
+ return position_ids
+ b_image_start_idx = []
+ b_image_nums = []
+ b_image_start_num = []
+ b_image_len = []
+ image_start_num = 0
+ b_image_thwd = []
+
+ # pad multimodal_params to batch size.
+ batch_size = self.b_q_seq_len.shape[0]
+ multimodal_params = multimodal_params + [
+ {"images": [], "audios": []} for _ in range(batch_size - len(multimodal_params))
+ ]
+
+ for _, p in enumerate(multimodal_params):
+ images = p.get("images", [])
+ for img in images:
+ b_image_start_idx.append(img["start_idx"])
+ b_image_len.append(img["token_num"])
+ b_image_thwd.append(img["grid_thwd"])
+ b_image_nums.append(len(images))
+ b_image_start_num.append(image_start_num)
+ image_start_num += len(images)
+
+ # 没有任何图片
+ if image_start_num == 0:
+ position_ids = self.position_ids.new_zeros((3, self.position_ids.size(0)))
+ position_ids[0].copy_(self.position_ids)
+ return position_ids.contiguous()
+ b_image_start_idx = torch.tensor(b_image_start_idx, device="cpu").cuda(non_blocking=True)
+ b_image_thwd = torch.tensor(b_image_thwd, device="cpu").cuda(non_blocking=True) # image_num x 4
+ b_image_nums = torch.tensor(b_image_nums, device="cpu").cuda(non_blocking=True)
+ b_image_start_num = torch.tensor(b_image_start_num, device="cpu").cuda(non_blocking=True)
+ b_image_len = torch.tensor(b_image_len, device="cpu").cuda(non_blocking=True)
+
+ position_ids = self.position_ids.new_zeros((3, self.position_ids.size(0)))
+ position_ids[0].copy_(self.position_ids)
+
+ get_neo_position_triton(
+ b_image_start_idx=b_image_start_idx,
+ b_image_thwd=b_image_thwd,
+ b_image_nums=b_image_nums,
+ b_image_start_num=b_image_start_num,
+ b_image_len=b_image_len,
+ position_ids=position_ids,
+ b_ready_cache_len=self.b_ready_cache_len,
+ b_q_seq_len=self.b_q_seq_len,
+ b_start_loc=self.b_q_start_loc,
+ b_image_token_end=self.b_image_token_end,
+ )
+ return position_ids
diff --git a/lightllm/models/neo_chat_moe/layer_infer/__init__.py b/lightllm/models/neo_chat_moe/layer_infer/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py
new file mode 100644
index 0000000000..9a791ff221
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py
@@ -0,0 +1,146 @@
+import os
+import torch
+from functools import partial
+from typing import Tuple
+from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
+from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo
+from lightllm.models.neo_chat_moe.triton_kernel.context_attention_fwd_neo import context_attention_fwd_neo
+from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd
+from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer
+from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight
+from lightllm.distributed import all_reduce
+import torch.distributed as dist
+from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
+from lightllm.common.basemodel.attention.base_att import AttControl
+from lightllm.common.basemodel.attention.fa3.fp import HAS_FLASH_ATTN_INTERFACE
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+_USE_TRITON_PREFILL = os.environ.get("LIGHTLLM_NEO_PREFILL_TRITON_BACKEND", "0").strip().lower() in ("1", "true")
+if not HAS_FLASH_ATTN_INTERFACE:
+ logger.warning(
+ "flash_attn_interface (fa3-neo) is not installed; falling back to triton prefill backend "
+ "for neo_chat_moe. Install fa3-neo or set LIGHTLLM_NEO_PREFILL_TRITON_BACKEND=1 to silence "
+ "this warning."
+ )
+ _USE_TRITON_PREFILL = True
+
+
+class NeoChatMOETransformerLayerInfer(Qwen3MOETransformerLayerInfer):
+ def __init__(self, data_type, network_config):
+ super().__init__(data_type, network_config)
+ return
+
+ def _bind_attention(self):
+ self._context_attention_kernel = self._context_attention_kernel
+ self._token_attention_kernel = self._token_attention_kernel
+ return
+
+ def _get_qkv(self, input, infer_state: NeoChatInferStateInfo, layer_weight: NeoChatMOETransformerLayerWeight):
+ input = input.view(-1, self.embed_dim_)
+
+ qkv = layer_weight.qkv_proj.mm(input)
+ q, cache_kv = qkv.split(
+ [self.tp_q_head_num_ * self.head_dim_, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], dim=-1
+ )
+ q = q.view(q.shape[0], self.tp_q_head_num_, self.head_dim_)
+ q_t, q_hw = q.chunk(2, dim=-1)
+
+ cache_kv = cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+ k = cache_kv[:, : self.tp_k_head_num_, :]
+ v = cache_kv[:, self.tp_k_head_num_ :, :]
+ k_t, k_hw = k.chunk(2, dim=-1)
+
+ q_t_2d = q_t.reshape(q.shape[0], -1)
+ q_hw_2d = q_hw.reshape(q.shape[0], -1)
+ k_t_2d = k_t.reshape(k.shape[0], -1)
+ k_hw_2d = k_hw.reshape(k.shape[0], -1)
+ layer_weight.qk_norm_weight_(q_t_2d, k_t_2d, eps=self.eps_)
+ layer_weight.qk_hw_norm_weight_(q_hw_2d, k_hw_2d, eps=self.eps_)
+
+ q_t = q_t_2d.view(q.shape[0], self.tp_q_head_num_, self.head_dim_ // 2)
+ q_hw = q_hw_2d.view(q.shape[0], self.tp_q_head_num_, self.head_dim_ // 2)
+ q_h, q_w = q_hw.chunk(2, dim=-1)
+
+ k_t = k_t_2d.view(k.shape[0], self.tp_k_head_num_, self.head_dim_ // 2)
+ k_hw = k_hw_2d.view(k.shape[0], self.tp_k_head_num_, self.head_dim_ // 2)
+ k_h, k_w = k_hw.chunk(2, dim=-1)
+
+ rotary_emb_fwd(
+ q_t,
+ k_t,
+ infer_state.position_cos,
+ infer_state.position_sin,
+ )
+ rotary_emb_fwd(
+ q_h,
+ k_h,
+ infer_state.position_cos_h,
+ infer_state.position_sin_h,
+ )
+ rotary_emb_fwd(
+ q_w,
+ k_w,
+ infer_state.position_cos_w,
+ infer_state.position_sin_w,
+ )
+
+ q = torch.cat([q_t, q_h, q_w], dim=-1)
+ q = q.reshape(q.shape[0], -1)
+
+ k = torch.cat([k_t, k_h, k_w], dim=-1)
+ cache_kv = torch.cat([k, v], dim=1)
+ return q, cache_kv
+
+ def _context_attention_kernel(
+ self, q, kv, infer_state: NeoChatInferStateInfo, layer_weight, out=None
+ ) -> torch.Tensor:
+ if _USE_TRITON_PREFILL:
+ o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
+ kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+ context_attention_fwd_neo(
+ q.view(-1, self.tp_q_head_num_, self.head_dim_),
+ kv[:, 0 : self.tp_k_head_num_, :],
+ kv[:, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :],
+ o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
+ infer_state.b_req_idx,
+ infer_state.b_q_start_loc,
+ infer_state.b_seq_len,
+ infer_state.b_ready_cache_len,
+ infer_state.max_q_seq_len,
+ infer_state.req_manager.req_to_token_indexs,
+ infer_state.b_image_token_end,
+ )
+ return o_tensor
+
+ _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+ _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+
+ att_control = AttControl()
+ att_control.image_token_tag = getattr(infer_state, "b_image_token_tag", None)
+
+ o_tensor = infer_state.prefill_att_state.prefill_att(
+ q=_q,
+ k=_k,
+ v=_v,
+ att_control=att_control,
+ alloc_func=self.alloc_tensor,
+ )
+ return o_tensor.view(q.shape)
+
+ def _token_attention_kernel(
+ self,
+ q: torch.Tensor,
+ infer_state: NeoChatInferStateInfo,
+ layer_weight: NeoChatMOETransformerLayerWeight,
+ ) -> torch.Tensor:
+ _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_)
+ _q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+ att_control = AttControl()
+ # att_control.mla_decode_dict["softmax_scale"] = 1.0 / (self.head_dim_ ** 0.5)
+ o_tensor = infer_state.decode_att_state.decode_att(
+ q=_q, k=_k, v=_v, att_control=att_control, alloc_func=self.alloc_tensor
+ )
+ o_tensor = o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_)[:, :, : self.head_dim_].contiguous()
+ return o_tensor
diff --git a/lightllm/models/neo_chat_moe/layer_weights/__init__.py b/lightllm/models/neo_chat_moe/layer_weights/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat_moe/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/neo_chat_moe/layer_weights/pre_and_post_layer_weight.py
new file mode 100644
index 0000000000..4b0eae91c3
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/layer_weights/pre_and_post_layer_weight.py
@@ -0,0 +1,23 @@
+import torch
+import numpy as np
+from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight
+
+# add key: language_model.xxx -> xxx
+# only change keys at PreAndPostLayerWeight load, TransformLayerWeight is correct now
+def rename_weight_keys(weights):
+ prefix = "language_model."
+ keys = list(weights.keys())
+ for k in keys:
+ if prefix in k:
+ weights[k.replace(prefix, "")] = weights.pop(k)
+
+
+class NeoChatMOEPreAndPostLayerWeight(Qwen2PreAndPostLayerWeight):
+ def __init__(self, data_type, network_config):
+ super().__init__(data_type, network_config)
+ return
+
+ def load_hf_weights(self, weights):
+ rename_weight_keys(weights)
+ super().load_hf_weights(weights)
+ return
diff --git a/lightllm/models/neo_chat_moe/layer_weights/transformer_layer_weight.py b/lightllm/models/neo_chat_moe/layer_weights/transformer_layer_weight.py
new file mode 100644
index 0000000000..d4f985db45
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,48 @@
+from lightllm.models.qwen3_moe.layer_weights.transformer_layer_weight import Qwen3MOETransformerLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+ QKRMSNORMWeight,
+ ROWMMWeight,
+ RMSNormWeight,
+)
+
+
+class NeoChatMOETransformerLayerWeight(Qwen3MOETransformerLayerWeight):
+ def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+ self._is_merge_kv = network_config.get("merge_kv", True)
+ super().__init__(layer_num, data_type, network_config, quant_cfg)
+ return
+
+ def _init_weight_names(self):
+ super()._init_weight_names()
+ self._q_norm_hw_name = f"model.layers.{self.layer_num_}.self_attn.q_norm_hw.weight"
+ self._k_norm_hw_name = f"model.layers.{self.layer_num_}.self_attn.k_norm_hw.weight"
+
+ def _init_qkv(self):
+ super()._init_qkv()
+
+ def _init_norm(self):
+ hidden_size = self.network_config_["hidden_size"]
+ self.att_norm_weight_ = RMSNormWeight(
+ dim=hidden_size,
+ weight_name=self._att_norm_weight_name,
+ data_type=self.data_type_,
+ )
+ self.ffn_norm_weight_ = RMSNormWeight(
+ dim=hidden_size,
+ weight_name=self._ffn_norm_weight_name,
+ data_type=self.data_type_,
+ )
+
+ self.qk_norm_weight_ = QKRMSNORMWeight(
+ dim=self.head_dim // 2,
+ q_weight_name=self._q_norm_name,
+ k_weight_name=self._k_norm_name,
+ data_type=self.data_type_,
+ )
+
+ self.qk_hw_norm_weight_ = QKRMSNORMWeight(
+ dim=self.head_dim // 2,
+ q_weight_name=self._q_norm_hw_name,
+ k_weight_name=self._k_norm_hw_name,
+ data_type=self.data_type_,
+ )
diff --git a/lightllm/models/neo_chat_moe/model.py b/lightllm/models/neo_chat_moe/model.py
new file mode 100644
index 0000000000..def6d42c15
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/model.py
@@ -0,0 +1,204 @@
+import os
+import json
+from lightllm.common.build_utils import repair_config
+from lightllm.models.registry import ModelRegistry, llm_model_type_is
+from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo
+from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer
+from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer
+from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight
+from lightllm.models.qwen2_vl.model import QWen2VLTokenizer
+from lightllm.models.qwen3.model import Qwen3TpPartModel
+from lightllm.server.core.objs import SamplingParams
+from lightllm.models.qwen3_moe.model import Qwen3MOEModel
+from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem
+from lightllm.models.neo_chat_moe.vision_process import smart_resize
+from lightllm.models.internvl.model import InternvlTokenizer
+from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
+from lightllm.models.neo_chat_moe.layer_infer.transformer_layer_infer import NeoChatMOETransformerLayerInfer
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight
+from lightllm.models.neo_chat_moe.layer_weights.pre_and_post_layer_weight import NeoChatMOEPreAndPostLayerWeight
+from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer
+from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo
+from lightllm.common.basemodel.attention import (
+ get_prefill_att_backend_class,
+ get_decode_att_backend_class,
+ BaseAttBackend,
+)
+
+IMG_START_TOKEN = "![]()
"
+IMG_END_TOKEN = ""
+IMG_TOKEN = ""
+AUDIO_START_TOKEN = ""
+
+
+class NeoChatTokenizer(BaseMultiModalTokenizer):
+ def __init__(self, tokenizer, model_cfg, **kwargs):
+ super().__init__(tokenizer)
+ self.tokenizer = tokenizer
+ self.min_pixel = model_cfg.get("vision_config").get("min_pixels")
+ self.max_pixel = model_cfg.get("vision_config").get("max_pixels")
+ self.patch_size = model_cfg.get("vision_config").get("patch_size")
+ self.downsample_ratio = model_cfg.get("vision_config").get("downsample_ratio")
+
+ self.image_token_id = model_cfg.get("image_token_id")
+ self.image_start_tag = IMG_START_TOKEN
+ self.image_start_id = tokenizer.convert_tokens_to_ids(self.image_start_tag)
+ self.image_end_tag = IMG_END_TOKEN
+ self.image_end_id = tokenizer.convert_tokens_to_ids(self.image_end_tag)
+ self.image_tag = IMG_TOKEN
+
+ def init_imageitem_extral_params(
+ self, img: ImageItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+ ):
+ img.extra_params["min_pixels"] = (
+ sampling_params.min_pixels if sampling_params.min_pixels > 0 else self.min_pixel
+ )
+ img.extra_params["max_pixels"] = (
+ sampling_params.max_pixels if sampling_params.max_pixels > 0 else self.max_pixel
+ )
+ assert (
+ img.extra_params["min_pixels"] <= img.extra_params["max_pixels"]
+ ), "min_pixels should be less than or equal to max_pixels"
+ return
+
+ def init_audioitem_extral_params(
+ self, audio: AudioItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+ ):
+ raise NotImplementedError
+
+ def get_audio_token_length(self, audio: AudioItem):
+ raise NotImplementedError
+
+ def get_image_token_length(self, img: ImageItem):
+ width, height = img.image_w, img.image_h
+ resized_height, resized_width = smart_resize(
+ height=height,
+ width=width,
+ factor=int(self.patch_size // self.downsample_ratio),
+ min_pixels=img.extra_params.get("min_pixels", self.min_pixel),
+ max_pixels=img.extra_params.get("max_pixels", self.max_pixel),
+ )
+ grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
+ token_num = int((grid_h * grid_w) * (self.downsample_ratio ** 2))
+ # 这里的grid_h和grid_w需要* self.downsample_ratio么？再仔细看下代码
+ img.grid_thwd = (1, int(grid_h * self.downsample_ratio), int(grid_w * self.downsample_ratio), 1 - token_num)
+ return token_num
+
+ def get_image_token_length_by_size(self, width: int, height: int):
+ assert (
+ width > 0 and height > 0 and width % self.patch_size == 0 and height % self.patch_size == 0
+ ), "width and height should be greater than 0 and divisible by patch_size"
+
+ grid_h, grid_w = height // self.patch_size, width // self.patch_size
+ token_num = int((grid_h * grid_w) * (self.downsample_ratio ** 2))
+ return token_num
+
+ # only change the impl of the encode func:
+ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
+ # TEXTTEXTTEXT --> TEXT![]()
TEXT![]()
TEXT
+ image_tokens = IMG_START_TOKEN + IMG_END_TOKEN
+ if multimodal_params is None:
+ add_special_tokens = kwargs.get("add_special_tokens", True)
+ return self.tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
+ image_count = len(multimodal_params.images)
+ if not kwargs.get("already_tokenized", False):
+ prompt = prompt.replace(IMG_TOKEN, image_tokens, image_count)
+ origin_ids = self.tokenizer.encode(prompt, add_special_tokens=kwargs["add_special_tokens"])
+ else:
+ origin_ids = prompt
+ # ![]()
--> ![]()
id,id+1...id+num
+ input_ids = []
+ image_id = 0
+ start_idx = 0
+ while True:
+ try:
+ start_idx = origin_ids.index(self.image_start_id)
+ if start_idx + 1 >= len(origin_ids):
+ break
+ if origin_ids[start_idx + 1] == self.image_end_id:
+ input_ids.extend(origin_ids[: start_idx + 1])
+ token_id = multimodal_params.images[image_id].token_id
+ token_num = multimodal_params.images[image_id].token_num
+ multimodal_params.images[image_id].start_idx = len(input_ids)
+ input_ids.extend(range(token_id, token_id + token_num))
+ input_ids.append(self.image_end_id)
+ origin_ids = origin_ids[start_idx + 2 :]
+ image_id += 1
+ else:
+ raise ValueError("image token error")
+ except ValueError:
+ break
+ input_ids.extend(origin_ids)
+ return input_ids
+
+ def _build_t2i_query(self, msg, thinking_content=""):
+ prompt = self.tokenizer.apply_chat_template(
+ conversation=[{"role": "user", "content": msg}], tokenize=False, add_generation_prompt=True
+ )
+
+ return prompt + thinking_content + IMG_START_TOKEN
+
+ def fix_prompt(self, prompt: str, img_len: int):
+ prompt_img_len = prompt.count(IMG_TOKEN)
+ assert prompt_img_len <= img_len, f"not enough images provided, need {prompt_img_len}, given {img_len}"
+ if prompt_img_len < img_len:
+ return f"{IMG_TOKEN}\n" * (img_len - prompt_img_len) + prompt
+ return prompt
+
+ def get_query_for_it2i(self, prompt: str):
+ image_len = prompt.count(IMG_TOKEN)
+ # query_condition = self._build_t2i_query(prompt, thinking_content="\n\n\n\n")
+ query_condition = prompt + IMG_START_TOKEN if not prompt.endswith(IMG_START_TOKEN) else prompt
+ query_text_uncondition = self._build_t2i_query(IMG_TOKEN * image_len)
+ question_img_uncondition = self._build_t2i_query("")
+ return query_condition, query_text_uncondition, question_img_uncondition
+
+ def get_query_for_t2i(self, prompt: str, input_image_num: int = 0):
+ # prompt is already applied
+ image_len = prompt.count(IMG_TOKEN)
+ query_condition = prompt + IMG_START_TOKEN if not prompt.endswith(IMG_START_TOKEN) else prompt
+ query_uncondition = self._build_t2i_query(
+ IMG_TOKEN * input_image_num, thinking_content=IMG_TOKEN * (image_len - input_image_num)
+ )
+ return query_condition, query_uncondition
+
+
+@ModelRegistry(["neo_chat"], is_multimodal=True, condition=llm_model_type_is("qwen3_moe"))
+class NeoTpMOEPartModel(Qwen3MOEModel):
+
+ pre_layer_infer_class = LlamaMultimodalPreLayerInfer
+ transformer_layer_infer_class = NeoChatMOETransformerLayerInfer
+
+ pre_and_post_weight_class = NeoChatMOEPreAndPostLayerWeight
+ transformer_weight_class = NeoChatMOETransformerLayerWeight
+
+ infer_state_class = NeoChatInferStateInfo
+
+ def __init__(self, kvargs):
+ super().__init__(kvargs)
+ return
+
+ def _init_inferstate_cls(self):
+ pass
+
+ def _init_att_backend(self):
+ self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0, priority_list=["fa3"])(
+ model=self
+ )
+ self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0, priority_list=["fa3"])(
+ model=self
+ )
+
+ def _init_config(self):
+ with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
+ all_config = json.load(json_file)
+ self.config = all_config["llm_config"]
+ # rename keys
+ repair_config(self.config, same_names=["num_attention_heads", "n_head"])
+ repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"])
+ repair_config(self.config, same_names=["num_hidden_layers", "n_layer"])
+ if self.finetune_config:
+ self.config["vocab_size"] = self.finetune_config.vocab_size
+ return
diff --git a/lightllm/models/neo_chat_moe/neo_visual.py b/lightllm/models/neo_chat_moe/neo_visual.py
new file mode 100644
index 0000000000..e0999079b0
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/neo_visual.py
@@ -0,0 +1,281 @@
+import os
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from typing import List
+from io import BytesIO
+import torch.nn as nn
+from transformers.activations import ACT2FN
+from safetensors import safe_open
+from lightllm.server.multimodal_params import ImageItem
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from lightllm.models.neo_chat_moe.vision_process import load_image_native
+from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data
+
+
+def apply_rotary_emb_1d(
+ x: torch.Tensor,
+ cos_cached: torch.Tensor,
+ sin_cached: torch.Tensor,
+ positions: torch.Tensor,
+):
+ """对输入张量的一部分应用1D RoPE。"""
+ # x: (..., seq_len, dim_part)
+ # positions: (..., seq_len)
+ # cos_cached: (max_pos, dim_part / 2)
+ cos_cached = cos_cached.to(device=positions.device)
+ sin_cached = sin_cached.to(device=positions.device)
+
+ cos = cos_cached[positions] # Shape: (positions.shape, dim_part / 2)
+ sin = sin_cached[positions] # Shape: (positions.shape, dim_part / 2)
+
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+
+ rotated_x1 = x1 * cos - x2 * sin
+ rotated_x2 = x1 * sin + x2 * cos
+
+ x_rotated = torch.empty_like(x)
+ x_rotated[..., 0::2] = rotated_x1
+ x_rotated[..., 1::2] = rotated_x2
+ return x_rotated
+
+
+def apply_2d_rotary_pos_emb(
+ x: torch.Tensor,
+ cos_cached_x: torch.Tensor,
+ sin_cached_x: torch.Tensor,
+ cos_cached_y: torch.Tensor,
+ sin_cached_y: torch.Tensor,
+ abs_positions_x: torch.Tensor,
+ abs_positions_y: torch.Tensor,
+):
+ """应用2D RoPE到输入张量x。"""
+ dim = x.shape[-1]
+ dim_half = dim // 2
+
+ # 假设我们将embedding的前半部分用于一个方向的RoPE，后半部分用于另一个方向
+ # 例如，前一半给X坐标，后一半给Y坐标 (或者反过来，但要保持一致)
+ x_part_1 = x[..., :dim_half]
+ x_part_2 = x[..., dim_half:]
+
+ # 将与 abs_positions_x 相关的旋转应用于 x_part_1
+ rotated_part_1 = apply_rotary_emb_1d(x_part_1, cos_cached_x, sin_cached_x, abs_positions_x)
+ # 将与 abs_positions_y 相关的旋转应用于 x_part_2
+ rotated_part_2 = apply_rotary_emb_1d(x_part_2, cos_cached_y, sin_cached_y, abs_positions_y)
+
+ # 将它们重新拼接起来。确保顺序与你分割时一致。
+ return torch.cat((rotated_part_1, rotated_part_2), dim=-1)
+
+
+def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None):
+ """
+ Compute patch coordinates (x, y)
+
+ Args:
+ grid_hw: (B, 2) tensor representing (H, W) per image
+ """
+ device = grid_hw.device
+ B = grid_hw.shape[0]
+
+ # Get the number of patches per image
+ H = grid_hw[:, 0]
+ W = grid_hw[:, 1]
+ N = H * W
+ N_total = N.sum()
+
+ # Create the batch index for each patch (B x patch count)
+ patch_to_sample = torch.repeat_interleave(torch.arange(B, device=device), N) # (N_total,)
+
+ # Generate intra-image patch index (row-major order)
+ patch_id_within_image = torch.arange(N_total, device=device)
+ patch_id_within_image = (
+ patch_id_within_image
+ - torch.cumsum(torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0)[patch_to_sample]
+ )
+
+ # Get H/W for each patch according to its image
+ W_per_patch = W[patch_to_sample]
+ abs_x = patch_id_within_image % W_per_patch
+ abs_y = patch_id_within_image // W_per_patch
+
+ return abs_x, abs_y
+
+
+class NeoVisionTransformerPretrainedModel(nn.Module):
+ def __init__(
+ self,
+ kvargs,
+ hidden_size: int = 1024,
+ llm_hidden_size: int = 2048,
+ downsample_ratio: float = 0.5,
+ patch_size: int = 16,
+ num_channels: int = 3,
+ max_position_embeddings_vision: int = 10000,
+ rope_theta_vision: float = 10000.0,
+ min_pixels: int = 65536,
+ max_pixels: int = 2408448,
+ **kwargs,
+ ):
+ super().__init__()
+ self.weight_dir = kvargs["weight_dir"]
+ self.data_type = kvargs.get("data_type", "bfloat16")
+ self.embed_dim = hidden_size
+ self.llm_hidden_size = llm_hidden_size
+ self.patch_size = patch_size
+ self.num_channels = num_channels
+ self.downsample_ratio = downsample_ratio
+ self.downsample_factor = int(1 / downsample_ratio)
+ self.max_position_embeddings_vision = max_position_embeddings_vision
+ self.rope_theta_vision = rope_theta_vision
+ self.rope_dim_part = self.embed_dim // 2
+ self.min_pixels = min_pixels
+ self.max_pixels = max_pixels
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=num_channels, out_channels=self.embed_dim, kernel_size=patch_size, stride=patch_size
+ )
+
+ self.dense_embedding = nn.Conv2d(
+ in_channels=self.embed_dim,
+ out_channels=self.llm_hidden_size,
+ kernel_size=self.downsample_factor,
+ stride=self.downsample_factor,
+ )
+ self.gelu = nn.GELU()
+
+ self.repe_dim_part = self.embed_dim // 2
+ self.cos_x, self.sin_x = self.precompute_rope_freqs_sincos()
+ self.cos_y, self.sin_y = self.precompute_rope_freqs_sincos()
+ self._init_datatype()
+
+ def _init_datatype(self):
+ if isinstance(self.data_type, torch.dtype):
+ return
+ if self.data_type in ["fp16", "float16"]:
+ self.data_type = torch.float16
+ elif self.data_type in ["bf16", "bfloat16"]:
+ self.data_type = torch.bfloat16
+ elif self.data_type in ["fp32", "float32"]:
+ self.data_type = torch.float32
+ else:
+ raise ValueError(f"Unsupport datatype {self.data_type}!")
+ return
+
+ def load_model(self, weight_dir):
+ bin_weight_files = [file_ for file_ in os.listdir(weight_dir) if file_.endswith(".bin")]
+ if bin_weight_files:
+ weight_dict = {}
+ for file_ in bin_weight_files:
+ f = torch.load(os.path.join(weight_dir, file_), "cpu")
+ for k, v in f.items():
+ if "vision_model" in k and "fm_modules" not in k:
+ weight_dict[k[len("vision_model.embeddings.") :]] = v
+ else:
+ hf_weight_files = [file_ for file_ in os.listdir(weight_dir) if file_.endswith(".safetensors")]
+ weight_dict = {}
+ for file_ in hf_weight_files:
+ f = safe_open(os.path.join(weight_dir, file_), "pt", "cpu")
+ for k in f.keys():
+ if "vision_model" in k and "fm_modules" not in k:
+ weight_dict[k[len("vision_model.embeddings.") :]] = f.get_tensor(k)
+ self.load_state_dict(weight_dict)
+
+ def precompute_rope_freqs_sincos(self):
+ inv_freq = 1.0 / (
+ self.rope_theta_vision ** (torch.arange(0, self.rope_dim_part, 2).float() / self.rope_dim_part)
+ )
+ t = torch.arange(self.max_position_embeddings_vision).type_as(inv_freq)
+ freqs = torch.outer(t, inv_freq)
+ return torch.cos(freqs), torch.sin(freqs)
+
+ def _apply_2d_rotary_pos_emb(self, patch_embeds, grid_hw):
+ """
+ Apply 2D Rotary Position Embedding to the patch embeddings.
+ """
+ abs_pos_x, abs_pos_y = build_abs_positions_from_grid_hw(grid_hw, device=patch_embeds.device)
+ embeddings = apply_2d_rotary_pos_emb(
+ patch_embeds.to(torch.float32), # RoPE calculations are often more stable in float32
+ self.cos_x,
+ self.sin_x,
+ self.cos_y,
+ self.sin_y,
+ abs_pos_x,
+ abs_pos_y,
+ ).to(self.patch_embedding.weight.dtype)
+ return embeddings
+
+ def forward(self, pixel_values: torch.Tensor, grid_hw: torch.Tensor) -> torch.Tensor:
+ pixel_values = pixel_values.view(
+ -1,
+ 3,
+ self.patch_size,
+ self.patch_size,
+ )
+ patch_embeds = self.gelu(self.patch_embedding(pixel_values)).view(-1, self.embed_dim)
+ patch_embeds = self._apply_2d_rotary_pos_emb(patch_embeds, grid_hw)
+ assert (grid_hw[:, 0] * grid_hw[:, 1]).sum() == patch_embeds.shape[
+ 0
+ ], "Grid size and patch embeds size mismatch."
+
+ patches_list = []
+ cur_position = 0
+ for i in range(grid_hw.shape[0]):
+ h, w = grid_hw[i]
+ patches_per_img = patch_embeds[cur_position : cur_position + h * w].view(h, w, -1).unsqueeze(0)
+ patches_per_img = self.dense_embedding(patches_per_img.permute(0, 3, 1, 2))
+ patches_per_img = patches_per_img.permute(0, 2, 3, 1)
+ patches_list.append(patches_per_img.view(-1, patches_per_img.shape[-1]))
+ cur_position += h * w
+
+ embeddings = torch.cat(patches_list, dim=0) # (N_total // downsample_factor**2, C)
+ assert cur_position == patch_embeds.shape[0]
+ assert embeddings.shape[0] == int(patch_embeds.shape[0] / self.downsample_factor ** 2)
+
+ return embeddings
+
+ def encode(self, images: List[ImageItem]):
+ img_tensors = []
+ valid_ids = []
+ valid_id = 0
+ img_grids = []
+ uuids = []
+
+ for i, img in enumerate(images):
+ if isinstance(img, ImageItem):
+ uuids.append(img.uuid)
+ image_data = read_shm(get_shm_name_data(img.uuid))
+ image_data = Image.open(BytesIO(image_data))
+ # a = img.extra_params["min_pixels"]
+ # b = img.extra_params["max_pixels"]
+ # print(f"self.min_pixels is {a} ,max_pixelx is {b}")
+ pixel_values, image_grid_hw = load_image_native(
+ image_data,
+ patch_size=self.patch_size,
+ downsample_ratio=self.downsample_ratio,
+ min_pixels=img.extra_params.get("min_pixels", self.min_pixels),
+ max_pixels=img.extra_params.get("max_pixels", self.max_pixels),
+ )
+ img_tensors.append(pixel_values)
+ img_grids.append(image_grid_hw)
+ else:
+ raise Exception("Unsupport input types: {} for {}".format(type(img), img))
+
+ # must devide merge_length
+ cur_num = int(img_tensors[-1].shape[0] * (self.downsample_ratio ** 2))
+ valid_ids.append([valid_id, valid_id + cur_num])
+ valid_id += cur_num
+
+ if len(img_tensors) <= 0:
+ return None
+
+ imgs = torch.cat(img_tensors, dim=0)
+ grid_hw = torch.cat(img_grids, dim=0)
+
+ pixel_values = imgs.to("cuda", dtype=self.data_type, non_blocking=True)
+ image_grid_hw = grid_hw.to("cuda", non_blocking=True)
+
+ all_img_embeds = self.forward(pixel_values, grid_hw=image_grid_hw)
+
+ return all_img_embeds, uuids, valid_ids
diff --git a/lightllm/models/neo_chat_moe/triton_kernel/__init__.py b/lightllm/models/neo_chat_moe/triton_kernel/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/neo_chat_moe/triton_kernel/context_attention_fwd_neo.py b/lightllm/models/neo_chat_moe/triton_kernel/context_attention_fwd_neo.py
new file mode 100644
index 0000000000..1016e5a3c1
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/triton_kernel/context_attention_fwd_neo.py
@@ -0,0 +1,413 @@
+import math
+import torch
+import triton
+import triton.language as tl
+
+from lightllm.utils.device_utils import is_tesla
+
+
+@triton.jit
+def _fwd_kernel(
+ Q,
+ K,
+ V,
+ sm_scale,
+ Out,
+ B_Start_Loc,
+ B_Seqlen,
+ Req_to_tokens,
+ B_req_idx,
+ stride_qbs,
+ stride_qh,
+ stride_qd,
+ stride_kbs,
+ stride_kh,
+ stride_kd,
+ stride_vbs,
+ stride_vh,
+ stride_vd,
+ stride_obs,
+ stride_oh,
+ stride_od,
+ stride_req_to_tokens_b,
+ stride_req_to_tokens_s,
+ kv_group_num,
+ b_prompt_cache_len,
+ b_image_token_end,
+ H: tl.constexpr,
+ QK_HEAD_DIM: tl.constexpr,
+ V_HEAD_DIM: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+):
+ start_m = tl.program_id(0)
+ cur_bh = tl.program_id(1)
+ cur_batch = cur_bh // H
+ cur_head = cur_bh % H
+
+ cur_kv_head = cur_head // kv_group_num
+
+ cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+ prompt_cache_len = tl.load(b_prompt_cache_len + cur_batch)
+ total_len = tl.load(B_Seqlen + cur_batch)
+ cur_batch_seq_len = total_len - prompt_cache_len # NEW len
+ cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+
+ block_start_loc = BLOCK_M * start_m
+ if block_start_loc >= cur_batch_seq_len:
+ return
+
+ offs_n = tl.arange(0, BLOCK_N)
+ offs_d_qk = tl.arange(0, QK_HEAD_DIM)
+ offs_d_v = tl.arange(0, V_HEAD_DIM)
+ offs_m = block_start_loc + tl.arange(0, BLOCK_M)
+
+ # Q pointers
+ off_q = (
+ (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
+ + cur_head * stride_qh
+ + offs_d_qk[None, :] * stride_qd
+ )
+
+ q_valid = offs_m < cur_batch_seq_len
+ q = tl.load(Q + off_q, mask=q_valid[:, None], other=0.0)
+
+ # online softmax state
+ m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+ l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+ acc = tl.zeros([BLOCK_M, V_HEAD_DIM], dtype=tl.float32)
+ # absolute q positions in the request
+ q_pos = prompt_cache_len + offs_m # [M]
+ q_image_end = tl.load(b_image_token_end + cur_batch_in_all_start_index + offs_m, mask=q_valid, other=0)
+
+ causal_end = tl.minimum(prompt_cache_len + block_start_loc + BLOCK_M, total_len)
+ block_image_end = tl.minimum(tl.max(q_image_end, axis=0), total_len)
+ block_end_loc = tl.maximum(causal_end, block_image_end)
+
+ for start_n in range(0, block_end_loc, BLOCK_N):
+ start_n = tl.multiple_of(start_n, BLOCK_N)
+
+ k_pos = start_n + offs_n # [N]
+ k_valid = k_pos < block_end_loc
+
+ # map logical pos -> mem_index (for K/V)
+ kv_loc = tl.load(
+ Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos,
+ mask=k_valid,
+ other=0,
+ ).to(tl.int64)
+
+ # load K
+ off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d_qk[:, None] * stride_kd
+ k = tl.load(K + off_k, mask=k_valid[None, :], other=0.0)
+ qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+ qk += tl.dot(q, k)
+
+ causal_mask = q_pos[:, None] >= k_pos[None, :]
+ image_mask = k_pos[None, :] < q_image_end[:, None]
+ mask = (causal_mask | image_mask) & k_valid[None, :]
+ qk = tl.where(mask, qk * sm_scale, -1.0e8)
+
+ # online softmax
+ m_ij = tl.maximum(m_i, tl.max(qk, 1))
+ qk -= m_ij[:, None]
+ p = tl.math.exp2(qk)
+ l_ij = tl.sum(p, 1)
+
+ alpha = tl.math.exp2(m_i - m_ij)
+ l_i = l_i * alpha + l_ij
+ acc = acc * alpha[:, None]
+
+ # load V
+ off_v = kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d_v[None, :] * stride_vd
+ v = tl.load(V + off_v, mask=k_valid[:, None], other=0.0)
+
+ p = p.to(v.dtype)
+ acc = tl.dot(p, v, acc)
+
+ m_i = m_ij
+
+ acc = acc / l_i[:, None]
+
+ off_o = (
+ (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
+ + cur_head * stride_oh
+ + offs_d_v[None, :] * stride_od
+ )
+ tl.store(Out + off_o, acc, mask=q_valid[:, None])
+
+
+@torch.no_grad()
+def context_attention_fwd_neo(
+ q,
+ k,
+ v,
+ o,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ max_input_len,
+ req_to_token_indexs,
+ b_image_token_end,
+):
+ BLOCK_M = 128 if not is_tesla() else 64
+
+ Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+ assert Lq == Lk and Lk == Lv
+ assert Lk in {16, 32, 64, 128, 256}
+ sm_scale = 1.0 / (Lq ** 0.5) * 1.4426950408889634
+
+ batch, head = b_seq_len.shape[0], q.shape[1]
+ kv_group_num = q.shape[1] // k.shape[1]
+
+ grid = lambda meta: (triton.cdiv(max_input_len, meta["BLOCK_M"]), batch * head, 1)
+
+ BLOCK_N = BLOCK_M
+ num_warps = 4 if Lk <= 64 else 8
+ num_stages = 1
+
+ _fwd_kernel[grid](
+ q,
+ k,
+ v,
+ sm_scale,
+ o,
+ b_start_loc,
+ b_seq_len,
+ req_to_token_indexs,
+ b_req_idx,
+ q.stride(0),
+ q.stride(1),
+ q.stride(2),
+ k.stride(0),
+ k.stride(1),
+ k.stride(2),
+ v.stride(0),
+ v.stride(1),
+ v.stride(2),
+ o.stride(0),
+ o.stride(1),
+ o.stride(2),
+ req_to_token_indexs.stride(0),
+ req_to_token_indexs.stride(1),
+ kv_group_num=kv_group_num,
+ b_prompt_cache_len=b_prompt_cache_len,
+ b_image_token_end=b_image_token_end,
+ H=head,
+ QK_HEAD_DIM=Lk,
+ V_HEAD_DIM=Lk,
+ BLOCK_M=BLOCK_M,
+ BLOCK_N=BLOCK_N,
+ num_warps=num_warps,
+ num_stages=num_stages,
+ )
+
+
+def reference_attention(
+ q,
+ k,
+ v,
+ b_image_token_end,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ req_to_token_indexs,
+):
+ device = q.device
+ dtype = q.dtype
+ sum_q, Hq, D = q.shape
+ Hk = k.shape[1]
+ kv_group_num = Hq // Hk
+
+ batch = b_seq_len.shape[0]
+ out = torch.empty_like(q)
+ scale = 1.0 / math.sqrt(D)
+
+ for b in range(batch):
+ req = int(b_req_idx[b].item())
+ total_len = int(b_seq_len[b].item())
+ prompt_len = int(b_prompt_cache_len[b].item())
+ new_len = total_len - prompt_len
+
+ q_start = int(b_start_loc[b].item())
+ q_blk = q[q_start : q_start + new_len] # [M, Hq, D]
+ image_end_new = b_image_token_end[q_start : q_start + new_len].to(torch.int64) # [M]
+
+ # gather K/V for full request by logical pos -> mem_index
+ token_locs = req_to_token_indexs[req, :total_len].to(torch.int64) # [L]
+ k_blk = k[token_locs] # [L, Hk, D]
+ v_blk = v[token_locs] # [L, Hk, D]
+
+ # expand kv heads to q heads (GQA)
+ k_hq = k_blk.repeat_interleave(kv_group_num, dim=1) # [L, Hq, D]
+ v_hq = v_blk.repeat_interleave(kv_group_num, dim=1) # [L, Hq, D]
+
+ # positions
+ q_pos = torch.arange(prompt_len, total_len, device=device, dtype=torch.int64) # [M]
+ k_pos = torch.arange(0, total_len, device=device, dtype=torch.int64) # [L]
+
+ # build allow mask:
+ # causal always
+ allow = k_pos[None, :] <= q_pos[:, None]
+ allow = allow | (k_pos[None, :] < image_end_new[:, None])
+
+ # scores: [Hq, M, L]
+ q_t = q_blk.permute(1, 0, 2).to(torch.float32) # [Hq, M, D]
+ k_t = k_hq.permute(1, 2, 0).to(torch.float32) # [Hq, D, L]
+ scores = torch.matmul(q_t, k_t) * scale # [Hq, M, L]
+
+ neg = torch.tensor(-1.0e9, device=device, dtype=torch.float32)
+ scores = torch.where(allow[None, :, :], scores, neg)
+
+ p = torch.softmax(scores, dim=-1).to(torch.float32) # [Hq, M, L]
+ v_t = v_hq.permute(1, 0, 2).to(torch.float32) # [Hq, L, D]
+ out_hq = torch.matmul(p, v_t) # [Hq, M, D]
+ out_blk = out_hq.permute(1, 0, 2).to(dtype) # [M, Hq, D]
+
+ out[q_start : q_start + new_len] = out_blk
+
+ return out
+
+
+def make_test_case(
+ device="cuda",
+ dtype=torch.float16,
+ batch=3,
+ Hq=8,
+ Hk=4,
+ D=64,
+ seed=0,
+ base_index=50000,
+):
+ torch.manual_seed(seed)
+
+ # prompt (cached) len and new len
+ prompt_lens = torch.randint(low=2, high=8, size=(batch,), device=device)
+ new_lens = torch.randint(low=1, high=8, size=(batch,), device=device)
+ total_lens = (prompt_lens + new_lens).to(torch.int32)
+
+ max_total_len = int(total_lens.max().item())
+ max_new_len = int(new_lens.max().item())
+
+ # packed q start
+ b_start_loc = torch.zeros((batch,), device=device, dtype=torch.int32)
+ cur = 0
+ for b in range(batch):
+ b_start_loc[b] = cur
+ cur += int(new_lens[b].item())
+ sum_q = cur
+
+ b_seq_len = total_lens
+ b_prompt_cache_len = prompt_lens.to(torch.int32)
+
+ # one req per batch
+ num_req = batch
+ b_req_idx = torch.arange(batch, device=device, dtype=torch.int32)
+
+ # global KV space large, indices not small
+ sum_kv = int(total_lens.sum().item())
+ kv_size = base_index + sum_kv + 1024
+ pool = torch.randperm(kv_size - base_index, device=device, dtype=torch.int64)[:sum_kv] + base_index
+
+ # Req_to_tokens [num_req, max_total_len]
+ req_to_token_indexs = torch.zeros((num_req, max_total_len), device=device, dtype=torch.int32)
+ p = 0
+ for r in range(num_req):
+ L = int(total_lens[r].item())
+ req_to_token_indexs[r, :L] = pool[p : p + L].to(torch.int32)
+ p += L
+
+ b_image_token_end = torch.zeros((sum_q,), device=device, dtype=torch.int32)
+ for b in range(batch):
+ M = int(new_lens[b].item())
+ P = int(prompt_lens[b].item())
+ start = int(b_start_loc[b].item())
+
+ # make one repeated block inside NEW part to simulate image tokens
+ if M >= 4 and torch.rand((), device=device).item() > 0.3:
+ s = int(torch.randint(0, M - 2, (1,), device=device).item())
+ e = min(M, s + 3)
+ b_image_token_end[start + s : start + e] = P + e
+
+ q = torch.randn((sum_q, Hq, D), device=device, dtype=dtype)
+ k = torch.randn((kv_size, Hk, D), device=device, dtype=dtype)
+ v = torch.randn((kv_size, Hk, D), device=device, dtype=dtype)
+ o = torch.empty((sum_q, Hq, D), device=device, dtype=dtype)
+
+ return (
+ q,
+ k,
+ v,
+ o,
+ b_image_token_end,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ max_new_len,
+ req_to_token_indexs,
+ )
+
+
+def check_once(device="cuda", dtype=torch.float16, seed=0):
+ (
+ q,
+ k,
+ v,
+ o,
+ b_image_token_end,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ max_new_len,
+ req_to_token_indexs,
+ ) = make_test_case(device=device, dtype=dtype, seed=seed)
+
+ context_attention_fwd_neo(
+ q,
+ k,
+ v,
+ o,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ max_new_len,
+ req_to_token_indexs,
+ b_image_token_end,
+ )
+
+ ref = reference_attention(
+ q,
+ k,
+ v,
+ b_image_token_end,
+ b_req_idx,
+ b_start_loc,
+ b_seq_len,
+ b_prompt_cache_len,
+ req_to_token_indexs,
+ )
+
+ diff = (o - ref).abs()
+ max_abs = diff.max().item()
+ denom = ref.abs().max().item() + 1e-6
+ max_rel = max_abs / denom
+
+ print(f"seed={seed}, dtype={dtype}")
+ print(f"max_abs_error = {max_abs:.6e}")
+ print(f"max_rel_error = {max_rel:.6e}")
+ print("allclose(fp16 tol)?", torch.allclose(o, ref, atol=5e-2, rtol=5e-2))
+
+
+if __name__ == "__main__":
+ if not torch.cuda.is_available():
+ print("No CUDA, skip.")
+ else:
+ torch.cuda.synchronize()
+ check_once(dtype=torch.bfloat16, seed=0)
+ check_once(dtype=torch.bfloat16, seed=1)
+ check_once(dtype=torch.bfloat16, seed=2)
diff --git a/lightllm/models/neo_chat_moe/triton_kernel/get_neo_position.py b/lightllm/models/neo_chat_moe/triton_kernel/get_neo_position.py
new file mode 100644
index 0000000000..b2e1af9f8f
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/triton_kernel/get_neo_position.py
@@ -0,0 +1,194 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _get_neo_position_triton(
+ b_image_start_idx: torch.Tensor,
+ b_image_thwd: torch.Tensor,
+ b_image_thwd_stride0: torch.Tensor,
+ b_image_nums: torch.Tensor,
+ b_image_start_num: torch.Tensor,
+ b_image_len: torch.Tensor,
+ position_ids: torch.Tensor,
+ position_ids_stride0: torch.Tensor,
+ b_ready_cache_len: torch.Tensor,
+ b_q_seq_len: torch.Tensor,
+ b_start_loc: torch.Tensor,
+ b_image_token_end: torch.Tensor,
+ BLOCK_SIZE: tl.constexpr,
+) -> torch.Tensor:
+ cur_batch = tl.program_id(0)
+ cache_len = tl.load(b_ready_cache_len + cur_batch)
+ q_seq_len = tl.load(b_q_seq_len + cur_batch)
+ image_num = tl.load(b_image_nums + cur_batch)
+ image_start_num = tl.load(b_image_start_num + cur_batch)
+ start_loc = tl.load(b_start_loc + cur_batch)
+
+ for i in range(image_num):
+ local_image_start_idx = tl.load(b_image_start_idx + image_start_num + i)
+ image_start_idx = start_loc + local_image_start_idx - cache_len
+ image_len = tl.load(b_image_len + image_start_num + i)
+ image_end = local_image_start_idx + image_len
+ # image_h = tl.load(b_image_thwd + (image_start_num + i) * b_image_thwd_stride0 + 1)
+ image_w = tl.load(b_image_thwd + (image_start_num + i) * b_image_thwd_stride0 + 2)
+
+ for j in range(0, image_len, BLOCK_SIZE):
+ off = j + tl.arange(0, BLOCK_SIZE)
+ # 目前没考虑视频，所以t 恒为 0
+ t_pos = local_image_start_idx + off * 0
+ h_pos = off // image_w
+ w_pos = off % image_w
+ tl.store(
+ b_image_token_end + off + image_start_idx,
+ image_end,
+ mask=(off < image_len)
+ & (off + local_image_start_idx - cache_len < q_seq_len)
+ & (local_image_start_idx - cache_len + off >= 0),
+ )
+ tl.store(
+ position_ids + off + image_start_idx,
+ t_pos,
+ mask=(off < image_len)
+ & (off + local_image_start_idx - cache_len < q_seq_len)
+ & (local_image_start_idx - cache_len + off >= 0),
+ )
+ tl.store(
+ position_ids + position_ids_stride0 + off + image_start_idx,
+ h_pos,
+ mask=(off < image_len)
+ & (off + local_image_start_idx - cache_len < q_seq_len)
+ & (local_image_start_idx - cache_len + off >= 0),
+ )
+ tl.store(
+ position_ids + position_ids_stride0 * 2 + off + image_start_idx,
+ w_pos,
+ mask=(off < image_len)
+ & (off + local_image_start_idx - cache_len < q_seq_len)
+ & (local_image_start_idx - cache_len + off >= 0),
+ )
+
+ for i in range(image_num):
+ local_image_start_idx = tl.load(b_image_start_idx + image_start_num + i)
+ image_len = tl.load(b_image_len + image_start_num + i)
+ image_delta = tl.load(b_image_thwd + (image_start_num + i) * b_image_thwd_stride0 + 3)
+ image_end = local_image_start_idx + image_len - cache_len
+ text_start = tl.maximum(0, image_end)
+ for j in range(text_start, q_seq_len, BLOCK_SIZE):
+ off = j + tl.arange(0, BLOCK_SIZE)
+ t_pos = tl.load(position_ids + off + start_loc, mask=(off < q_seq_len), other=0.0) + image_delta
+ h_pos = tl.load(position_ids + position_ids_stride0 + off + start_loc, mask=(off < q_seq_len), other=0.0)
+ w_pos = tl.load(
+ position_ids + position_ids_stride0 * 2 + off + start_loc, mask=(off < q_seq_len), other=0.0
+ )
+ tl.store(position_ids + off + start_loc, t_pos, mask=(off < q_seq_len))
+ tl.store(position_ids + position_ids_stride0 + off + start_loc, h_pos, mask=(off < q_seq_len))
+ tl.store(position_ids + position_ids_stride0 * 2 + off + start_loc, w_pos, mask=(off < q_seq_len))
+ return
+
+
+def get_neo_position_triton(
+ b_image_start_idx: torch.Tensor,
+ b_image_thwd: torch.Tensor,
+ b_image_nums: torch.Tensor,
+ b_image_start_num: torch.Tensor,
+ b_image_len: torch.Tensor,
+ position_ids: torch.Tensor,
+ b_ready_cache_len: torch.Tensor,
+ b_q_seq_len: torch.Tensor,
+ b_start_loc: torch.Tensor,
+ b_image_token_end: torch.Tensor,
+) -> torch.Tensor:
+
+ batch_size = b_q_seq_len.shape[0]
+ assert batch_size == b_image_nums.shape[0]
+ grid = (batch_size,)
+ BLOCK_SIZE = 64
+ _get_neo_position_triton[grid](
+ b_image_start_idx=b_image_start_idx,
+ b_image_thwd=b_image_thwd,
+ b_image_thwd_stride0=b_image_thwd.stride(0),
+ b_image_nums=b_image_nums,
+ b_image_start_num=b_image_start_num,
+ b_image_len=b_image_len,
+ position_ids=position_ids,
+ position_ids_stride0=position_ids.stride(0),
+ b_ready_cache_len=b_ready_cache_len,
+ b_q_seq_len=b_q_seq_len,
+ b_start_loc=b_start_loc,
+ b_image_token_end=b_image_token_end,
+ BLOCK_SIZE=BLOCK_SIZE,
+ )
+
+
+def test():
+ b_image_start_idx = torch.tensor([0, 0, 4], dtype=torch.int32, device="cuda")
+ b_image_thwd = torch.tensor([[1, 2, 2, -3], [1, 2, 2, -3], [1, 2, 4, -7]], dtype=torch.int32, device="cuda")
+ b_image_nums = torch.tensor([1, 2], dtype=torch.int32, device="cuda")
+ b_image_start_num = torch.tensor([0, 1], dtype=torch.int32, device="cuda")
+ b_image_len = torch.tensor([4, 4, 8], dtype=torch.int32, device="cuda")
+ position_ids = (
+ torch.tensor([0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=torch.int32, device="cuda")
+ .unsqueeze(0)
+ .expand(3, -1)
+ .contiguous()
+ )
+ b_image_token_end = torch.zeros([position_ids.size(1)], dtype=torch.int32, device="cuda")
+ position_ids[1:].zero_()
+ b_ready_cache_len = torch.tensor([0, 0], dtype=torch.int32, device="cuda")
+ b_q_seq_len = torch.tensor([7, 13], dtype=torch.int32, device="cuda")
+ b_start_loc = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
+ get_neo_position_triton(
+ b_image_start_idx,
+ b_image_thwd,
+ b_image_nums,
+ b_image_start_num,
+ b_image_len,
+ position_ids,
+ b_ready_cache_len,
+ b_q_seq_len,
+ b_start_loc,
+ b_image_token_end,
+ )
+
+ print(b_image_token_end)
+ print(position_ids)
+ # old_value = torch.cat([position_ids[:, 2:7], position_ids[:, 7 + 2 :]], dim=1)
+
+ # position_ids = (
+ # torch.tensor([2, 3, 4, 5, 6, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=torch.int32, device="cuda")
+ # .unsqueeze(0)
+ # .expand(3, -1)
+ # .contiguous()
+ # )
+ # b_ready_cache_len = torch.tensor([2, 2], dtype=torch.int32, device="cuda")
+ # b_q_seq_len = torch.tensor([5, 11], dtype=torch.int32, device="cuda")
+ # b_start_loc = torch.tensor([0, 5], dtype=torch.int32, device="cuda")
+
+ # get_neo_position_triton(
+ # b_image_start_idx,
+ # b_image_thwd,
+ # b_image_nums,
+ # b_image_start_num,
+ # b_image_len,
+ # position_ids,
+ # b_ready_cache_len,
+ # b_q_seq_len,
+ # b_start_loc,
+ # )
+
+ # print(f"old_value:\n{old_value}")
+ # print(f"position_ids:\n{position_ids}")
+ # assert torch.equal(old_value, position_ids)
+
+ """
+ tensor([[0, 0, 0, 0, 2, 3, 4, 0, 0, 0, 0, 2, 2, 2, 2, 4, 5, 6, 7, 8],
+ [0, 0, 1, 1, 2, 3, 4, 0, 0, 1, 1, 2, 2, 3, 3, 4, 5, 6, 7, 8],
+ [0, 1, 0, 1, 2, 3, 4, 0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 6, 7, 8]],
+ device='cuda:0', dtype=torch.int32)
+ """
+
+
+if __name__ == "__main__":
+ test()
diff --git a/lightllm/models/neo_chat_moe/vision_process.py b/lightllm/models/neo_chat_moe/vision_process.py
new file mode 100644
index 0000000000..fbd57a5e9c
--- /dev/null
+++ b/lightllm/models/neo_chat_moe/vision_process.py
@@ -0,0 +1,141 @@
+import re
+import math
+import torch
+import string
+import numpy as np
+import pandas as pd
+from PIL import Image
+import torch.distributed as dist
+import torchvision.transforms as T
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def round_by_factor(number: int, factor: int) -> int:
+ """Returns the closest integer to 'number' that is divisible by 'factor'."""
+ return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+ """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+ return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+ """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+ return math.floor(number / factor) * factor
+
+
+# copy from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60
+def smart_resize(
+ height: int, width: int, factor: int = 32, min_pixels: int = 65536, max_pixels: int = 4194304
+) -> tuple[int, int]:
+ """
+ Rescales the image so that the following conditions are met:
+
+ 1. Both dimensions (height and width) are divisible by 'factor'.
+
+ 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+ 3. The aspect ratio of the image is maintained as closely as possible.
+ """
+ if max(height, width) / min(height, width) > 200:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than {200}, got {max(height, width) / min(height, width)}"
+ )
+ h_bar = max(factor, round_by_factor(height, factor))
+ w_bar = max(factor, round_by_factor(width, factor))
+ if h_bar * w_bar > max_pixels:
+ beta = math.sqrt((height * width) / max_pixels)
+ h_bar = max(factor, floor_by_factor(height / beta, factor))
+ w_bar = max(factor, floor_by_factor(width / beta, factor))
+ elif h_bar * w_bar < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ h_bar = ceil_by_factor(height * beta, factor)
+ w_bar = ceil_by_factor(width * beta, factor)
+ return h_bar, w_bar
+
+
+def dynamic_preprocess_native_resolution(image, size_factor=32, min_pixels=65536, max_pixels=4194304, **kwargs):
+ width, height = image.size
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=size_factor,
+ min_pixels=min_pixels,
+ max_pixels=max_pixels,
+ )
+ image = image.resize((resized_width, resized_height))
+
+ return image
+
+
+def preprocess_pixel_values(pixel_values, patch_size=16):
+ c, h, w = pixel_values.shape
+ grid_h = h // patch_size
+ grid_w = w // patch_size
+
+ flatten_pixel_values = (
+ pixel_values.view(c, grid_h, patch_size, grid_w, patch_size)
+ .permute(1, 3, 0, 2, 4) # [grid_h, grid_w, c, patch_size, patch_size]
+ .reshape(grid_h * grid_w, c * patch_size ** 2)
+ )
+
+ grid_hw = torch.tensor([[grid_h, grid_w]]).to(device=pixel_values.device)
+
+ return flatten_pixel_values, grid_hw
+
+
+def get_contrasting_background(image):
+ """
+ Calculate the color (white or black) that is different from the average foreground color
+ to use as the background color
+ """
+ image_np = np.array(image)
+ if (image_np[:, :, 3] == 0).any():
+ non_transparent_pixels = image_np[:, :, :3][image_np[:, :, 3] > 0]
+ if non_transparent_pixels.size == 0:
+ return None
+ pixel_mean = non_transparent_pixels.mean()
+ contrasting_color = (0, 0, 0) if pixel_mean > 382.5 else (255, 255, 255)
+ return contrasting_color
+ else:
+ return None
+
+
+def load_image_native(image, patch_size=16, downsample_ratio=0.5, min_pixels=65536, max_pixels=4194304, upscale=False):
+ """
+ Load and preprocess an image file, converting it to RGB mode,
+ resizing, normalizing, and optionally adding a thumbnail version.
+ """
+ if image.mode == "RGBA":
+ bg_color = get_contrasting_background(image)
+ if bg_color:
+ background = Image.new("RGB", image.size, bg_color)
+ background.paste(image, mask=image.split()[3])
+ image = background.convert("RGB")
+ else:
+ image = image.convert("RGB")
+ else:
+ image = image.convert("RGB")
+
+ if upscale:
+ image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
+
+ transform = T.Compose(
+ [
+ T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
+ T.ToTensor(),
+ T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+ ]
+ )
+
+ new_image = dynamic_preprocess_native_resolution(
+ image, size_factor=int(patch_size // downsample_ratio), min_pixels=min_pixels, max_pixels=max_pixels
+ )
+ pixel_values, grid_hw = preprocess_pixel_values(transform(new_image).to(torch.float32), patch_size=patch_size)
+
+ # print(f"Transfer image_size from ({image.height, image.width}) to ({new_image.height, new_image.width})")
+
+ return pixel_values, grid_hw
diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py
index 1878e60f1c..fbcf23cc74 100644
--- a/lightllm/server/api_openai.py
+++ b/lightllm/server/api_openai.py
@@ -396,7 +396,9 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
prompt_tokens = prompt_tokens_dict[sub_ids[0]]
completion_tokens = sum(count_output_tokens_dict[sub_req_id] for sub_req_id in sub_ids)
cached_tokens = prompt_cache_len_dict.get(sub_ids[0], 0)
- reasoning_tokens = sum(reasoning_parser_dict[sub_req_id].reasoning_tokens for sub_req_id in sub_ids)
+ reasoning_tokens = sum(
+ getattr(reasoning_parser_dict.get(sub_req_id), "reasoning_tokens", 0) for sub_req_id in sub_ids
+ )
for i in range(request.n):
sub_req_id = sub_ids[i]
diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py
index 8e31c50624..b790058750 100644
--- a/lightllm/server/core/objs/sampling_params.py
+++ b/lightllm/server/core/objs/sampling_params.py
@@ -278,6 +278,8 @@ class SamplingParams(ctypes.Structure):
("ignore_eos", ctypes.c_bool),
# the max number of image patches to be used in the internvl model, for the test
("image_max_patch_num", ctypes.c_int),
+ ("min_pixels", ctypes.c_int),
+ ("max_pixels", ctypes.c_int),
("max_new_tokens", ctypes.c_int),
("min_new_tokens", ctypes.c_int),
# Whether to count input tokens for presence_penalty, frequency_penalty and repetition_penalty
@@ -333,6 +335,8 @@ def init(self, tokenizer, **kwargs):
self.top_k = kwargs.get("top_k", SamplingParams._top_k)
self.ignore_eos = kwargs.get("ignore_eos", False)
self.image_max_patch_num = kwargs.get("image_max_patch_num", -1)
+ self.min_pixels = kwargs.get("min_pixels", -1)
+ self.max_pixels = kwargs.get("max_pixels", -1)
self.max_new_tokens = kwargs.get("max_new_tokens", 65535)
self.min_new_tokens = kwargs.get("min_new_tokens", 1)
self.input_penalty = kwargs.get("input_penalty", DEFAULT_INPUT_PENALTY)
@@ -491,6 +495,8 @@ def to_dict(self):
"top_k": self.top_k,
"ignore_eos": self.ignore_eos,
"image_max_patch_num": self.image_max_patch_num,
+ "min_pixels": self.min_pixels,
+ "max_pixels": self.max_pixels,
"max_new_tokens": self.max_new_tokens,
"min_new_tokens": self.min_new_tokens,
"exponential_decay_length_penalty": self.exponential_decay_length_penalty.to_tuple(),
diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py
index e1a4e421d1..b2a58e0cfb 100644
--- a/lightllm/server/tokenizer.py
+++ b/lightllm/server/tokenizer.py
@@ -33,6 +33,7 @@
from ..models.gemma3.model import Gemma3Tokenizer
from ..models.gemma4.tokenizer import Gemma4Tokenizer
from ..models.qwen3_omni_moe_thinker.model import QWen3OmniTokenizer
+from ..models.neo_chat_moe.model import NeoChatTokenizer
from ..models import deepseek3_2 # noqa: F401 # registers the deepseek_v32 config with transformers
# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file.
@@ -140,5 +141,7 @@ def get_tokenizer(
processor = AutoProcessor.from_pretrained(tokenizer_name)
image_processor = processor.image_processor
tokenizer = Gemma4Tokenizer(tokenizer, model_cfg, image_processor=image_processor)
+ elif model_type == "neo_chat":
+ tokenizer = NeoChatTokenizer(tokenizer, model_cfg, weight_dir=tokenizer_name)
return tokenizer
diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py
index 68e0a97ca1..b3dad0ef71 100644
--- a/lightllm/server/visualserver/model_infer/model_rpc.py
+++ b/lightllm/server/visualserver/model_infer/model_rpc.py
@@ -21,6 +21,7 @@
from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel
from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel
from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel
+from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel
from lightllm.utils.infer_utils import set_random_seed
from lightllm.utils.dist_utils import init_vision_distributed_env
from lightllm.utils.envs_utils import get_env_start_args
@@ -110,6 +111,8 @@ def exposed_init_model(self, kvargs):
.eval()
.bfloat16()
)
+ elif self.model_type == "neo_chat":
+ self.model = NeoVisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16()
else:
raise Exception(f"can not support {self.model_type} now")
diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py
index 6695f0ec44..f659af2253 100644
--- a/lightllm/utils/config_utils.py
+++ b/lightllm/utils/config_utils.py
@@ -417,6 +417,8 @@ def has_vision_module(model_path: str) -> bool:
):
# Qwen3OmniMoeVisionTransformerPretrainedModel
return True
+ elif model_type == "neo_chat":
+ return True
elif model_type in ["qwen3_5", "qwen3_5_moe"]:
return True
else:
diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py
new file mode 100644
index 0000000000..5cb8bd58ff
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py
@@ -0,0 +1,342 @@
+"""Unit test for ``context_attention_fwd_neo``.
+
+Torch reference expresses the *semantics* of the attention, not the kernel's
+internal block structure — it has no notion of BLOCK_N / BLOCK_M. For each
+batch element we gather K/V for the whole request (prompt + new tokens) via
+``req_to_token_indexs`` and apply::
+
+ allow[m, k] = (k <= q_pos[m]) OR (k < image_end[m])
+
+i.e. normal queries are causal, and image-token queries can only see future
+tokens from the same image span. If the Triton kernel disagrees with this
+reference, the kernel is wrong.
+
+Run directly for quick debugging:
+
+ python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\
+ test_context_attention_fwd_neo.py
+
+or via pytest:
+
+ pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\
+ test_context_attention_fwd_neo.py -x -s
+"""
+
+import math
+import pytest
+import torch
+
+from lightllm.models.neo_chat_moe.triton_kernel.context_attention_fwd_neo import (
+ context_attention_fwd_neo,
+)
+
+
+def torch_reference_context_attention_neo(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ b_req_idx: torch.Tensor,
+ b_start_loc: torch.Tensor,
+ b_seq_len: torch.Tensor,
+ b_prompt_cache_len: torch.Tensor,
+ req_to_token_indexs: torch.Tensor,
+ b_image_token_end: torch.Tensor,
+) -> torch.Tensor:
+ device = q.device
+ dtype = q.dtype
+ _, Hq, D = q.shape
+ Hk = k.shape[1]
+ kv_group = Hq // Hk
+ scale = 1.0 / math.sqrt(D)
+
+ out = torch.empty_like(q)
+
+ for b in range(b_seq_len.shape[0]):
+ req_idx = int(b_req_idx[b].item())
+ total = int(b_seq_len[b].item())
+ prompt = int(b_prompt_cache_len[b].item())
+ new = total - prompt
+ if new <= 0:
+ continue
+
+ q_start = int(b_start_loc[b].item())
+ q_blk = q[q_start : q_start + new] # [M, Hq, D]
+ q_image_end = b_image_token_end[q_start : q_start + new].to(torch.int64)
+
+ token_locs = req_to_token_indexs[req_idx, :total].to(torch.int64)
+ k_blk = k[token_locs] # [total, Hk, D]
+ v_blk = v[token_locs]
+
+ q_pos = torch.arange(prompt, total, device=device, dtype=torch.int64) # [M]
+ k_pos = torch.arange(0, total, device=device, dtype=torch.int64) # [total]
+ causal = k_pos[None, :] <= q_pos[:, None]
+ allow = causal | (k_pos[None, :] < q_image_end[:, None])
+
+ out_blk = torch.empty_like(q_blk)
+ for h in range(Hq):
+ h_k = h // kv_group
+ q_h = q_blk[:, h, :].to(torch.float32)
+ k_h = k_blk[:, h_k, :].to(torch.float32)
+ v_h = v_blk[:, h_k, :].to(torch.float32)
+
+ scores = (q_h @ k_h.transpose(0, 1)) * scale
+ scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8))
+ probs = torch.softmax(scores, dim=-1)
+ out_h = (probs @ v_h).to(dtype)
+ out_blk[:, h, :] = out_h
+
+ out[q_start : q_start + new] = out_blk
+
+ return out
+
+
+def _build_inputs(
+ batch: int,
+ Hq: int,
+ Hk: int,
+ D: int,
+ dtype: torch.dtype,
+ device: str,
+ max_new: int = 256,
+ max_prompt: int = 512,
+ image_prob: float = 0.7,
+ num_image_spans_max: int = 3,
+ image_span_len_max: int = 24,
+ kv_pool_slack: int = 4096,
+ seed: int = 0,
+):
+ g = torch.Generator(device="cpu").manual_seed(seed)
+
+ new_lens = torch.randint(low=1, high=max_new + 1, size=(batch,), generator=g)
+ prompt_lens = torch.randint(low=0, high=max_prompt + 1, size=(batch,), generator=g)
+ total_lens = new_lens + prompt_lens
+
+ sum_new = int(new_lens.sum().item())
+ sum_total = int(total_lens.sum().item())
+ max_total_len = int(total_lens.max().item())
+ max_new_len = int(new_lens.max().item())
+
+ b_start_loc = torch.zeros(batch, dtype=torch.int32)
+ cur = 0
+ for i in range(batch):
+ b_start_loc[i] = cur
+ cur += int(new_lens[i].item())
+
+ # Permute so batch idx != request idx: exercises the Req_to_tokens indexing.
+ b_req_idx = torch.randperm(batch, generator=g).to(torch.int32)
+
+ # Global KV pool with scattered, non-contiguous slot assignment per request.
+ base = 1024
+ kv_pool_size = base + sum_total + kv_pool_slack
+ pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base
+
+ req_to_token_indexs = torch.zeros((batch, max_total_len), dtype=torch.int32)
+ p = 0
+ for r_logical, req_id in enumerate(b_req_idx.tolist()):
+ L = int(total_lens[r_logical].item())
+ req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32)
+ p += L
+
+ b_image_token_end = torch.zeros(sum_new, dtype=torch.int32)
+ for i in range(batch):
+ M = int(new_lens[i].item())
+ P = int(prompt_lens[i].item())
+ start_pack = int(b_start_loc[i].item())
+ if M < 2:
+ continue
+ if torch.rand((), generator=g).item() > image_prob:
+ continue
+ n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item())
+ cursor = 0
+ for _ in range(n_spans):
+ remaining = M - cursor
+ if remaining <= 0:
+ break
+ gap = int(torch.randint(0, remaining, (1,), generator=g).item())
+ s_rel = cursor + gap
+ max_span_len = min(image_span_len_max, M - s_rel)
+ if max_span_len <= 0:
+ break
+ span_len = int(torch.randint(1, max_span_len + 1, (1,), generator=g).item())
+ e_rel = s_rel + span_len
+ image_end = P + e_rel
+ b_image_token_end[start_pack + s_rel : start_pack + e_rel] = image_end
+ cursor = e_rel
+
+ b_seq_len = total_lens.to(torch.int32)
+ b_prompt_cache_len = prompt_lens.to(torch.int32)
+
+ q = torch.randn((sum_new, Hq, D), dtype=dtype, device=device)
+ k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device)
+ v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device)
+ o = torch.empty_like(q)
+
+ return dict(
+ q=q,
+ k=k,
+ v=v,
+ o=o,
+ b_req_idx=b_req_idx.to(device),
+ b_start_loc=b_start_loc.to(device),
+ b_seq_len=b_seq_len.to(device),
+ b_prompt_cache_len=b_prompt_cache_len.to(device),
+ max_new_len=max_new_len,
+ req_to_token_indexs=req_to_token_indexs.to(device),
+ b_image_token_end=b_image_token_end.to(device),
+ new_lens=new_lens,
+ prompt_lens=prompt_lens,
+ )
+
+
+def _report_per_batch_error(out_triton, out_ref, new_lens, b_start_loc, image_token_end, tag=""):
+ print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):")
+ for i in range(new_lens.shape[0]):
+ s = int(b_start_loc[i].item())
+ m = int(new_lens[i].item())
+ if m == 0:
+ continue
+ a = out_triton[s : s + m].float()
+ b = out_ref[s : s + m].float()
+ abs_err = (a - b).abs().max().item()
+ denom = b.abs().max().item() + 1e-6
+ rel_err = abs_err / denom
+ cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+ n_img = int((image_token_end[s : s + m] > 0).sum().item())
+ print(
+ f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | "
+ f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}"
+ )
+
+
+def _run_case(
+ batch: int,
+ Hq: int,
+ Hk: int,
+ D: int,
+ dtype: torch.dtype,
+ seed: int,
+ max_new: int,
+ max_prompt: int,
+ atol: float = 5e-2,
+ rtol: float = 5e-2,
+ cos_threshold: float = 0.99,
+ verbose: bool = True,
+):
+ assert Hq % Hk == 0
+ device = "cuda"
+
+ inputs = _build_inputs(
+ batch=batch,
+ Hq=Hq,
+ Hk=Hk,
+ D=D,
+ dtype=dtype,
+ device=device,
+ max_new=max_new,
+ max_prompt=max_prompt,
+ seed=seed,
+ )
+
+ context_attention_fwd_neo(
+ inputs["q"],
+ inputs["k"],
+ inputs["v"],
+ inputs["o"],
+ inputs["b_req_idx"],
+ inputs["b_start_loc"],
+ inputs["b_seq_len"],
+ inputs["b_prompt_cache_len"],
+ inputs["max_new_len"],
+ inputs["req_to_token_indexs"],
+ inputs["b_image_token_end"],
+ )
+ out_triton = inputs["o"]
+
+ out_ref = torch_reference_context_attention_neo(
+ inputs["q"],
+ inputs["k"],
+ inputs["v"],
+ inputs["b_req_idx"],
+ inputs["b_start_loc"],
+ inputs["b_seq_len"],
+ inputs["b_prompt_cache_len"],
+ inputs["req_to_token_indexs"],
+ inputs["b_image_token_end"],
+ )
+
+ a = out_triton.float()
+ b = out_ref.float()
+ abs_err = (a - b).abs().max().item()
+ denom = b.abs().max().item() + 1e-6
+ rel_err = abs_err / denom
+ cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+
+ n_image = int((inputs["b_image_token_end"] > 0).sum().item())
+ n_tokens = int(inputs["b_image_token_end"].numel())
+ if verbose:
+ print(
+ f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} "
+ f"seed={seed} image_tokens={n_image}/{n_tokens}"
+ )
+ print(
+ f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} "
+ f"(allclose atol={atol}, rtol={rtol}? "
+ f"{torch.allclose(a, b, atol=atol, rtol=rtol)})"
+ )
+ _report_per_batch_error(
+ out_triton,
+ out_ref,
+ inputs["new_lens"],
+ inputs["b_start_loc"],
+ inputs["b_image_token_end"],
+ tag=f"seed={seed}",
+ )
+
+ return abs_err, rel_err, cos
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA")
+@pytest.mark.parametrize(
+ "batch,Hq,Hk,D,dtype,seed,max_new,max_prompt",
+ [
+ (4, 8, 2, 128, torch.bfloat16, 0, 128, 256),
+ (4, 8, 2, 128, torch.bfloat16, 1, 256, 512),
+ (8, 16, 4, 128, torch.bfloat16, 2, 256, 512),
+ (16, 28, 4, 128, torch.bfloat16, 3, 128, 256),
+ (4, 8, 2, 128, torch.float16, 4, 256, 512),
+ (4, 8, 8, 64, torch.bfloat16, 5, 128, 256),
+ (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024),
+ ],
+)
+def test_context_attention_fwd_neo(batch, Hq, Hk, D, dtype, seed, max_new, max_prompt):
+ abs_err, rel_err, cos = _run_case(
+ batch=batch,
+ Hq=Hq,
+ Hk=Hk,
+ D=D,
+ dtype=dtype,
+ seed=seed,
+ max_new=max_new,
+ max_prompt=max_prompt,
+ verbose=True,
+ )
+ assert cos > 0.99, f"cosine similarity too low: {cos}"
+ assert rel_err < 5e-2, f"max relative error too large: {rel_err}"
+
+
+if __name__ == "__main__":
+ if not torch.cuda.is_available():
+ print("No CUDA available.")
+ raise SystemExit(0)
+
+ torch.manual_seed(0)
+
+ cases = [
+ dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_new=128, max_prompt=256),
+ dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_new=256, max_prompt=512),
+ dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_new=128, max_prompt=256),
+ dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_new=256, max_prompt=512),
+ ]
+
+ for cfg in cases:
+ _run_case(**cfg, verbose=True)
diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py
new file mode 100644
index 0000000000..41f6b476ba
--- /dev/null
+++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py
@@ -0,0 +1,523 @@
+"""Unit test for the FA3-based prefill path with image-token support.
+
+This test pre-wires a call to ``flash_attn_with_kvcache`` with an
+``image_token_tag`` keyword argument. The expectation is that ``fa3-neo``'s
+``flash_attn_with_kvcache`` will be extended with an optional
+``image_token_tag`` parameter that, for queries flagged as image tokens,
+relaxes the causal mask so they can attend bidirectionally to every real key
+in the request.
+
+Torch reference expresses the *semantics* of the attention, not FA3's internal
+tiling — it has no notion of BLOCK_N / BLOCK_M. For each batch element we
+gather K/V for the whole request (prompt + new tokens) and apply::
+
+ allow[m, k] = (k <= q_pos[m]) OR image_tag[m] for k in [0, total)
+
+i.e. normal queries are causal, image-token queries can see every real key in
+the request. If FA3 disagrees with this reference, the kernel is wrong.
+
+Run directly for quick debugging:
+
+ python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\
+ test_fa3_neo.py
+
+or via pytest:
+
+ pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\
+ test_fa3_neo.py -x -s
+"""
+
+import math
+import pytest
+import torch
+
+from flash_attn_interface import flash_attn_with_kvcache
+
+try:
+ import triton
+ import triton.testing as triton_testing
+except ImportError:
+ triton = None
+ triton_testing = None
+
+try:
+ from lightllm.models.neo_chat_moe.triton_kernel.context_attention_fwd_neo import (
+ context_attention_fwd_neo,
+ )
+except ImportError:
+ context_attention_fwd_neo = None
+
+
+def torch_reference_context_attention_neo(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ b_req_idx: torch.Tensor,
+ b_q_start_loc: torch.Tensor,
+ b_seq_len: torch.Tensor,
+ b_prompt_cache_len: torch.Tensor,
+ req_to_token_indexs: torch.Tensor,
+ b_image_token_tag: torch.Tensor,
+) -> torch.Tensor:
+ device = q.device
+ dtype = q.dtype
+ _, Hq, D = q.shape
+ Hk = k.shape[1]
+ kv_group = Hq // Hk
+ scale = 1.0 / math.sqrt(D)
+
+ out = torch.empty_like(q)
+
+ for b in range(b_seq_len.shape[0]):
+ req_idx = int(b_req_idx[b].item())
+ seq_len = int(b_seq_len[b].item())
+ prompt_cache_len = int(b_prompt_cache_len[b].item())
+ q_seq_len = seq_len - prompt_cache_len
+ if q_seq_len <= 0:
+ continue
+
+ q_start = int(b_q_start_loc[b].item())
+ q_blk = q[q_start : q_start + q_seq_len] # [M, Hq, D]
+ image_tag = b_image_token_tag[q_start : q_start + q_seq_len].to(torch.bool)
+
+ token_locs = req_to_token_indexs[req_idx, :seq_len].to(torch.int64)
+ k_blk = k[token_locs] # [seq_len, Hk, D]
+ v_blk = v[token_locs]
+
+ q_pos = torch.arange(prompt_cache_len, seq_len, device=device, dtype=torch.int64) # [M]
+ k_pos = torch.arange(0, seq_len, device=device, dtype=torch.int64) # [seq_len]
+ causal = k_pos[None, :] <= q_pos[:, None]
+ allow = causal | image_tag[:, None]
+
+ out_blk = torch.empty_like(q_blk)
+ for h in range(Hq):
+ h_k = h // kv_group
+ q_h = q_blk[:, h, :].to(torch.float32)
+ k_h = k_blk[:, h_k, :].to(torch.float32)
+ v_h = v_blk[:, h_k, :].to(torch.float32)
+
+ scores = (q_h @ k_h.transpose(0, 1)) * scale
+ scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8))
+ probs = torch.softmax(scores, dim=-1)
+ out_h = (probs @ v_h).to(dtype)
+ out_blk[:, h, :] = out_h
+
+ out[q_start : q_start + q_seq_len] = out_blk
+
+ return out
+
+
+def _build_inputs(
+ batch: int,
+ Hq: int,
+ Hk: int,
+ D: int,
+ dtype: torch.dtype,
+ device: str,
+ max_q_seq_len: int = 256,
+ max_prompt_cache_len: int = 512,
+ image_prob: float = 0.7,
+ num_image_spans_max: int = 3,
+ image_span_len_max: int = 24,
+ kv_pool_slack: int = 4096,
+ seed: int = 0,
+):
+ """Build one realistic prefill batch.
+
+ Naming matches lightllm's infer_state:
+ - ``q_seq_len`` = number of new Q tokens in this prefill call
+ - ``prompt_cache_len`` = length of the already-cached prefix for this req
+ - ``seq_len`` = prompt_cache_len + q_seq_len (total KV length)
+ """
+ g = torch.Generator(device="cpu").manual_seed(seed)
+
+ q_seq_lens = torch.randint(low=1, high=max_q_seq_len + 1, size=(batch,), generator=g)
+ prompt_cache_lens = torch.randint(low=0, high=max_prompt_cache_len + 1, size=(batch,), generator=g)
+ seq_lens = q_seq_lens + prompt_cache_lens
+
+ sum_q = int(q_seq_lens.sum().item())
+ sum_total = int(seq_lens.sum().item())
+ max_seq_len_in_batch = int(seq_lens.max().item())
+ max_q_seq_len_in_batch = int(q_seq_lens.max().item())
+
+ b_q_start_loc = torch.zeros(batch, dtype=torch.int32)
+ cur = 0
+ for i in range(batch):
+ b_q_start_loc[i] = cur
+ cur += int(q_seq_lens[i].item())
+
+ # Permute so batch idx != request idx: exercises the page_table indexing.
+ b_req_idx = torch.randperm(batch, generator=g).to(torch.int32)
+
+ # Global KV pool with scattered, non-contiguous slot assignment per request.
+ base = 1024
+ kv_pool_size = base + sum_total + kv_pool_slack
+ pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base
+
+ req_to_token_indexs = torch.zeros((batch, max_seq_len_in_batch), dtype=torch.int32)
+ p = 0
+ for r_logical, req_id in enumerate(b_req_idx.tolist()):
+ L = int(seq_lens[r_logical].item())
+ req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32)
+ p += L
+
+ # Randomly place contiguous image-token spans inside each batch's new-Q region.
+ b_image_token_tag = torch.zeros(sum_q, dtype=torch.bool)
+ for i in range(batch):
+ M = int(q_seq_lens[i].item())
+ if M < 2:
+ continue
+ if torch.rand((), generator=g).item() > image_prob:
+ continue
+ n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item())
+ start_pack = int(b_q_start_loc[i].item())
+ for _ in range(n_spans):
+ span_len = int(torch.randint(1, max(2, image_span_len_max) + 1, (1,), generator=g).item())
+ span_len = min(span_len, M)
+ s_rel = int(torch.randint(0, M - span_len + 1, (1,), generator=g).item())
+ b_image_token_tag[start_pack + s_rel : start_pack + s_rel + span_len] = True
+
+ b_seq_len = seq_lens.to(torch.int32)
+ b_prompt_cache_len = prompt_cache_lens.to(torch.int32)
+
+ q = torch.randn((sum_q, Hq, D), dtype=dtype, device=device)
+ k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device)
+ v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device)
+
+ return dict(
+ q=q,
+ k=k,
+ v=v,
+ b_req_idx=b_req_idx.to(device),
+ b_q_start_loc=b_q_start_loc.to(device),
+ b_seq_len=b_seq_len.to(device),
+ b_prompt_cache_len=b_prompt_cache_len.to(device),
+ max_seq_len_in_batch=max_seq_len_in_batch,
+ max_q_seq_len_in_batch=max_q_seq_len_in_batch,
+ req_to_token_indexs=req_to_token_indexs.to(device),
+ b_image_token_tag=b_image_token_tag.to(device),
+ q_seq_lens=q_seq_lens,
+ prompt_cache_lens=prompt_cache_lens,
+ )
+
+
+def _fa3_prefill_with_image_tag(inputs: dict) -> torch.Tensor:
+ """Drive ``flash_attn_with_kvcache`` with the same prefill semantics as
+ ``Fa3PrefillAttState._nomarl_prefill_att`` plus an optional
+ ``image_token_tag`` kwarg for image-token bidirectional attention.
+ """
+ q = inputs["q"]
+ k = inputs["k"]
+ v = inputs["v"]
+ device = q.device
+
+ # Build page_table[b, p] = req_to_token_indexs[b_req_idx[b], p].
+ page_table = inputs["req_to_token_indexs"][inputs["b_req_idx"].long(), : inputs["max_seq_len_in_batch"]].to(
+ torch.int32
+ )
+
+ q_seq_lens_t = inputs["b_seq_len"].to(torch.int32) - inputs["b_prompt_cache_len"].to(torch.int32)
+
+ cu_seqlens_q = torch.zeros(q_seq_lens_t.shape[0] + 1, dtype=torch.int32, device=device)
+ cu_seqlens_q[1:] = q_seq_lens_t.cumsum(0).to(torch.int32)
+
+ cu_seqlens_k = torch.zeros(inputs["b_seq_len"].shape[0] + 1, dtype=torch.int32, device=device)
+ cu_seqlens_k[1:] = inputs["b_seq_len"].cumsum(0).to(torch.int32)
+
+ sm_scale = 1.0 / math.sqrt(q.shape[-1])
+
+ # page_size = 1 paged KV cache view.
+ k_cache = k.view(k.shape[0], 1, k.shape[1], k.shape[2])
+ v_cache = v.view(v.shape[0], 1, v.shape[1], v.shape[2])
+
+ o = flash_attn_with_kvcache(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ page_table=page_table,
+ cache_seqlens=inputs["b_seq_len"],
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_k_new=cu_seqlens_k,
+ max_seqlen_q=inputs["max_q_seq_len_in_batch"],
+ softmax_scale=sm_scale,
+ causal=True,
+ window_size=(-1, -1),
+ softcap=0.0,
+ k_descale=None,
+ v_descale=None,
+ return_softmax_lse=False,
+ # image-token bidirectional attention. Packed like q (shape [sum_q],
+ # bool). Rows where the tag is True are allowed to attend to every
+ # real key in the request (not just the causal prefix).
+ # The kernel uses warp OR reduce to detect image tokens per M-block
+ # and extends n_block_max for full attention automatically.
+ image_token_tag=inputs["b_image_token_tag"],
+ )
+ return o
+
+
+def _report_per_batch_error(out_fa3, out_ref, q_seq_lens, b_q_start_loc, image_tag, tag=""):
+ print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):")
+ for i in range(q_seq_lens.shape[0]):
+ s = int(b_q_start_loc[i].item())
+ m = int(q_seq_lens[i].item())
+ if m == 0:
+ continue
+ a = out_fa3[s : s + m].float()
+ b = out_ref[s : s + m].float()
+ abs_err = (a - b).abs().max().item()
+ denom = b.abs().max().item() + 1e-6
+ rel_err = abs_err / denom
+ cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+ n_img = int(image_tag[s : s + m].sum().item())
+ print(
+ f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | "
+ f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}"
+ )
+
+
+def _run_case(
+ batch: int,
+ Hq: int,
+ Hk: int,
+ D: int,
+ dtype: torch.dtype,
+ seed: int,
+ max_q_seq_len: int,
+ max_prompt_cache_len: int,
+ atol: float = 5e-2,
+ rtol: float = 5e-2,
+ cos_threshold: float = 0.99,
+ verbose: bool = True,
+):
+ assert Hq % Hk == 0
+ device = "cuda"
+
+ inputs = _build_inputs(
+ batch=batch,
+ Hq=Hq,
+ Hk=Hk,
+ D=D,
+ dtype=dtype,
+ device=device,
+ max_q_seq_len=max_q_seq_len,
+ max_prompt_cache_len=max_prompt_cache_len,
+ seed=seed,
+ )
+
+ out_fa3 = _fa3_prefill_with_image_tag(inputs)
+
+ out_ref = torch_reference_context_attention_neo(
+ inputs["q"],
+ inputs["k"],
+ inputs["v"],
+ inputs["b_req_idx"],
+ inputs["b_q_start_loc"],
+ inputs["b_seq_len"],
+ inputs["b_prompt_cache_len"],
+ inputs["req_to_token_indexs"],
+ inputs["b_image_token_tag"],
+ )
+
+ a = out_fa3.float().reshape_as(out_ref.float())
+ b = out_ref.float()
+ abs_err = (a - b).abs().max().item()
+ denom = b.abs().max().item() + 1e-6
+ rel_err = abs_err / denom
+ cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+
+ n_image = int(inputs["b_image_token_tag"].sum().item())
+ n_tokens = int(inputs["b_image_token_tag"].numel())
+ if verbose:
+ print(
+ f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} "
+ f"seed={seed} image_tokens={n_image}/{n_tokens}"
+ )
+ print(
+ f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} "
+ f"(allclose atol={atol}, rtol={rtol}? "
+ f"{torch.allclose(a, b, atol=atol, rtol=rtol)})"
+ )
+ _report_per_batch_error(
+ out_fa3,
+ out_ref,
+ inputs["q_seq_lens"],
+ inputs["b_q_start_loc"],
+ inputs["b_image_token_tag"],
+ tag=f"seed={seed}",
+ )
+
+ return abs_err, rel_err, cos
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA")
+@pytest.mark.skipif(flash_attn_with_kvcache is None, reason="fa3 not available")
+@pytest.mark.parametrize(
+ "batch,Hq,Hk,D,dtype,seed,max_q_seq_len,max_prompt_cache_len",
+ [
+ (4, 8, 2, 128, torch.bfloat16, 0, 128, 256),
+ (4, 8, 2, 128, torch.bfloat16, 1, 256, 512),
+ (8, 16, 4, 128, torch.bfloat16, 2, 256, 512),
+ (16, 28, 4, 128, torch.bfloat16, 3, 128, 256),
+ (4, 8, 2, 128, torch.float16, 4, 256, 512),
+ (4, 8, 8, 64, torch.bfloat16, 5, 128, 256),
+ (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024),
+ ],
+)
+def test_fa3_neo_prefill_with_image_tag(batch, Hq, Hk, D, dtype, seed, max_q_seq_len, max_prompt_cache_len):
+ abs_err, rel_err, cos = _run_case(
+ batch=batch,
+ Hq=Hq,
+ Hk=Hk,
+ D=D,
+ dtype=dtype,
+ seed=seed,
+ max_q_seq_len=max_q_seq_len,
+ max_prompt_cache_len=max_prompt_cache_len,
+ verbose=True,
+ )
+ assert cos > 0.99, f"cosine similarity too low: {cos}"
+ assert rel_err < 5e-2, f"max relative error too large: {rel_err}"
+
+
+def _bench_case(
+ batch: int,
+ Hq: int,
+ Hk: int,
+ D: int,
+ dtype: torch.dtype,
+ seed: int,
+ max_q_seq_len: int,
+ max_prompt_cache_len: int,
+ rep_ms: int = 100,
+ warmup_iters: int = 3,
+):
+ """Compare FA3 (with image_token_tag) vs the original Triton
+ ``context_attention_fwd_neo`` using ``triton.testing.do_bench_cudagraph``.
+
+ Both kernels are captured into a CUDA graph so scheduling/launch overhead
+ is minimized and the measurement reflects the kernel cost.
+ """
+ assert Hq % Hk == 0
+ device = "cuda"
+
+ inputs = _build_inputs(
+ batch=batch,
+ Hq=Hq,
+ Hk=Hk,
+ D=D,
+ dtype=dtype,
+ device=device,
+ max_q_seq_len=max_q_seq_len,
+ max_prompt_cache_len=max_prompt_cache_len,
+ seed=seed,
+ )
+
+ # --- fa3 runner: output tensor is allocated inside flash_attn_with_kvcache.
+ def fa3_run():
+ return _fa3_prefill_with_image_tag(inputs)
+
+ # --- triton runner: pre-allocate o & position_ids so the graph captures
+ # only the kernel launch.
+ o_triton = torch.empty_like(inputs["q"])
+ # Kernel signature requires position_ids but the current masking path does
+ # not read it; zeros are fine for perf measurement.
+ position_ids_0 = torch.zeros(inputs["q"].shape[0], dtype=torch.int32, device=inputs["q"].device)
+
+ def triton_run():
+ context_attention_fwd_neo(
+ inputs["q"],
+ inputs["k"],
+ inputs["v"],
+ o_triton,
+ position_ids_0,
+ inputs["b_req_idx"],
+ inputs["b_q_start_loc"],
+ inputs["b_seq_len"],
+ inputs["b_prompt_cache_len"],
+ inputs["max_q_seq_len_in_batch"],
+ inputs["req_to_token_indexs"],
+ inputs["b_image_token_tag"],
+ )
+
+ # Warm up outside the graph capture so lazy allocations / autotune happen.
+ for _ in range(warmup_iters):
+ fa3_run()
+ triton_run()
+ torch.cuda.synchronize()
+
+ fa3_ms = triton_testing.do_bench_cudagraph(fa3_run, rep=rep_ms)
+ triton_ms = triton_testing.do_bench_cudagraph(triton_run, rep=rep_ms)
+
+ n_image = int(inputs["b_image_token_tag"].sum().item())
+ n_tokens = int(inputs["b_image_token_tag"].numel())
+ sum_kv = int(inputs["b_seq_len"].sum().item())
+ speedup = triton_ms / fa3_ms if fa3_ms > 0 else float("inf")
+
+ print(
+ f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} "
+ f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} "
+ f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | "
+ f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | "
+ f"speedup {speedup:5.2f}x"
+ )
+
+ return fa3_ms, triton_ms
+
+
+if __name__ == "__main__":
+ if not torch.cuda.is_available():
+ print("No CUDA available.")
+ raise SystemExit(0)
+ if flash_attn_with_kvcache is None:
+ print("fa3 flash_attn_with_kvcache not available (sgl_kernel missing?).")
+ raise SystemExit(0)
+
+ torch.manual_seed(0)
+
+ cases = [
+ dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_q_seq_len=128, max_prompt_cache_len=256),
+ dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_q_seq_len=256, max_prompt_cache_len=512),
+ dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_q_seq_len=128, max_prompt_cache_len=256),
+ # FP16 case disabled: current build compiled without FP16 support
+ # dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_q_seq_len=256, max_prompt_cache_len=512),
+ ]
+
+ print("=" * 100)
+ print("Correctness")
+ print("=" * 100)
+ for cfg in cases:
+ _run_case(**cfg, verbose=True)
+
+ if triton_testing is None or context_attention_fwd_neo is None:
+ print("\nSkipping benchmark: triton or context_attention_fwd_neo not available.")
+ raise SystemExit(0)
+
+ print("\n" + "=" * 100)
+ print("Benchmark (triton.testing.do_bench_cudagraph)")
+ print("=" * 100)
+
+ # Cold-prefill sweep: max_prompt_cache_len=0 so seq_len == q_seq_len.
+ # Head shape matches neo_chat_moe / Qwen3 llm_config:
+ # num_attention_heads = 32, num_key_value_heads = 8, head_dim = 128
+ # (GQA ratio 4:1)
+ bench_batches = [8, 16, 32, 64, 128]
+ bench_q_seq_lens = [1024, 4096, 8192]
+
+ bench_cases = [
+ dict(
+ batch=b,
+ Hq=32,
+ Hk=8,
+ D=128,
+ dtype=torch.bfloat16,
+ seed=0,
+ max_q_seq_len=s,
+ max_prompt_cache_len=0,
+ )
+ for b in bench_batches
+ for s in bench_q_seq_lens
+ ]
+
+ for cfg in bench_cases:
+ _bench_case(**cfg)