dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity
Files
900ae0b1efc19e46cd313eb3ba7eb694fa596f6b
project_6/qwen3_6_scripts/qwen3_5.py
claude 051b02d3cd feat: ix_moe_bridge + ix_attn_bridge — dlopen bridges for full ixformer::infer API 
Bridge architecture (from xllm/core/kernels/ilu/ixformer.h):

ix_moe_bridge.so (MoE 7-step fused pipeline):
  - topk_softmax → moe_compute_token_index_api → moe_expand_input
  - moe_w16a16_group_gemm (x2) → silu_and_mul → moe_output_reduce_sum
  - fused_moe_forward(): replaces entire Python expert loop
  - Fix: group_gemm format NT→TN (match xllm trans_b=true)

ix_attn_bridge.so (attention + linear):
  - ixinfer_flash_attn_unpad_with_block_tables (fused prefill)
  - xllm_paged_attention (fused paged decode)
  - ixformer_linear (matmul + activation)
  - residual_rms_norm (fused residual + norm)

Integration:
  - ix_fused_moe.py: Python loader (prebuilt .so → JIT → unavailable)
  - qwen3_5.py: Tier 0 dispatch in _pure_pytorch_experts()
  - patch_ops.sh: deploys ix_fused_moe.py + all prebuilt/*.so

Source: jd-opensource/xllm (fresh clone, all ILU kernels verified SAME)
Sync: upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h (+32 lines)

Build on real machine:
  bash qwen3_6_scripts/build_ix_moe_bridge.sh
  bash qwen3_6_scripts/build_ix_attn_bridge.sh
2026-08-14 07:32:31 +00:00

2697 lines
115 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# Inference-only Qwen3.6-35B-A3B (Qwen3_5 MoE architecture) for Iluvatar BI-V100.
# Pure-PyTorch DeltaNet (no fla / causal_conv1d dependency).
# Includes the native Qwen3.6 vision tower; MTP remains unsupported.
from functools import lru_cache, partial
import hashlib
import os
import sys
import time
from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional,
Tuple, TypedDict, Union)
def _bi100_model_trace(message: str) -> None:
if os.getenv("BI100_EXECUTOR_STARTUP_DEBUG") == "1":
stamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
rank = os.getenv("RANK", os.getenv("LOCAL_RANK", "?"))
print(f"[BI100 STARTUP] {stamp} pid={os.getpid()} rank={rank} {message}",
file=sys.stderr, flush=True)
_bi100_model_trace("qwen3_5 stdlib imports complete; importing torch and vLLM")
import torch
import torch.nn.functional as F
from torch import nn
from PIL import Image
from transformers.image_utils import (ChannelDimension, get_image_size,
infer_channel_dimension_format,
to_numpy_array)
from transformers.models.qwen2_vl import (
image_processing_qwen2_vl as _qwen2_vl_image_processing)
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
Qwen2VLImageProcessor, smart_resize)
def _compat_make_batched_images(images):
return images if isinstance(images, list) else [images]
def _compat_make_batched_videos(videos):
if isinstance(videos, list) and videos and isinstance(videos[0], list):
return videos
return [videos]
# The CoreX image pins transformers 4.55.3, while its vLLM Qwen2-VL module
# imports helpers introduced by another transformers build.
if not hasattr(_qwen2_vl_image_processing, "make_batched_images"):
_qwen2_vl_image_processing.make_batched_images = \
_compat_make_batched_images
if not hasattr(_qwen2_vl_image_processing, "make_batched_videos"):
_qwen2_vl_image_processing.make_batched_videos = \
_compat_make_batched_videos
from vllm.attention import Attention, AttentionMetadata
from vllm.config import (CacheConfig, LoRAConfig, MultiModalConfig,
SchedulerConfig)
from vllm.distributed import (get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce)
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
MergedColumnParallelLinear,
ReplicatedLinear,
RowParallelLinear)
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import (
MRotaryEmbedding, _apply_rotary_emb)
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead, VocabParallelEmbedding)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader, sharded_weight_loader)
from vllm.model_executor.models.mamba_cache import MambaCacheManager
from vllm.model_executor.models.qwen2_vl import (Qwen2VisionAttention,
Qwen2VisionRotaryEmbedding)
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.model_executor.utils import set_weight_attrs
from vllm.inputs import INPUT_REGISTRY, InputContext, LLMInputs
from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalDataDict,
MultiModalInputs)
from vllm.multimodal.base import MultiModalData
from vllm.sequence import IntermediateTensors, SequenceData
from vllm.transformers_utils.tokenizer import get_tokenizer
from vllm.worker.model_runner import (_BATCH_SIZES_TO_CAPTURE,
_get_graph_batch_size)
from vllm.logger import init_logger
from vllm.bi100_env import env_bool, env_int
from vllm.bi100_profile import (bi100_profile_event_enabled,
bi100_profile_flush,
bi100_profile_transaction, bi100_timer)
try:
from vllm import corex_gdn_causal_conv as _corex_gdn_causal_conv
except ImportError:
_corex_gdn_causal_conv = None
try:
from vllm import corex_gdn_gated_norm as _corex_gdn_gated_norm
except ImportError:
_corex_gdn_gated_norm = None
try:
from vllm import corex_gdn_beta_decay as _corex_gdn_beta_decay
except ImportError:
_corex_gdn_beta_decay = None
try:
from vllm import corex_gdn_qk_map as _corex_gdn_qk_map
except ImportError:
_corex_gdn_qk_map = None
try:
from vllm import corex_gdn_packed_decode as _corex_gdn_packed_decode
except ImportError:
_corex_gdn_packed_decode = None
try:
from vllm import corex_attn_head_rms_norm as _corex_attn_head_rms_norm
except ImportError:
_corex_attn_head_rms_norm = None
try:
from vllm import corex_moe_exact_reduce as _corex_moe_exact_reduce
except ImportError:
_corex_moe_exact_reduce = None
try:
from vllm import corex_moe_weight_gather as _corex_moe_weight_gather
except ImportError:
_corex_moe_weight_gather = None
try:
from vllm import corex_moe_direct_routed as _corex_moe_direct_routed
except ImportError:
_corex_moe_direct_routed = None
try:
from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax
except ImportError:
_corex_moe_topk_softmax = None
try:
from vllm import corex_moe_index_combine as _corex_moe_index_combine
except ImportError:
_corex_moe_index_combine = None
try:
from vllm import corex_gdn_chunk_recurrent as _corex_gdn_chunk_recurrent
except ImportError:
_corex_gdn_chunk_recurrent = None
_HAS_COREX_GDN_CHUNK = _corex_gdn_chunk_recurrent is not None
from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA,
SupportsMultiModal)
logger = init_logger(__name__)
_bi100_model_trace("qwen3_5 runtime imports complete")
_ALLOW_GDN_NAN_ZERO = env_bool("BI100_GDN_ALLOW_NAN_ZERO", False)
_GDN_FINITE_CHECK = (env_bool("BI100_GDN_FINITE_CHECK", False)
or _ALLOW_GDN_NAN_ZERO)
_DNN_CHUNK_SIZE = env_int("BI100_DNN_CHUNK", 4096, 64, 65536)
_USE_COREX_GDN_CAUSAL_CONV = (
_corex_gdn_causal_conv is not None
and env_bool("BI100_GDN_COREX_CAUSAL_CONV", True))
_USE_COREX_GDN_GATED_NORM = (
_corex_gdn_gated_norm is not None
and env_bool("BI100_GDN_COREX_GATED_NORM", True))
_USE_COREX_GDN_BETA_DECAY = (
_corex_gdn_beta_decay is not None
and env_bool("BI100_GDN_COREX_BETA_DECAY", True))
_USE_COREX_GDN_QK_MAP = (
_corex_gdn_qk_map is not None
and env_bool("BI100_GDN_COREX_QK_MAP", True))
_USE_COREX_GDN_COMBINED_QK_NORM = (
_USE_COREX_GDN_QK_MAP
and env_bool("BI100_GDN_COMBINED_QK_NORM", False))
_USE_COREX_GDN_PACKED_DECODE = (
_corex_gdn_packed_decode is not None
and env_bool("BI100_GDN_COREX_PACKED_DECODE", False))
_USE_COREX_ATTN_HEAD_RMS_NORM = (
_corex_attn_head_rms_norm is not None
and env_bool("BI100_ATTN_COREX_HEAD_RMS_NORM", True))
_USE_COREX_MOE_EXACT_REDUCE = (
_corex_moe_exact_reduce is not None
and env_bool("BI100_MOE_COREX_EXACT_REDUCE", True))
_USE_COREX_MOE_WEIGHT_GATHER = (
_corex_moe_weight_gather is not None
and env_bool("BI100_MOE_COREX_WEIGHT_GATHER", True))
_USE_COREX_MOE_DIRECT_ROUTED = (
_corex_moe_direct_routed is not None
and env_bool("BI100_MOE_COREX_DIRECT_ROUTED", False))
_USE_COREX_MOE_TOPK_SOFTMAX = (
_corex_moe_topk_softmax is not None
and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
_USE_COREX_MOE_INDEX_COMBINE = (
_corex_moe_index_combine is not None
and env_bool("BI100_MOE_COREX_INDEX_COMBINE", True))
_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
# ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API
# Source: xllm/core/layers/ilu/fused_moe.cpp → ix_moe_bridge.so
try:
from vllm.model_executor.models import ix_fused_moe as _ix_fused_moe
_HAS_IX_FUSED_MOE = _ix_fused_moe.is_available()
except ImportError:
try:
import ix_fused_moe as _ix_fused_moe
_HAS_IX_FUSED_MOE = _ix_fused_moe.is_available()
except ImportError:
_ix_fused_moe = None
_HAS_IX_FUSED_MOE = False
_USE_IX_FUSED_MOE = (
_HAS_IX_FUSED_MOE
and env_bool("BI100_MOE_IX_FUSED", True))
if _USE_IX_FUSED_MOE:
logger.info("ix_fused_moe ENABLED — full 7-step fused MoE pipeline")
else:
logger.info("ix_fused_moe unavailable — using point-optimized Python MoE")
# ---------------------------------------------------------------------------
# Qwen3.6 vision tower and vLLM 0.6 multimodal input integration
# ---------------------------------------------------------------------------
_MAX_IMAGE_TOKENS = 1280
@lru_cache(maxsize=None)
def _cached_get_qwen36_image_processor(model_path: str):
# The fast processor in transformers 4.55 calls torch.compiler APIs that
# are absent from the evaluator's torch 2.1 CoreX build.
return Qwen2VLImageProcessor.from_pretrained(model_path)
@lru_cache(maxsize=None)
def _cached_get_qwen36_tokenizer(model_path: str, trust_remote_code: bool):
return get_tokenizer(model_path, trust_remote_code=trust_remote_code)
def _image_cache_marker_tokens(image, tokenizer) -> List[int]:
array = to_numpy_array(image)
digest = hashlib.sha256()
digest.update(str(array.shape).encode("ascii"))
digest.update(str(array.dtype).encode("ascii"))
digest.update(array.tobytes())
marker = f"[image-cache-key:{digest.hexdigest()[:16]}]"
return tokenizer.encode(marker, add_special_tokens=False)
def _make_batched_images(images):
if isinstance(images, list):
if images and isinstance(images[0], list):
return [image for batch in images for image in batch]
return images
return [images]
class Qwen3_5ImagePixelInputs(TypedDict):
type: Literal["pixel_values"]
data: torch.Tensor
image_grid_thw: torch.Tensor
class Qwen3_5ImageEmbeddingInputs(TypedDict):
type: Literal["image_embeds"]
data: torch.Tensor
Qwen3_5ImageInputs = Union[Qwen3_5ImagePixelInputs,
Qwen3_5ImageEmbeddingInputs]
def _vision_pos_embed_interpolate(
embed_weight: torch.Tensor,
t: int,
h: int,
w: int,
num_grid_per_side: int,
merge_size: int,
dtype: torch.dtype,
) -> torch.Tensor:
if h % merge_size or w % merge_size:
raise ValueError(
f"vision grid {(t, h, w)} is not divisible by merge_size="
f"{merge_size}")
hidden_dim = embed_weight.shape[1]
device = embed_weight.device
h_idxs = torch.linspace(0, num_grid_per_side - 1, h,
dtype=torch.float32, device=device)
w_idxs = torch.linspace(0, num_grid_per_side - 1, w,
dtype=torch.float32, device=device)
h_floor = h_idxs.long()
w_floor = w_idxs.long()
h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
dh = h_idxs - h_floor
dw = w_idxs - w_floor
dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
hf_grid, wf_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
hc_grid, wc_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
w11 = dh_grid * dw_grid
w10 = dh_grid - w11
w01 = dw_grid - w11
w00 = 1 - dh_grid - w01
h_grid = torch.stack([hf_grid, hf_grid, hc_grid, hc_grid])
w_grid = torch.stack([wf_grid, wc_grid, wf_grid, wc_grid])
indices = (h_grid * num_grid_per_side + w_grid).reshape(4, -1)
weights = torch.stack([w00, w01, w10, w11], dim=0)
weights = weights.reshape(4, -1, 1).to(dtype=dtype)
combined = (embed_weight[indices] * weights).sum(dim=0)
combined = combined.reshape(
h // merge_size, merge_size,
w // merge_size, merge_size, hidden_dim)
combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim)
return combined.expand(t, -1, -1).reshape(-1, hidden_dim).to(dtype)
class Qwen3_5VisionPatchEmbed(nn.Module):
def __init__(self, vision_config) -> None:
super().__init__()
self.patch_size = vision_config.patch_size
self.temporal_patch_size = vision_config.temporal_patch_size
self.hidden_size = vision_config.hidden_size
kernel = (self.temporal_patch_size, self.patch_size, self.patch_size)
self.proj = nn.Conv3d(
vision_config.in_channels,
self.hidden_size,
kernel_size=kernel,
stride=kernel,
bias=True,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
length = x.shape[0]
x = x.view(length, -1, self.temporal_patch_size,
self.patch_size, self.patch_size)
return self.proj(x).view(length, self.hidden_size)
class Qwen3_5VisionMLP(nn.Module):
def __init__(self, vision_config,
quant_config: Optional[QuantizationConfig] = None) -> None:
super().__init__()
self.linear_fc1 = ColumnParallelLinear(
vision_config.hidden_size,
vision_config.intermediate_size,
bias=True,
quant_config=quant_config,
)
self.linear_fc2 = RowParallelLinear(
vision_config.intermediate_size,
vision_config.hidden_size,
bias=True,
quant_config=quant_config,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x, _ = self.linear_fc1(x)
x = F.gelu(x, approximate="tanh")
x, _ = self.linear_fc2(x)
return x
class Qwen3_5VisionBlock(nn.Module):
def __init__(self, vision_config,
quant_config: Optional[QuantizationConfig] = None) -> None:
super().__init__()
dim = vision_config.hidden_size
self.norm1 = nn.LayerNorm(dim, eps=1e-6)
self.norm2 = nn.LayerNorm(dim, eps=1e-6)
self.attn = Qwen2VisionAttention(
embed_dim=dim,
num_heads=vision_config.num_heads,
projection_size=dim,
quant_config=quant_config,
)
self.mlp = Qwen3_5VisionMLP(vision_config, quant_config)
def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor) -> torch.Tensor:
x = x + self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
)
return x + self.mlp(self.norm2(x))
class Qwen3_5VisionPatchMerger(nn.Module):
def __init__(self, vision_config,
quant_config: Optional[QuantizationConfig] = None) -> None:
super().__init__()
self.hidden_size = (vision_config.hidden_size
* vision_config.spatial_merge_size ** 2)
self.norm = nn.LayerNorm(vision_config.hidden_size, eps=1e-6)
self.linear_fc1 = ColumnParallelLinear(
self.hidden_size,
self.hidden_size,
bias=True,
quant_config=quant_config,
)
self.linear_fc2 = RowParallelLinear(
self.hidden_size,
vision_config.out_hidden_size,
bias=True,
quant_config=quant_config,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.norm(x).view(-1, self.hidden_size)
x, _ = self.linear_fc1(x)
x = F.gelu(x)
x, _ = self.linear_fc2(x)
return x
class Qwen3_5VisionTransformer(nn.Module):
def __init__(self, vision_config,
quant_config: Optional[QuantizationConfig] = None) -> None:
super().__init__()
self.hidden_size = vision_config.hidden_size
self.num_heads = vision_config.num_heads
self.spatial_merge_size = vision_config.spatial_merge_size
self.num_grid_per_side = int(vision_config.num_position_embeddings ** .5)
self.patch_embed = Qwen3_5VisionPatchEmbed(vision_config)
self.pos_embed = nn.Embedding(
vision_config.num_position_embeddings, self.hidden_size)
head_dim = self.hidden_size // self.num_heads
self.rotary_pos_emb = Qwen2VisionRotaryEmbedding(head_dim // 2)
self.blocks = nn.ModuleList([
Qwen3_5VisionBlock(vision_config, quant_config)
for _ in range(vision_config.depth)
])
self.merger = Qwen3_5VisionPatchMerger(vision_config, quant_config)
@property
def dtype(self) -> torch.dtype:
return self.patch_embed.proj.weight.dtype
@property
def device(self) -> torch.device:
return self.patch_embed.proj.weight.device
def _rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
pos_ids = []
for t, h, w in grid_thw.tolist():
h_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
w_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
h_ids = h_ids.reshape(
h // self.spatial_merge_size, self.spatial_merge_size,
w // self.spatial_merge_size, self.spatial_merge_size,
).permute(0, 2, 1, 3).flatten()
w_ids = w_ids.reshape(
h // self.spatial_merge_size, self.spatial_merge_size,
w // self.spatial_merge_size, self.spatial_merge_size,
).permute(0, 2, 1, 3).flatten()
pos_ids.append(torch.stack([h_ids, w_ids], dim=-1).repeat(t, 1))
pos_ids_t = torch.cat(pos_ids, dim=0).to(self.device)
max_grid_size = int(grid_thw[:, 1:].max().item())
return self.rotary_pos_emb(max_grid_size)[pos_ids_t].flatten(1)
def _absolute_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
return torch.cat([
_vision_pos_embed_interpolate(
self.pos_embed.weight, int(t), int(h), int(w),
self.num_grid_per_side, self.spatial_merge_size, self.dtype)
for t, h, w in grid_thw.tolist()
], dim=0)
def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
x = x.to(device=self.device, dtype=self.dtype)
grid_thw = grid_thw.to(device=self.device)
x = self.patch_embed(x)
x = x + self._absolute_pos_emb(grid_thw)
rotary_pos_emb = self._rot_pos_emb(grid_thw)
cu_seqlens = torch.repeat_interleave(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0],
).cumsum(dim=0, dtype=torch.int32)
cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
x = x.unsqueeze(1)
for block in self.blocks:
x = block(x, cu_seqlens, rotary_pos_emb)
return self.merger(x)
class Qwen3_5InterleavedMRotaryEmbedding(MRotaryEmbedding):
"""Qwen3.5 frequency-interleaved T/H/W rotary embedding."""
def forward(self, positions: torch.Tensor, query: torch.Tensor,
key: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
if positions.ndim not in (1, 2):
raise ValueError(f"invalid MRoPE positions shape {positions.shape}")
num_tokens = positions.shape[-1]
cos_sin = self.cos_sin_cache[positions]
cos_all, sin_all = cos_sin.chunk(2, dim=-1)
if positions.ndim == 2:
if not self.mrope_section:
raise ValueError("mrope_section is required")
cos = cos_all[0].clone()
sin = sin_all[0].clone()
for dim, offset in enumerate((1, 2), start=1):
stop = self.mrope_section[dim] * 3
cos[..., offset:stop:3] = cos_all[dim, ..., offset:stop:3]
sin[..., offset:stop:3] = sin_all[dim, ..., offset:stop:3]
else:
cos, sin = cos_all, sin_all
query_shape = query.shape
query = query.view(num_tokens, -1, self.head_size)
query_rot = _apply_rotary_emb(
query[..., :self.rotary_dim], cos, sin, self.is_neox_style)
query = torch.cat((query_rot, query[..., self.rotary_dim:]), dim=-1)
key_shape = key.shape
key = key.view(num_tokens, -1, self.head_size)
key_rot = _apply_rotary_emb(
key[..., :self.rotary_dim], cos, sin, self.is_neox_style)
key = torch.cat((key_rot, key[..., self.rotary_dim:]), dim=-1)
return query.reshape(query_shape), key.reshape(key_shape)
def _qwen36_pixel_limits(image_processor) -> Tuple[int, int]:
min_pixels = 256 * 256
configured_max = 4096 * 4096
runtime_max = _MAX_IMAGE_TOKENS * (
image_processor.patch_size * image_processor.merge_size) ** 2
return min_pixels, min(configured_max, runtime_max)
def _qwen36_image_token_count(image, image_processor) -> int:
if isinstance(image, Image.Image):
image = image.convert("RGB")
image_array = to_numpy_array(image)
height, width = get_image_size(
image_array, channel_dim=ChannelDimension.LAST)
min_pixels, max_pixels = _qwen36_pixel_limits(image_processor)
if getattr(image_processor, "do_resize", True):
height, width = smart_resize(
height=height,
width=width,
factor=image_processor.patch_size * image_processor.merge_size,
min_pixels=min_pixels,
max_pixels=max_pixels,
)
return (height // image_processor.patch_size
* width // image_processor.patch_size
// image_processor.merge_size ** 2)
def qwen36_image_input_mapper(
ctx: InputContext,
data: MultiModalData[object],
) -> MultiModalInputs:
if isinstance(data, dict):
return MultiModalInputs({
"image_embeds": data.get("image_embeds"),
"image_grid_thw": data.get("image_grid_thw"),
})
image_processor = _cached_get_qwen36_image_processor(
ctx.model_config.model)
min_pixels, max_pixels = _qwen36_pixel_limits(image_processor)
batch_data = image_processor.preprocess(
images=data,
return_tensors="pt",
size={"shortest_edge": min_pixels, "longest_edge": max_pixels},
do_convert_rgb=True,
input_data_format=ChannelDimension.LAST,
).data
return MultiModalInputs(batch_data)
def get_max_qwen36_image_tokens(_ctx: InputContext) -> int:
return _MAX_IMAGE_TOKENS
def dummy_data_for_qwen36(
ctx: InputContext,
seq_len: int,
mm_counts: Mapping[str, int],
) -> Tuple[SequenceData, Optional[MultiModalDataDict]]:
num_images = mm_counts.get("image", 0)
image_tokens = _MAX_IMAGE_TOKENS * num_images
if seq_len < image_tokens + 2:
raise RuntimeError(
f"Qwen3.6 needs {image_tokens + 2} tokens for {num_images} "
f"max-size image(s), but max_model_len is {seq_len}")
config = ctx.model_config.hf_config
seq_data = SequenceData.from_token_counts(
(config.vision_start_token_id, 1),
(config.image_token_id, image_tokens),
(config.vision_end_token_id, 1),
(0, seq_len - image_tokens - 2),
)
dummy_image = Image.new("RGB", (1280, 1024), color=0)
return seq_data, {
"image": (dummy_image if num_images == 1
else [dummy_image] * num_images)
}
def input_processor_for_qwen36(ctx: InputContext,
llm_inputs: LLMInputs) -> LLMInputs:
multi_modal_data = llm_inputs.get("multi_modal_data")
if not multi_modal_data or "image" not in multi_modal_data:
return llm_inputs
images = multi_modal_data["image"]
prompt_token_ids = llm_inputs.get("prompt_token_ids")
if prompt_token_ids is None:
raise ValueError("Qwen3.6 image requests require tokenized prompt input")
config = ctx.model_config.hf_config
image_processor = _cached_get_qwen36_image_processor(
ctx.model_config.model)
tokenizer = _cached_get_qwen36_tokenizer(
ctx.model_config.tokenizer,
ctx.model_config.trust_remote_code,
)
batched_images = _make_batched_images(images)
image_indices = [
idx for idx, token in enumerate(prompt_token_ids)
if token == config.image_token_id
]
if len(image_indices) != len(batched_images):
raise ValueError(
f"found {len(image_indices)} image placeholders for "
f"{len(batched_images)} image(s)")
expanded = []
previous = 0
for index, image in zip(image_indices, batched_images):
vision_start = index - 1
if (vision_start < previous
or prompt_token_ids[vision_start]
!= config.vision_start_token_id):
raise ValueError("image token is not preceded by vision_start")
expanded.extend(prompt_token_ids[previous:vision_start])
expanded.extend(_image_cache_marker_tokens(image, tokenizer))
expanded.extend(prompt_token_ids[vision_start:index])
expanded.extend([config.image_token_id]
* _qwen36_image_token_count(image, image_processor))
previous = index + 1
expanded.extend(prompt_token_ids[previous:])
return LLMInputs(
prompt_token_ids=expanded,
prompt=llm_inputs["prompt"],
multi_modal_data=multi_modal_data,
)
# ---------------------------------------------------------------------------
# Pure-PyTorch DeltaNet kernels (fallbacks from transformers 5.2.0)
# ---------------------------------------------------------------------------
def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
def _check_gdn_finite(tensor: torch.Tensor, *, layer_idx: int,
stage: str) -> torch.Tensor:
if not _GDN_FINITE_CHECK:
return tensor
if torch.isfinite(tensor).all():
return tensor
bad = (~torch.isfinite(tensor)).float().mean().item()
msg = (
f"non-finite values in {stage} GatedDeltaNet layer {layer_idx} "
f"(frac={bad:.4f})"
)
if not _ALLOW_GDN_NAN_ZERO:
raise RuntimeError(msg)
logger.warning("%s; replacing with zeros because BI100_GDN_ALLOW_NAN_ZERO=1",
msg)
return torch.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0)
def _gdn_segment_ends(seq_len: int, chunk_size: int,
capture_offsets: Iterable[int]) -> List[int]:
ends = list(range(chunk_size, seq_len, chunk_size))
ends.append(seq_len)
ends.extend(offset for offset in capture_offsets
if 0 < offset < seq_len)
return sorted(set(ends))
def _validate_gdn_prefix_key(key: Any) -> Tuple[int, bytes]:
if (not isinstance(key, tuple) or len(key) != 2
or not isinstance(key[0], int) or key[0] <= 0
or not isinstance(key[1], bytes) or len(key[1]) != 32):
raise RuntimeError(f"invalid GDN prefix key: {key!r}")
return key
def _torch_causal_conv1d_update(
hidden_states: torch.Tensor, # (batch, channels, seq=1)
conv_state: torch.Tensor, # (batch, channels, state_len) modified in-place
weight: torch.Tensor, # (channels, kernel_size)
bias: Optional[torch.Tensor] = None,
activation: Optional[str] = None,
) -> torch.Tensor:
_, channels, seq_len = hidden_states.shape
state_len = conv_state.shape[-1]
cat = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
conv_state.copy_(cat[:, :, -state_len:])
out = F.conv1d(cat, weight.unsqueeze(1), bias, padding=0, groups=channels)
out = out[:, :, -seq_len:]
if activation is not None:
out = F.silu(out)
return out.to(hidden_states.dtype)
def _torch_chunk_gated_delta_rule(
query: torch.Tensor, # (batch, seq, num_heads, head_k_dim)
key: torch.Tensor,
value: torch.Tensor, # (batch, seq, num_heads, head_v_dim)
g: torch.Tensor, # (batch, seq, num_heads)
beta: torch.Tensor, # (batch, seq, num_heads)
chunk_size: int = 64,
initial_state: Optional[torch.Tensor] = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
if use_qk_l2norm_in_kernel:
query = _l2norm(query)
key = _l2norm(key)
# Transpose to (batch, num_heads, seq, dim)
query, key, value, beta, g = [
x.transpose(1, 2).contiguous().to(torch.float32)
for x in (query, key, value, beta, g)
]
batch, num_heads, seq_len, k_dim = key.shape
v_dim = value.shape[-1]
pad = (chunk_size - seq_len % chunk_size) % chunk_size
query = F.pad(query, (0, 0, 0, pad))
key = F.pad(key, (0, 0, 0, pad))
value = F.pad(value, (0, 0, 0, pad))
beta = F.pad(beta, (0, pad))
g = F.pad(g, (0, pad))
total_len = seq_len + pad
scale = 1.0 / (query.shape[-1] ** 0.5)
query = query * scale
v_beta = value * beta.unsqueeze(-1)
k_beta = key * beta.unsqueeze(-1)
query, key, value, k_beta, v_beta = [
x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
for x in (query, key, value, k_beta, v_beta)
]
g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
mask_upper = torch.triu(
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
diagonal=0)
g = g.cumsum(dim=-1)
decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0)
for i in range(1, chunk_size):
row = attn[..., i, :i].clone()
sub = attn[..., :i, :i].clone()
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
last_state = (
torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device)
if initial_state is None
else initial_state.to(value)
)
core_out = torch.zeros_like(value)
mask_upper2 = torch.triu(
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
diagonal=1)
for i in range(total_len // chunk_size):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
v_prime = k_cumdecay[:, :, i] @ last_state
v_new = v_i - v_prime
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
core_out[:, :, i] = attn_inter + attn_i @ v_new
last_state = (
last_state * g[:, :, i, -1, None, None].exp()
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
.transpose(-1, -2) @ v_new
)
if not output_final_state:
last_state = None
core_out = core_out.reshape(batch, num_heads, -1, v_dim)[:, :, :seq_len]
core_out = core_out.transpose(1, 2).contiguous()
return core_out, last_state
def _torch_recurrent_gated_delta_rule(
query: torch.Tensor, # (batch, 1, num_heads, head_k_dim)
key: torch.Tensor,
value: torch.Tensor,
g: torch.Tensor, # (batch, 1, num_heads)
beta: torch.Tensor,
initial_state: Optional[torch.Tensor] = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
if use_qk_l2norm_in_kernel:
query = _l2norm(query)
key = _l2norm(key)
query, key, value, beta, g = [
x.transpose(1, 2).contiguous().to(torch.float32)
for x in (query, key, value, beta, g)
]
batch, num_heads, seq_len, k_dim = key.shape
v_dim = value.shape[-1]
scale = 1.0 / (query.shape[-1] ** 0.5)
query = query * scale
core_out = torch.zeros(batch, num_heads, seq_len, v_dim,
dtype=value.dtype, device=value.device)
last_state = (
torch.zeros(batch, num_heads, k_dim, v_dim,
dtype=value.dtype, device=value.device)
if initial_state is None
else initial_state.to(value)
)
for t in range(seq_len):
q_t = query[:, :, t]
k_t = key[:, :, t]
v_t = value[:, :, t]
g_t = g[:, :, t].exp().unsqueeze(-1).unsqueeze(-1)
beta_t = beta[:, :, t].unsqueeze(-1)
last_state = last_state * g_t
kv_mem = (last_state * k_t.unsqueeze(-1)).sum(dim=-2)
delta = (v_t - kv_mem) * beta_t
last_state = last_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
core_out[:, :, t] = (last_state * q_t.unsqueeze(-1)).sum(dim=-2)
if not output_final_state:
last_state = None
core_out = core_out.transpose(1, 2).contiguous()
return core_out, last_state
# ---------------------------------------------------------------------------
# Gated RMSNorm (for DeltaNet output normalisation)
# ---------------------------------------------------------------------------
class Qwen3_5RMSNormGated(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor,
gate: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hs = hidden_states.to(torch.float32)
variance = hs.pow(2).mean(-1, keepdim=True)
hs = hs * torch.rsqrt(variance + self.variance_epsilon)
hs = self.weight * hs.to(input_dtype)
return (hs * F.silu(gate.to(torch.float32))).to(input_dtype)
def forward_decode(self, hidden_states: torch.Tensor,
gate: torch.Tensor) -> torch.Tensor:
if (_USE_COREX_GDN_GATED_NORM
and hidden_states.dtype == torch.float32
and gate.dtype == torch.float16
and self.weight.dtype == torch.float16
and hidden_states.shape[-1] == 128):
hs = hidden_states.float()
inverse = torch.rsqrt(
hs.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
return _corex_gdn_gated_norm.apply_inverse(
hs, gate, self.weight, inverse)
return self.forward(hidden_states, gate).to(gate.dtype)
def _load_gdn_projection_weight(params_dict, name: str,
loaded_weight: torch.Tensor,
text_cfg) -> bool:
projections = {
"in_proj_qkv": None,
"in_proj_z": 3,
"in_proj_b": 4,
"in_proj_a": 5,
}
source = next((projection for projection in projections
if f".linear_attn.{projection}." in name), None)
if source is None:
return False
target_name = name.replace(
f".linear_attn.{source}.",
".linear_attn.in_proj_qkvzba.",
)
if target_name not in params_dict:
raise ValueError(f"missing fused GDN projection parameter: {target_name}")
param = params_dict[target_name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
if source == "in_proj_qkv":
key_dim = (text_cfg.linear_num_key_heads
* text_cfg.linear_key_head_dim)
value_dim = (text_cfg.linear_num_value_heads
* text_cfg.linear_value_head_dim)
shard_sizes = (key_dim, key_dim, value_dim)
if loaded_weight.shape[0] != sum(shard_sizes):
raise ValueError(
"unexpected fused QKV output size: "
f"{loaded_weight.shape[0]} != {sum(shard_sizes)}")
for shard_id, shard in enumerate(
torch.split(loaded_weight, shard_sizes, dim=0)):
weight_loader(param, shard, shard_id)
else:
weight_loader(param, loaded_weight, projections[source])
return True
def _load_full_attention_qgkv_weight(params_dict, name: str,
loaded_weight: torch.Tensor,
text_cfg) -> bool:
projections = {"q_proj": 0, "k_proj": 1, "v_proj": 2}
source = next((projection for projection in projections
if f".self_attn.{projection}." in name), None)
if source is None:
return False
target_name = name.replace(
f".self_attn.{source}.", ".self_attn.qgkv_proj.")
if target_name not in params_dict:
return False
tp_size = get_tensor_model_parallel_world_size()
tp_rank = get_tensor_model_parallel_rank()
qg_dim = text_cfg.num_attention_heads * text_cfg.head_dim * 2
if qg_dim % tp_size != 0:
raise ValueError(f"QG output size {qg_dim} is not divisible by TP {tp_size}")
local_qg_dim = qg_dim // tp_size
kv_dim = text_cfg.num_key_value_heads * text_cfg.head_dim
expected_rows = qg_dim if source == "q_proj" else kv_dim
if loaded_weight.shape[0] != expected_rows:
raise ValueError(
f"unexpected full-attention {source} output size: "
f"{loaded_weight.shape[0]} != {expected_rows}")
if source == "q_proj":
loaded_weight = loaded_weight.narrow(
0, tp_rank * local_qg_dim, local_qg_dim)
offset = 0
elif source == "k_proj":
offset = local_qg_dim
else:
offset = local_qg_dim + kv_dim
param = params_dict[target_name]
default_weight_loader(
param[offset:offset + loaded_weight.shape[0]], loaded_weight)
return True
# ---------------------------------------------------------------------------
# Gated DeltaNet (linear_attention layers)
# ---------------------------------------------------------------------------
class GatedDeltaNet(nn.Module):
def __init__(
self,
text_cfg,
layer_idx: int,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.layer_idx = layer_idx
self.hidden_size = text_cfg.hidden_size
self.num_v_heads = text_cfg.linear_num_value_heads # checkpoint: 32
self.num_k_heads = text_cfg.linear_num_key_heads # checkpoint: 16
self.head_k_dim = text_cfg.linear_key_head_dim # 128
self.head_v_dim = text_cfg.linear_value_head_dim # 128
self.key_dim = self.num_k_heads * self.head_k_dim # 2048
self.value_dim = self.num_v_heads * self.head_v_dim # checkpoint: 4096
self.conv_dim = self.key_dim * 2 + self.value_dim # checkpoint: 8192
self.conv_kernel_size = text_cfg.linear_conv_kernel_dim # 4
self.head_expand_ratio = self.num_v_heads // self.num_k_heads # checkpoint: 2
tp_size = get_tensor_model_parallel_world_size()
# Keep each logical projection independently TP-sharded while executing
# one GEMM. Per-rank output order is [q, k, v, z, beta, decay].
self.in_proj_qkvzba = MergedColumnParallelLinear(
self.hidden_size,
[self.key_dim, self.key_dim, self.value_dim, self.value_dim,
self.num_v_heads, self.num_v_heads],
bias=False, quant_config=quant_config)
self.out_proj = RowParallelLinear(
self.value_dim, self.hidden_size,
bias=False, quant_config=quant_config)
# Depthwise conv weight — sharded along channel dim (dim 0)
local_conv_dim = self.conv_dim // tp_size
self.conv1d_weight = nn.Parameter(
torch.empty(local_conv_dim, 1, self.conv_kernel_size))
set_weight_attrs(self.conv1d_weight, {
"weight_loader": self._conv1d_weight_loader})
# Per-head scalar parameters — sharded along dim 0
local_num_v = self.num_v_heads // tp_size
self.A_log = nn.Parameter(torch.zeros(local_num_v))
self.dt_bias = nn.Parameter(torch.zeros(local_num_v))
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
# Gated RMSNorm on head_v_dim — replicated (head_v_dim=128 is small)
self.norm = Qwen3_5RMSNormGated(self.head_v_dim,
eps=text_cfg.rms_norm_eps)
self.captured_conv_states: Dict[int, torch.Tensor] = {}
self.captured_temporal_states: Dict[int, torch.Tensor] = {}
def _conv1d_weight_loader(self, param: torch.Tensor,
loaded_weight: torch.Tensor) -> None:
# loaded_weight is ordered as [q, k, v] along its channel dimension.
# Must gather channels in the same non-contiguous pattern that
# MergedColumnParallelLinear uses for in_proj_qkv, so that each rank's
# conv1d_weight[i] applies to the correct in_proj_qkv output channel.
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
key_local = self.key_dim // tp_size # 512 with TP=4
val_local = self.value_dim // tp_size # 1024 with TP=4
q_s = loaded_weight[tp_rank * key_local : (tp_rank + 1) * key_local]
k_s = loaded_weight[self.key_dim + tp_rank * key_local :
self.key_dim + (tp_rank + 1) * key_local]
v_s = loaded_weight[2 * self.key_dim + tp_rank * val_local :
2 * self.key_dim + (tp_rank + 1) * val_local]
param.data.copy_(torch.cat([q_s, k_s, v_s], dim=0))
def forward(
self,
hidden_states: torch.Tensor, # (total_tokens, hidden_size)
attn_metadata: AttentionMetadata,
conv_state: torch.Tensor, # (batch, local_conv_dim, kernel-1) in-place
temporal_state: torch.Tensor, # (batch, local_v_heads, k_dim, v_dim) in-place
capture_offsets: Optional[Iterable[int]] = None,
segment_offsets: Optional[Iterable[int]] = None,
) -> torch.Tensor:
tp_size = get_tensor_model_parallel_world_size()
local_key_dim = self.key_dim // tp_size
local_val_dim = self.value_dim // tp_size
local_num_v = self.num_v_heads // tp_size
local_num_k = self.num_k_heads // tp_size
local_conv_dim = self.conv_dim // tp_size
self.captured_conv_states = {}
self.captured_temporal_states = {}
is_prefill = attn_metadata.num_prefill_tokens > 0
projected, _ = self.in_proj_qkvzba(hidden_states)
mixed_qkv_all, z_all, b_all, a_all = torch.split(
projected,
[local_conv_dim, local_val_dim, local_num_v, local_num_v],
dim=-1,
)
if is_prefill:
seq_starts = attn_metadata.query_start_loc.tolist()
outputs = []
state_len = self.conv_kernel_size - 1
weight_2d = self.conv1d_weight.squeeze(1) # (local_conv_dim, kernel)
for si in range(len(seq_starts) - 1):
s, e = int(seq_starts[si]), int(seq_starts[si + 1])
seq_len = e - s
# Shape: (1, local_conv_dim, seq_len)
mixed_qkv = (mixed_qkv_all[s:e]
.transpose(0, 1).unsqueeze(0)
.to(weight_2d.dtype))
# Load prev conv state BEFORE overwriting (needed for causal conv padding).
# For first prefill of a request: mamba_cache is zeros → correct.
# For chunked prefill chunk 2+: carries last state_len tokens from prev chunk.
prev_conv = conv_state[si:si + 1].clone().to(weight_2d.dtype) # [1, local_conv_dim, state_len]
# Save conv state (last state_len positions)
if seq_len >= state_len:
conv_state[si].copy_(mixed_qkv[0, :, -state_len:])
else:
conv_state[si, :, state_len - seq_len:].copy_(
mixed_qkv[0])
conv_state[si, :, :state_len - seq_len] = 0
# Causal conv: left-pad with previous conv state (not zeros).
padded = torch.cat([prev_conv, mixed_qkv], dim=2)
seq_capture_offsets = (set(capture_offsets or ())
if si == 0 else set())
seq_segment_offsets = (set(segment_offsets or ())
if si == 0 else set())
for capture_offset in seq_capture_offsets:
if 0 < capture_offset < seq_len:
self.captured_conv_states[capture_offset] = padded[
0, :, capture_offset:
capture_offset + state_len].clone()
mixed_qkv_conv = F.conv1d(
padded, self.conv1d_weight,
bias=None, padding=0, groups=local_conv_dim)
mixed_qkv_conv = F.silu(mixed_qkv_conv)
# (1, seq_len, local_conv_dim)
mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0)
q, k, v = torch.split(
mixed_qkv_conv,
[local_key_dim, local_key_dim, local_val_dim], dim=-1)
q = q.reshape(1, seq_len, local_num_k, self.head_k_dim)
k = k.reshape(1, seq_len, local_num_k, self.head_k_dim)
v = v.reshape(1, seq_len, local_num_v, self.head_v_dim)
beta = b_all[s:e].sigmoid().unsqueeze(0) # (1, seq_len, local_num_v)
g = (-self.A_log.float().exp()
* F.softplus(a_all[s:e].float() + self.dt_bias)
).unsqueeze(0) # (1, seq_len, local_num_v)
# Expand k/q to match num_v_heads
q = q.repeat_interleave(self.head_expand_ratio, dim=2)
k = k.repeat_interleave(self.head_expand_ratio, dim=2)
# Sub-sequence chunking: call _torch_chunk_gated_delta_rule
# on _DNN_CHUNK tokens at a time to cap peak memory.
# Full 18K: tensors [1,6,282,64,64]=220 MB each → ~990 MB/call.
# With _DNN_CHUNK=4096: [1,6,64,64,64]=6 MB each → ~137 MB/call.
# State is chained via initial_state / output_final_state.
cur_state = temporal_state[si:si + 1].clone()
core_out_parts = []
segment_ends = _gdn_segment_ends(
seq_len, _DNN_CHUNK_SIZE,
seq_capture_offsets | seq_segment_offsets)
sc_start = 0
_chunk_fn = (
_corex_gdn_chunk_recurrent.torch_chunk_gated_delta_rule
if _HAS_COREX_GDN_CHUNK
else _torch_chunk_gated_delta_rule
)
with bi100_timer(f"L{self.layer_idx}.gdn.prefill"):
for sc_end in segment_ends:
c_out, cur_state = _chunk_fn(
q[:, sc_start:sc_end],
k[:, sc_start:sc_end],
v[:, sc_start:sc_end],
g[:, sc_start:sc_end],
beta[:, sc_start:sc_end],
initial_state=cur_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
)
core_out_parts.append(c_out)
if sc_end in seq_capture_offsets:
self.captured_temporal_states[sc_end] = (
cur_state[0].clone())
sc_start = sc_end
if cur_state is not None:
temporal_state[si].copy_(cur_state[0])
# [1, seq_len, num_v_heads, head_v_dim]
core_out = torch.cat(core_out_parts, dim=1)
# Gate + norm + output proj
z = z_all[s:e].reshape(seq_len, local_num_v, self.head_v_dim)
core_out = core_out.reshape(seq_len, local_num_v, self.head_v_dim)
normed = self.norm(
core_out.reshape(-1, self.head_v_dim),
z.reshape(-1, self.head_v_dim))
normed = _check_gdn_finite(
normed, layer_idx=self.layer_idx,
stage="prefill-norm").reshape(seq_len, -1)
normed = normed.to(z_all.dtype)
out, _ = self.out_proj(normed)
outputs.append(out)
result = torch.cat(outputs, dim=0)
return _check_gdn_finite(
result, layer_idx=self.layer_idx, stage="prefill-output")
else:
# Decode: one token per sequence
num_seqs = hidden_states.shape[0]
weight_2d = self.conv1d_weight.squeeze(1)
# (num_seqs, local_conv_dim, 1)
mixed_qkv = (mixed_qkv_all
.to(weight_2d.dtype)
.unsqueeze(-1))
if _USE_COREX_GDN_CAUSAL_CONV:
mixed_qkv_conv = _corex_gdn_causal_conv.causal_conv_update(
conv_state, mixed_qkv, weight_2d)
else:
mixed_qkv_conv = _torch_causal_conv1d_update(
mixed_qkv, conv_state, weight_2d,
bias=None, activation='silu')
# (num_seqs, local_conv_dim, 1) → (num_seqs, 1, local_conv_dim)
mixed_qkv_conv = mixed_qkv_conv.squeeze(-1).unsqueeze(1)
packed_mixed_qkv = mixed_qkv_conv.squeeze(1)
use_corex_packed_decode = (
_USE_COREX_GDN_PACKED_DECODE
and num_seqs == 1
and local_num_k == 4
and local_num_v == 8
and self.head_k_dim == 128
and self.head_v_dim == 128
and packed_mixed_qkv.dtype == torch.float16
and packed_mixed_qkv.shape == (1, 2048)
and packed_mixed_qkv.is_contiguous()
and b_all.dtype == torch.float16
and b_all.shape == (1, 8)
and b_all.is_contiguous()
and a_all.dtype == torch.float16
and a_all.shape == (1, 8)
and a_all.is_contiguous()
and self.A_log.dtype == torch.float16
and self.A_log.shape == (8,)
and self.A_log.is_contiguous()
and self.dt_bias.dtype == torch.float16
and self.dt_bias.shape == (8,)
and self.dt_bias.is_contiguous()
and temporal_state.dtype == torch.float32
and temporal_state.shape == (1, 8, 128, 128)
and temporal_state.is_contiguous())
if use_corex_packed_decode:
with bi100_timer(f"L{self.layer_idx}.gdn.decode"):
core_out = _corex_gdn_packed_decode.packed_decode(
temporal_state, packed_mixed_qkv, b_all, a_all,
self.A_log, self.dt_bias)
else:
q, k, v = torch.split(
mixed_qkv_conv,
[local_key_dim, local_key_dim, local_val_dim], dim=-1)
q = q.reshape(num_seqs, 1, local_num_k, self.head_k_dim)
k = k.reshape(num_seqs, 1, local_num_k, self.head_k_dim)
v = v.reshape(num_seqs, 1, local_num_v, self.head_v_dim)
use_corex_beta_decay = (
_USE_COREX_GDN_BETA_DECAY
and b_all.dtype == torch.float16
and a_all.dtype == torch.float16
and self.A_log.dtype == torch.float16
and self.dt_bias.dtype == torch.float16
and b_all.is_contiguous()
and a_all.is_contiguous())
if use_corex_beta_decay:
beta_decay = _corex_gdn_beta_decay.beta_decay(
b_all, a_all, self.A_log, self.dt_bias)
bt = beta_decay[0]
g_t = beta_decay[1]
else:
beta = b_all.sigmoid()
g = (-self.A_log.float().exp()
* F.softplus(a_all.float() + self.dt_bias))
bt = beta.float()
g_t = g.float().exp_()
# Inlined decode recurrent step (seq_len=1).
# Uses bmm/baddbmm_ to avoid large intermediate tensors.
_scale = self.head_k_dim ** -0.5
q_raw = q.squeeze(1)
k_raw = k.squeeze(1)
use_corex_qk_map = (
_USE_COREX_GDN_QK_MAP
and q_raw.dtype == torch.float16
and k_raw.dtype == torch.float16
and self.head_k_dim == 128
and q_raw.is_contiguous()
and k_raw.is_contiguous())
if use_corex_qk_map:
use_combined_qk_norm = (
_USE_COREX_GDN_COMBINED_QK_NORM
and num_seqs == 1
and local_num_k == 4
and local_num_v == 8
and packed_mixed_qkv.dtype == torch.float16
and packed_mixed_qkv.shape == (1, 2048)
and packed_mixed_qkv.is_contiguous())
if use_combined_qk_norm:
raw_qk = packed_mixed_qkv.narrow(
1, 0, 2 * local_key_dim).view(
num_seqs, 2 * local_num_k,
self.head_k_dim)
normalized_qk = _l2norm(raw_qk)
normalized_q, normalized_k = torch.split(
normalized_qk, local_num_k, dim=1)
else:
normalized_q = _l2norm(q_raw)
normalized_k = _l2norm(k_raw)
qk_mapped = _corex_gdn_qk_map.qk_map(
normalized_q, normalized_k, local_num_v)
q_t = qk_mapped[0]
k_t = qk_mapped[1]
else:
q_expanded = q_raw.repeat_interleave(
self.head_expand_ratio, dim=1)
k_expanded = k_raw.repeat_interleave(
self.head_expand_ratio, dim=1)
q_t = _l2norm(q_expanded).float() * _scale
k_t = _l2norm(k_expanded).float()
v_t = v.squeeze(1).float()
with bi100_timer(f"L{self.layer_idx}.gdn.decode"):
# State shape is (B, H_v, k_dim, v_dim).
temporal_state.mul_(g_t[:, :, None, None])
ts_flat = temporal_state.view(
-1, self.head_k_dim, self.head_v_dim)
BH = ts_flat.shape[0]
kv_mem = torch.bmm(
k_t.view(BH, 1, self.head_k_dim), ts_flat
).view(num_seqs, local_num_v, self.head_v_dim)
delta = (v_t - kv_mem) * bt[:, :, None]
ts_flat.baddbmm_(
k_t.view(BH, self.head_k_dim, 1),
delta.view(BH, 1, self.head_v_dim),
)
core_out = torch.bmm(
q_t.view(BH, 1, self.head_k_dim), ts_flat
).view(num_seqs, local_num_v, self.head_v_dim)
# core_out: (B, H_v, v_dim) = (num_seqs, local_num_v, head_v_dim) already
z = z_all.reshape(num_seqs, local_num_v, self.head_v_dim)
normed = self.norm.forward_decode(
core_out.reshape(-1, self.head_v_dim),
z.reshape(-1, self.head_v_dim))
normed = _check_gdn_finite(
normed, layer_idx=self.layer_idx,
stage="decode-norm").reshape(num_seqs, -1)
out, _ = self.out_proj(normed)
return _check_gdn_finite(
out, layer_idx=self.layer_idx, stage="decode-output")
# ---------------------------------------------------------------------------
# Full Attention (with gated q — unique to Qwen3.5)
# ---------------------------------------------------------------------------
class Qwen3_5AttentionHeadRMSNorm(GemmaRMSNorm):
def forward_cuda(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
):
if (_USE_COREX_ATTN_HEAD_RMS_NORM
and residual is None
and x.dtype == torch.float16
and self.weight.dtype == torch.float16
and x.dim() == 3
and x.shape[0] == 1
and x.shape[-1] == 256
and x.is_contiguous()
and self.weight.is_contiguous()):
original_shape = x.shape
converted, squares = _corex_attn_head_rms_norm.prepare(
x.view(-1, 256))
inverse = torch.rsqrt(
squares.mean(dim=-1, keepdim=True)
+ self.variance_epsilon)
return _corex_attn_head_rms_norm.apply_inverse(
converted, self.weight, inverse).view(original_shape)
return super().forward_cuda(x, residual)
class Qwen3_5FullAttention(nn.Module):
def __init__(
self,
text_cfg,
layer_idx: int,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.layer_idx = layer_idx
self.hidden_size = text_cfg.hidden_size # 5120
self.num_heads = text_cfg.num_attention_heads # 24
self.num_kv_heads = text_cfg.num_key_value_heads # 4
self.head_dim = text_cfg.head_dim # 256
self.rms_norm_eps = text_cfg.rms_norm_eps
tp_size = get_tensor_model_parallel_world_size()
self.local_num_heads = self.num_heads // tp_size
self.scaling = self.head_dim ** -0.5
self.use_packed_local_qgkv = tp_size > self.num_kv_heads
# When num_kv_heads < tp_size we cannot shard KV further (would give
# fractional heads per rank). Use ReplicatedLinear so every rank holds
# all KV heads; local_num_kv_heads equals the full count.
# When num_kv_heads >= tp_size standard ColumnParallel sharding applies.
if tp_size > self.num_kv_heads:
# GQA-aware TP sharding: ixformer kernel only supports num_kv_heads=1
# per rank. With num_kv_heads=2 < tp_size=4 we cannot shard KV
# evenly, but we CAN assign each rank the ONE KV head that serves
# its Q heads:
# q_per_kv = num_heads // num_kv_heads (e.g. 16//2 = 8)
# Rank r uses KV head r * local_num_heads // q_per_kv
# e.g. ranks 0,1 → KV head 0; ranks 2,3 → KV head 1.
# We replicate all KV heads to every rank and select in forward().
self.proj_kv_heads = self.num_kv_heads # heads available from projection
self.local_num_kv_heads = 1 # heads after rank-local selection
self.q_per_kv_global = self.num_heads // self.num_kv_heads
local_qg_dim = self.local_num_heads * self.head_dim * 2
replicated_kv_dim = self.num_kv_heads * self.head_dim
self.qgkv_proj = ReplicatedLinear(
self.hidden_size, local_qg_dim + 2 * replicated_kv_dim,
bias=False, quant_config=quant_config,
prefix=f"{prefix}.qgkv_proj")
else:
# Standard sharding: each rank gets num_kv_heads // tp_size heads.
self.local_num_kv_heads = self.num_kv_heads // tp_size
self.proj_kv_heads = self.local_num_kv_heads # already sharded
self.q_per_kv_global = None
self.k_proj = ColumnParallelLinear(
self.hidden_size, self.num_kv_heads * self.head_dim,
bias=False, quant_config=quant_config,
prefix=f"{prefix}.k_proj")
self.v_proj = ColumnParallelLinear(
self.hidden_size, self.num_kv_heads * self.head_dim,
bias=False, quant_config=quant_config,
prefix=f"{prefix}.v_proj")
self.local_q_dim = self.local_num_heads * self.head_dim
self.local_kv_dim = self.local_num_kv_heads * self.head_dim
if not self.use_packed_local_qgkv:
# q_proj includes gate: output = num_heads * head_dim * 2
self.q_proj = ColumnParallelLinear(
self.hidden_size, self.num_heads * self.head_dim * 2,
bias=False, quant_config=quant_config,
prefix=f"{prefix}.q_proj")
self.o_proj = RowParallelLinear(
self.num_heads * self.head_dim, self.hidden_size,
bias=False, quant_config=quant_config,
prefix=f"{prefix}.o_proj")
self.q_norm = Qwen3_5AttentionHeadRMSNorm(
self.head_dim, eps=self.rms_norm_eps)
self.k_norm = Qwen3_5AttentionHeadRMSNorm(
self.head_dim, eps=self.rms_norm_eps)
# Partial RoPE: rotary_dim = head_dim * partial_rotary_factor = 256 * 0.25 = 64
rope_params = getattr(text_cfg, "rope_parameters", {}) or {}
rope_theta = rope_params.get("rope_theta", 10_000_000)
partial_factor = rope_params.get("partial_rotary_factor", 0.25)
rotary_dim = int(self.head_dim * partial_factor)
self.rotary_emb = Qwen3_5InterleavedMRotaryEmbedding(
head_size=self.head_dim,
rotary_dim=rotary_dim,
max_position_embeddings=text_cfg.max_position_embeddings,
base=rope_theta,
is_neox_style=True,
dtype=torch.get_default_dtype(),
mrope_section=rope_params.get("mrope_section", [11, 11, 10]),
)
self.attn = Attention(
self.local_num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.local_num_kv_heads,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: torch.Tensor,
attn_metadata: AttentionMetadata,
) -> torch.Tensor:
total_tokens = hidden_states.shape[0]
with bi100_timer("full_attn.project_qgkv"):
if self.use_packed_local_qgkv:
projected, _ = self.qgkv_proj(hidden_states)
qg, k, v = torch.split(
projected,
[self.local_num_heads * self.head_dim * 2,
self.proj_kv_heads * self.head_dim,
self.proj_kv_heads * self.head_dim],
dim=-1)
else:
qg, _ = self.q_proj(hidden_states)
k, _ = self.k_proj(hidden_states)
v, _ = self.v_proj(hidden_states)
with bi100_timer("full_attn.norm_rope"):
# q projection output includes gate (dim doubled).
qg = qg.view(total_tokens, self.local_num_heads,
self.head_dim * 2)
q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
q = self.q_norm.forward_cuda(
q.view(total_tokens, self.local_num_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
# Select the one rank-local KV head before k_norm and RoPE.
if self.q_per_kv_global is not None:
tp_rank = get_tensor_model_parallel_rank()
kv_idx = ((tp_rank * self.local_num_heads)
// self.q_per_kv_global)
k = (k.view(total_tokens, self.proj_kv_heads, self.head_dim)
[:, kv_idx, :].contiguous())
v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
[:, kv_idx, :].contiguous())
k = self.k_norm.forward_cuda(
k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
q, k = self.rotary_emb(positions, q, k)
with bi100_timer("full_attn.attention"):
with bi100_timer(f"L{self.layer_idx}.full_attn"):
attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
with bi100_timer("full_attn.gate"):
attn_out = (attn_out
* torch.sigmoid(gate.float()).to(attn_out.dtype))
with bi100_timer("full_attn.output_proj"):
output, _ = self.o_proj(attn_out)
return output
# ---------------------------------------------------------------------------
# MLP (SwiGLU, same as Qwen2/Qwen3)
# ---------------------------------------------------------------------------
class Qwen3_5MLP(nn.Module):
def __init__(
self,
hidden_size: int,
intermediate_size: int,
hidden_act: str,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size, [intermediate_size] * 2,
bias=False, quant_config=quant_config)
self.down_proj = RowParallelLinear(
intermediate_size, hidden_size,
bias=False, quant_config=quant_config)
if hidden_act != "silu":
raise ValueError(f"Unsupported activation: {hidden_act}")
self.act_fn = SiluAndMul()
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(x)
return x
# ---------------------------------------------------------------------------
# MoE sparse block (Qwen3.5-MoE / Qwen3.6-35B-A3B)
# ---------------------------------------------------------------------------
class Qwen3_5MoeSparseBlock(nn.Module):
"""Replaces Qwen3_5MLP for qwen3_5_moe_text layers.
FusedMoE is used ONLY for weight storage and loading (create_weights /
weight_loader are pure PyTorch). Its forward kernel is bypassed because
ixformer on BI-V100 lacks vllm_moe_topk_softmax / vllm_invoke_fused_moe_kernel.
Routing and expert computation use a pure-PyTorch loop instead.
Shared expert uses RowParallelLinear(reduce_results=False) so both paths
produce partial (pre-all-reduce) outputs that are combined before a single
all-reduce.
"""
def __init__(
self,
text_cfg,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
hidden_size = text_cfg.hidden_size
self.num_experts = text_cfg.num_experts
self.top_k = text_cfg.num_experts_per_tok
# Router and scalar shared-expert gate read the same hidden state. Keep
# their checkpoint shards in one replicated weight so forward needs a
# single GEMM for 256 + 1 outputs.
self.router_shared_gate = ReplicatedLinear(
hidden_size, text_cfg.num_experts + 1,
bias=False, quant_config=quant_config)
self.router_shared_gate.weight.weight_loader = \
self._router_shared_gate_weight_loader
# FusedMoE: only used for weight storage + weight_loader.
# Forward is bypassed — see _pure_pytorch_experts().
self.experts = FusedMoE(
num_experts=text_cfg.num_experts,
top_k=text_cfg.num_experts_per_tok,
hidden_size=hidden_size,
intermediate_size=text_cfg.moe_intermediate_size,
reduce_results=False, # we do the all-reduce ourselves below
renormalize=True,
quant_config=quant_config,
)
# Shared expert: defer all-reduce to combine with routed output first
shared_size = text_cfg.shared_expert_intermediate_size
self.shared_expert_gate_up = MergedColumnParallelLinear(
hidden_size, [shared_size] * 2, bias=False,
quant_config=quant_config)
self.shared_expert_down = RowParallelLinear(
shared_size, hidden_size, bias=False, reduce_results=False,
quant_config=quant_config)
self.act_fn = SiluAndMul()
def _router_shared_gate_weight_loader(
self,
param: torch.Tensor,
loaded_weight: torch.Tensor,
shard_id: int,
) -> None:
if shard_id == 0:
offset = 0
rows = self.num_experts
elif shard_id == 1:
offset = self.num_experts
rows = 1
else:
raise ValueError(f"unexpected router/shared gate shard: {shard_id}")
expected = (rows, param.shape[1])
if tuple(loaded_weight.shape) != expected:
raise ValueError(
"unexpected router/shared gate weight shape: "
f"expected {expected}, got {tuple(loaded_weight.shape)}")
param.data.narrow(0, offset, rows).copy_(loaded_weight)
def _pure_pytorch_experts(
self,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor:
"""MoE expert dispatch — fused C++ pipeline when available.
w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded]
w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded]
Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False.
"""
# ---------------------------------------------------------------
# Tier 0: Full fused MoE via ix_moe_bridge (xllm 7-step pipeline)
# topk → gen_idx → expand → group_gemm → silu → group_gemm → combine
# Source: xllm/core/layers/ilu/fused_moe.cpp
# ---------------------------------------------------------------
if _USE_IX_FUSED_MOE:
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
return _ix_fused_moe.fused_moe_forward(
hidden_states, router_logits,
w13, w2,
self.top_k, w13.shape[0],
True) # renormalize
# ---------------------------------------------------------------
# Tier 1: Point-optimized Python loop (individual corex .so)
# ---------------------------------------------------------------
# Fused topk+softmax: single CUB kernel vs 2 PyTorch ops.
# Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh
if _USE_COREX_MOE_TOPK_SOFTMAX:
topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax(
router_logits.float(), self.top_k, True)
topk_ids = topk_ids.to(torch.int64)
topk_weights = topk_weights.to(hidden_states.dtype)
else:
topk_logits, topk_ids = torch.topk(
router_logits.float(), self.top_k, dim=-1) # (T, top_k)
topk_weights = torch.softmax(topk_logits, dim=-1)
topk_weights = topk_weights.to(hidden_states.dtype)
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)
T = hidden_states.shape[0]
if T == 1:
# Fast path: single token (decode).
# Batched GEMM: replace top_k separate F.linear calls with 2 fused ops.
# gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I)
# down: 1 bmm (K,H,I) @ (K,I,1) → (K,H)
# Total: 3 kernel launches vs previous 16 (top_k*2).
eids = topk_ids[0] # (K,)
ws = topk_weights[0].to(hidden_states.dtype) # (K,)
use_corex_direct = (
_USE_COREX_MOE_DIRECT_ROUTED
and hidden_states.dtype == torch.float16
and w13.dtype == torch.float16
and w2.dtype == torch.float16
and ws.dtype == torch.float16
and hidden_states.is_cuda and w13.is_cuda and w2.is_cuda
and eids.is_cuda and ws.is_cuda
and hidden_states.is_contiguous()
and w13.is_contiguous() and w2.is_contiguous()
and eids.is_contiguous() and ws.is_contiguous()
and hidden_states.shape == (1, 2048)
and w13.shape == (256, 256, 2048)
and w2.shape == (256, 2048, 128)
and eids.shape == (8,) and ws.shape == (8,))
if use_corex_direct:
gate_up = _corex_moe_direct_routed.w13(
hidden_states, w13, eids)
act = self.act_fn(gate_up)
return _corex_moe_direct_routed.w2_reduce(
act, w2, eids, ws)
use_corex_gather = (
_USE_COREX_MOE_WEIGHT_GATHER
and hidden_states.dtype == torch.float16
and w13.dtype == torch.float16
and w2.dtype == torch.float16
and w13.is_cuda and w2.is_cuda and eids.is_cuda
and w13.is_contiguous() and w2.is_contiguous()
and eids.is_contiguous()
and w13.dim() == 3 and w2.dim() == 3
and eids.dim() == 1 and eids.numel() == 8
and w13.shape[0] == w2.shape[0]
and w13.shape[2] == w2.shape[1]
and w13.shape[1] == 2 * w2.shape[2]
and w13.shape[1] * w13.shape[2] % 8 == 0
and w2.shape[1] * w2.shape[2] % 8 == 0)
if use_corex_gather:
w13_sel, w2_sel = _corex_moe_weight_gather.gather(
w13, w2, eids)
else:
w13_sel = w13[eids] # (K, 2*I, H)
w2_sel = w2[eids] # (K, H, I)
H = hidden_states.shape[-1]
gate_up = F.linear(
hidden_states,
w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing
) # (1, K*2*I)
gate_up = gate_up.view(self.top_k, -1) # (K, 2*I)
if _USE_FUSED_MOE_ACTIVATION:
act = self.act_fn(gate_up) # (K, I)
else:
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up
# bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H)
expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
if (_USE_COREX_MOE_EXACT_REDUCE
and expert_out.dtype == torch.float16
and ws.dtype == torch.float16
and expert_out.shape[0] == 8):
out = _corex_moe_exact_reduce.serial_float(expert_out, ws)
else:
out = (expert_out * ws.unsqueeze(-1)).sum(
0, keepdim=True).to(hidden_states.dtype) # (1, H)
else:
# General path (prefill / multi-seq): group assignments once.
out = torch.zeros_like(hidden_states)
flat_eids = topk_ids.reshape(-1)
if _USE_COREX_MOE_INDEX_COMBINE:
# Fused CUDA: histogram + prefix_sum + place (11.5x faster)
src_dst, dst_src, expert_sizes = \
_corex_moe_index_combine.moe_compute_index(
flat_eids, w13.shape[0])
sorted_tok_ids = torch.arange(
T, device=topk_ids.device
).repeat_interleave(self.top_k)[dst_src.long()]
sorted_weights = topk_weights.reshape(-1)[dst_src.long()]
expert_counts = expert_sizes.tolist()
else:
order = torch.argsort(flat_eids, stable=True)
sorted_tok_ids = torch.arange(
T, device=topk_ids.device
).repeat_interleave(self.top_k)[order]
sorted_weights = topk_weights.reshape(-1)[order]
expert_counts = torch.bincount(
flat_eids, minlength=w13.shape[0]).tolist()
start = 0
for eid, count in enumerate(expert_counts):
end = start + count
if count == 0:
start = end
continue
tok_ids = sorted_tok_ids[start:end]
tokens = hidden_states[tok_ids] # (n, H)
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up # (n, I)
expert_out = F.linear(act, w2[eid]) # (n, H)
weights = sorted_weights[start:end].unsqueeze(-1)
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
start = end
return out # partial, all-reduce done in forward()
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
with bi100_timer("moe.router"):
router_and_shared_gate, _ = self.router_shared_gate(hidden_states)
router_logits = router_and_shared_gate[..., :self.num_experts]
gate_score = router_and_shared_gate[..., self.num_experts:]
with bi100_timer("moe.routed"):
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
with bi100_timer("moe.shared"):
gate_up, _ = self.shared_expert_gate_up(hidden_states)
shared_out = self.act_fn(gate_up)
shared_out, _ = self.shared_expert_down(shared_out)
shared_out = shared_out * torch.sigmoid(gate_score)
with bi100_timer("moe.combine"):
out = routed_out + shared_out
if self.experts.tp_size > 1:
with bi100_timer("moe.all_reduce"):
out = tensor_model_parallel_all_reduce(out)
return out
# ---------------------------------------------------------------------------
# Decoder layer (dispatches to GatedDeltaNet or Qwen3_5FullAttention)
# ---------------------------------------------------------------------------
class Qwen3_5DecoderLayer(nn.Module):
def __init__(
self,
text_cfg,
layer_idx: int,
layer_type: str,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.layer_idx = layer_idx
self.layer_type = layer_type
self._diagnostic_trace_pending = (
os.getenv("BI100_DIAGNOSTIC_LAYER_TRACE") == "1")
self.input_layernorm = GemmaRMSNorm(text_cfg.hidden_size,
eps=text_cfg.rms_norm_eps)
self.post_attention_layernorm = GemmaRMSNorm(text_cfg.hidden_size,
eps=text_cfg.rms_norm_eps)
if layer_type == "linear_attention":
self.linear_attn = GatedDeltaNet(text_cfg, layer_idx,
quant_config=quant_config)
else:
self.self_attn = Qwen3_5FullAttention(
text_cfg, layer_idx,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"layers.{layer_idx}.self_attn",
)
if getattr(text_cfg, 'model_type', '') == 'qwen3_5_moe_text':
self.mlp = Qwen3_5MoeSparseBlock(text_cfg, quant_config=quant_config)
else:
self.mlp = Qwen3_5MLP(
hidden_size=text_cfg.hidden_size,
intermediate_size=text_cfg.intermediate_size,
hidden_act=text_cfg.hidden_act,
quant_config=quant_config,
)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: Optional[torch.Tensor],
attn_metadata: AttentionMetadata,
residual: Optional[torch.Tensor],
# Only for linear_attention layers:
conv_state: Optional[torch.Tensor] = None,
temporal_state: Optional[torch.Tensor] = None,
gdn_capture_offsets: Optional[Iterable[int]] = None,
gdn_segment_offsets: Optional[Iterable[int]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
with bi100_timer("layer.input_norm"):
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(
hidden_states, residual)
if self.layer_type == "linear_attention":
with bi100_timer("layer.gdn"):
hidden_states = self.linear_attn(
hidden_states, attn_metadata, conv_state, temporal_state,
capture_offsets=gdn_capture_offsets,
segment_offsets=gdn_segment_offsets)
else:
with bi100_timer("layer.full_attn"):
hidden_states = self.self_attn(
positions, hidden_states, kv_cache, attn_metadata)
with bi100_timer("layer.post_attn_norm"):
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual)
with bi100_timer("layer.moe"):
hidden_states = self.mlp(hidden_states)
if self._diagnostic_trace_pending:
self._diagnostic_trace_pending = False
rank = os.getenv("RANK", os.getenv("LOCAL_RANK", "?"))
print(
"[BI100 DIAGNOSTIC] "
f"rank={rank} layer={self.layer_idx} "
f"attention={self.layer_type} "
f"mlp={type(self.mlp).__name__} stage=completed",
file=sys.stderr,
flush=True,
)
return hidden_states, residual
# ---------------------------------------------------------------------------
# Full transformer model
# ---------------------------------------------------------------------------
def _validate_qwen_kv_cache_count(configured_count, kv_caches):
if len(kv_caches) != configured_count:
raise RuntimeError(
"Qwen3.5 allocated KV cache count mismatch: "
f"configured {configured_count}, received {len(kv_caches)}")
class Qwen3_5Model(nn.Module):
def __init__(
self,
text_cfg,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
kv_cache_count: Optional[int] = None,
) -> None:
super().__init__()
self.text_cfg = text_cfg
full_attention_count = sum(
layer_type == "full_attention"
for layer_type in text_cfg.layer_types)
if kv_cache_count is None:
kv_cache_count = full_attention_count
if (not isinstance(kv_cache_count, int) or isinstance(kv_cache_count, bool)
or kv_cache_count < full_attention_count):
raise RuntimeError(
"Qwen3.5 configured KV cache count must cover every "
f"full-attention layer: configured {kv_cache_count}, "
f"required {full_attention_count}")
self.kv_cache_count = kv_cache_count
self.embed_tokens = VocabParallelEmbedding(
text_cfg.vocab_size, text_cfg.hidden_size)
self.layers = nn.ModuleList([
Qwen3_5DecoderLayer(
text_cfg, i, text_cfg.layer_types[i],
cache_config=cache_config, quant_config=quant_config)
for i in range(text_cfg.num_hidden_layers)
])
self.norm = GemmaRMSNorm(text_cfg.hidden_size, eps=text_cfg.rms_norm_eps)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[torch.Tensor],
attn_metadata: AttentionMetadata,
conv_states: torch.Tensor, # (num_linear_layers, batch, ...)
temporal_states: torch.Tensor, # (num_linear_layers, batch, ...)
inputs_embeds: Optional[torch.Tensor] = None,
gdn_capture_offsets: Optional[Iterable[int]] = None,
gdn_segment_offsets: Optional[Iterable[int]] = None,
) -> torch.Tensor:
_validate_qwen_kv_cache_count(self.kv_cache_count, kv_caches)
with bi100_timer("model.embed"):
hidden_states = (self.embed_tokens(input_ids)
if inputs_embeds is None else inputs_embeds)
residual = None
attn_idx = 0
linear_idx = 0
capture_offsets = tuple(gdn_capture_offsets or ())
captured_conv_states: Dict[int, List[torch.Tensor]] = {
offset: [] for offset in capture_offsets
}
captured_temporal_states: Dict[int, List[torch.Tensor]] = {
offset: [] for offset in capture_offsets
}
for layer in self.layers:
if layer.layer_type == "linear_attention":
hidden_states, residual = layer(
positions, hidden_states,
kv_cache=None,
attn_metadata=attn_metadata,
residual=residual,
conv_state=conv_states[linear_idx],
temporal_state=temporal_states[linear_idx],
gdn_capture_offsets=capture_offsets,
gdn_segment_offsets=gdn_segment_offsets,
)
for offset in capture_offsets:
captured_conv_states[offset].append(
layer.linear_attn.captured_conv_states[offset])
captured_temporal_states[offset].append(
layer.linear_attn.captured_temporal_states[offset])
linear_idx += 1
else:
kv_cache = kv_caches[attn_idx]
hidden_states, residual = layer(
positions, hidden_states,
kv_cache=kv_cache,
attn_metadata=attn_metadata,
residual=residual,
)
attn_idx += 1
with bi100_timer("model.final_norm"):
hidden_states, _ = self.norm(hidden_states, residual)
self.captured_conv_states = {
offset: torch.stack(states)
for offset, states in captured_conv_states.items()
}
self.captured_temporal_states = {
offset: torch.stack(states)
for offset, states in captured_temporal_states.items()
}
return hidden_states
# ---------------------------------------------------------------------------
# Top-level CausalLM wrapper with MambaCacheManager
# ---------------------------------------------------------------------------
class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
SupportsMultiModal):
has_inner_state = True
supports_lora = True
packed_modules_mapping = {
"gate_up_proj": ["gate_proj", "up_proj"],
}
supported_lora_modules = [
"gate_up_proj",
"down_proj",
"o_proj",
]
embedding_modules = {}
embedding_padding_modules = []
def __init__(
self,
config, # Qwen3_5Config (top-level)
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
lora_config: Optional[LoRAConfig] = None,
scheduler_config: Optional[SchedulerConfig] = None,
multimodal_config: Optional[MultiModalConfig] = None,
prefix: str = "",
) -> None:
_bi100_model_trace("Qwen3_5ForCausalLM initialization begin")
super().__init__()
self.config = config
self.scheduler_config = scheduler_config
self.multimodal_config = multimodal_config
# The text config holds all architecture parameters
text_cfg = config.text_config
self.text_cfg = text_cfg
rope_parameters = getattr(text_cfg, "rope_parameters", {}) or {}
mrope_sections = rope_parameters.get("mrope_section", [11, 11, 10])
if getattr(config, "rope_scaling", None) is None:
config.rope_scaling = {
"type": "mrope",
"mrope_section": mrope_sections,
}
# Pre-compute counts
self.num_linear_layers = sum(
1 for lt in text_cfg.layer_types if lt == "linear_attention")
self.num_attn_layers = sum(
1 for lt in text_cfg.layer_types if lt == "full_attention")
layers_block_type = getattr(
config, "layers_block_type",
["attention"] * text_cfg.num_hidden_layers)
self.num_kv_cache_layers = sum(
layer_type == "attention" for layer_type in layers_block_type)
if self.num_kv_cache_layers < self.num_attn_layers:
raise RuntimeError(
"Qwen3.5 KV accounting provides fewer caches than "
f"full-attention layers: {self.num_kv_cache_layers} < "
f"{self.num_attn_layers}")
accounting_mode = getattr(
config, "bi100_hybrid_kv_accounting_mode", "legacy40")
accounting_env = os.getenv("BI100_HYBRID_KV_ACCOUNTING", "<unset>")
tp_rank = get_tensor_model_parallel_rank()
full_attention_ordinals = ",".join(
str(index) for index, layer_type in enumerate(text_cfg.layer_types)
if layer_type == "full_attention")
logger.info(
"[BI100] Qwen hybrid KV accounting; tp_rank=%d "
"env_mode=%s config_mode=%s "
"configured_kv_layers=%d full_attention_layers=%d "
"full_attention_ordinals=%s",
tp_rank,
accounting_env,
accounting_mode,
self.num_kv_cache_layers,
self.num_attn_layers,
full_attention_ordinals,
)
# DeltaNet state dimensions (per layer, per sequence, TP-sharded)
tp_size = get_tensor_model_parallel_world_size()
self.conv_dim = (text_cfg.linear_num_key_heads * text_cfg.linear_key_head_dim * 2
+ text_cfg.linear_num_value_heads * text_cfg.linear_value_head_dim)
self.num_v_heads = text_cfg.linear_num_value_heads
self.head_k_dim = text_cfg.linear_key_head_dim
self.head_v_dim = text_cfg.linear_value_head_dim
self.conv_kernel_size = text_cfg.linear_conv_kernel_dim
self.model = Qwen3_5Model(
text_cfg,
cache_config=cache_config,
quant_config=quant_config,
kv_cache_count=self.num_kv_cache_layers,
)
self.visual = Qwen3_5VisionTransformer(
config.vision_config,
quant_config=None,
)
self.lm_head = ParallelLMHead(
text_cfg.vocab_size, text_cfg.hidden_size,
quant_config=quant_config,
)
self.logits_processor = LogitsProcessor(text_cfg.vocab_size)
self.sampler = Sampler()
# Lazy initialised in first forward call
self.mamba_cache: Optional[MambaCacheManager] = None
# Scheduler-owned recurrent prefix states. Keys are stable chained
# content hashes, never recyclable physical KV block ids.
self._gdn_prefix_cache: Dict[
Tuple[int, bytes], Tuple[torch.Tensor, torch.Tensor]] = {}
self._block_size: int = (cache_config.block_size
if cache_config is not None else 16)
self._startup_forward_traced = False
_bi100_model_trace("Qwen3_5ForCausalLM initialization complete")
def _get_mamba_cache_shape(self):
tp_size = get_tensor_model_parallel_world_size()
# Each sequence's state is stored in float32
conv_state_shape = (self.conv_dim // tp_size, self.conv_kernel_size - 1)
temporal_state_shape = (
self.num_v_heads // tp_size, self.head_k_dim, self.head_v_dim)
return conv_state_shape, temporal_state_shape
@staticmethod
def _validate_and_reshape_mm_tensor(
mm_input: Union[torch.Tensor, List[torch.Tensor]],
name: str,
) -> torch.Tensor:
if isinstance(mm_input, list):
return torch.cat(mm_input)
if not isinstance(mm_input, torch.Tensor):
raise ValueError(f"incorrect type for {name}: {type(mm_input)}")
if mm_input.ndim == 2:
return mm_input
if mm_input.ndim == 3:
return torch.cat(list(mm_input))
raise ValueError(
f"{name} must be a 2D tensor or batched 3D tensor, got "
f"shape={tuple(mm_input.shape)}")
def _parse_and_validate_image_input(
self,
**kwargs: object,
) -> Optional[Qwen3_5ImageInputs]:
pixel_values = kwargs.get("pixel_values")
image_embeds = kwargs.get("image_embeds")
image_grid_thw = kwargs.get("image_grid_thw")
if pixel_values is None and image_embeds is None:
return None
if pixel_values is not None:
if image_grid_thw is None:
raise ValueError("image_grid_thw is required with pixel_values")
return Qwen3_5ImagePixelInputs(
type="pixel_values",
data=self._validate_and_reshape_mm_tensor(
pixel_values, "image pixel values"),
image_grid_thw=self._validate_and_reshape_mm_tensor(
image_grid_thw, "image grid_thw"),
)
return Qwen3_5ImageEmbeddingInputs(
type="image_embeds",
data=self._validate_and_reshape_mm_tensor(
image_embeds, "image embeddings"),
)
def _process_image_input(
self,
image_input: Qwen3_5ImageInputs,
) -> torch.Tensor:
if image_input["type"] == "image_embeds":
return image_input["data"].to(dtype=self.visual.dtype,
device=self.visual.device)
return self.visual(
image_input["data"],
grid_thw=image_input["image_grid_thw"],
)
@bi100_profile_transaction
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[torch.Tensor],
attn_metadata: AttentionMetadata,
intermediate_tensors: Optional[IntermediateTensors] = None,
**kwargs,
) -> torch.Tensor:
if not self._startup_forward_traced:
self._startup_forward_traced = True
_bi100_model_trace("first model forward entered")
if self.mamba_cache is None:
if self.scheduler_config is not None:
max_batch_size = _get_graph_batch_size(
self.scheduler_config.max_num_seqs)
else:
max_batch_size = max(_BATCH_SIZES_TO_CAPTURE) + 2
self.mamba_cache = MambaCacheManager(
torch.float32,
self.num_linear_layers,
max_batch_size,
*self._get_mamba_cache_shape(),
)
gdn_restore_key = kwargs.pop("gdn_restore_key", None)
gdn_capture_points = kwargs.pop("gdn_capture_points", None) or []
gdn_evict_keys = kwargs.pop("gdn_evict_keys", None) or []
gdn_segment_offsets = kwargs.pop("gdn_segment_offsets", None) or []
mamba_tensors = self.mamba_cache.current_run_tensors(
input_ids, attn_metadata, **kwargs)
# conv_states: (num_linear_layers, batch, local_conv_dim, kernel-1)
# temporal_states: (num_linear_layers, batch, local_num_v, k_dim, v_dim)
conv_states, temporal_states = mamba_tensors
_is_single_seq_prefill = (
attn_metadata is not None
and attn_metadata.num_prefill_tokens > 0
and conv_states.shape[1] == 1 # batch == 1
and getattr(attn_metadata, 'context_lens_tensor', None) is not None
)
has_gdn_actions = (gdn_restore_key is not None
or bool(gdn_capture_points)
or bool(gdn_evict_keys)
or bool(gdn_segment_offsets))
if has_gdn_actions and not _is_single_seq_prefill:
raise RuntimeError(
"GDN prefix-cache actions require a single-sequence prefill")
for evict_key in gdn_evict_keys:
self._gdn_prefix_cache.pop(_validate_gdn_prefix_key(evict_key),
None)
if gdn_restore_key is not None:
restore_key = _validate_gdn_prefix_key(gdn_restore_key)
saved_state = self._gdn_prefix_cache.get(restore_key)
if saved_state is None:
raise RuntimeError(
"scheduler requested a missing GDN prefix state: "
f"blocks={restore_key[0]} digest={restore_key[1].hex()}")
saved_conv, saved_temporal = saved_state
with bi100_timer("gdn_prefix.restore"):
conv_states[:, 0].copy_(
saved_conv.to(device=conv_states.device,
dtype=conv_states.dtype),
non_blocking=True)
temporal_states[:, 0].copy_(
saved_temporal.to(device=temporal_states.device,
dtype=temporal_states.dtype),
non_blocking=True)
query_len = (int(attn_metadata.num_prefill_tokens)
if _is_single_seq_prefill else 0)
capture_keys: Dict[int, Tuple[int, bytes]] = {}
for capture_point in gdn_capture_points:
if not isinstance(capture_point, tuple) or len(capture_point) != 2:
raise RuntimeError(
f"invalid GDN capture point: {capture_point!r}")
offset, capture_key = capture_point
if (not isinstance(offset, int) or offset <= 0
or offset > query_len or offset in capture_keys):
raise RuntimeError(
f"invalid GDN capture offset: {offset!r} "
f"for query_len={query_len}")
capture_keys[offset] = _validate_gdn_prefix_key(capture_key)
if len(capture_keys) > 2:
raise RuntimeError("at most two GDN capture points are supported")
interior_capture_offsets = tuple(
offset for offset in capture_keys if offset < query_len)
segment_offsets = set()
for offset in gdn_segment_offsets:
if (not isinstance(offset, int) or offset <= 0
or offset >= query_len):
raise RuntimeError(
f"invalid GDN segment offset: {offset!r} "
f"for query_len={query_len}")
segment_offsets.add(offset)
if len(segment_offsets) > 128:
raise RuntimeError("at most 128 GDN segment offsets are supported")
interior_segment_offsets = tuple(sorted(segment_offsets))
inputs_embeds = None
image_input = self._parse_and_validate_image_input(**kwargs)
if image_input is not None:
image_mask = input_ids == self.config.image_token_id
num_placeholders = int(image_mask.sum().item())
if num_placeholders:
inputs_embeds = self.model.embed_tokens(input_ids)
image_embeds = self._process_image_input(image_input)
if num_placeholders > image_embeds.shape[0]:
raise ValueError(
f"image token count ({num_placeholders}) exceeds "
f"vision embeddings ({image_embeds.shape[0]})")
# Prefix caching can consume the leading image tokens while
# vLLM 0.6 still supplies the full pixel tensor. The query's
# remaining placeholders always form a suffix of the flattened
# visual token stream.
image_embeds = image_embeds[-num_placeholders:]
inputs_embeds[image_mask, :] = image_embeds.to(
inputs_embeds.dtype)
with bi100_timer("model.forward"):
hidden_states = self.model(
input_ids, positions, kv_caches, attn_metadata,
conv_states, temporal_states,
inputs_embeds=inputs_embeds,
gdn_capture_offsets=interior_capture_offsets,
gdn_segment_offsets=interior_segment_offsets)
for offset, capture_key in capture_keys.items():
if offset == query_len:
captured_conv = conv_states[:, 0]
captured_temporal = temporal_states[:, 0]
else:
captured_conv = self.model.captured_conv_states[offset]
captured_temporal = self.model.captured_temporal_states[offset]
with bi100_timer("gdn_prefix.save"):
self._gdn_prefix_cache[capture_key] = (
captured_conv.detach().cpu().clone(),
captured_temporal.detach().cpu().clone(),
)
if bi100_profile_event_enabled():
profile_prefill_tokens = int(
getattr(attn_metadata, "num_prefill_tokens", 0) or 0)
profile_decode_tokens = int(
getattr(attn_metadata, "num_decode_tokens", 0) or 0)
profile_context_len = 0
if profile_prefill_tokens > 0:
profile_seq_lens = getattr(attn_metadata, "seq_lens", None)
if (not isinstance(profile_seq_lens, list)
or len(profile_seq_lens) != 1
or not isinstance(profile_seq_lens[0], int)):
raise RuntimeError(
"BI100 profile requires one host-visible prefill "
"sequence length")
profile_context_len = (
profile_seq_lens[0] - profile_prefill_tokens)
if profile_context_len < 0:
raise RuntimeError(
"BI100 profile observed a negative prefill context")
bi100_profile_flush(
tp_rank=get_tensor_model_parallel_rank(),
phase=("prefill" if profile_prefill_tokens > 0 else "decode"),
prefill_tokens=profile_prefill_tokens,
decode_tokens=profile_decode_tokens,
context_len=profile_context_len,
gdn_restore=bool(gdn_restore_key is not None),
gdn_capture_points=len(gdn_capture_points),
gdn_evict_keys=len(gdn_evict_keys),
)
return hidden_states
def compute_logits(
self,
hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
# All TP ranks must call logits_processor to participate in the NCCL
# gather inside lm_head. Non-driver ranks return None after the gather.
# With chunked prefill, intermediate chunks have seq_groups=None on all
# ranks; _apply_logits_processors is guarded against this in
# logits_processor.py (patched by patch_xformers_sdpa_seq.py).
logits = self.logits_processor(self.lm_head, hidden_states,
sampling_metadata)
return logits
def sample(
self,
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[SamplerOutput]:
return self.sampler(logits, sampling_metadata)
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
return self.mamba_cache.copy_inputs_before_cuda_graphs(
input_buffers, **kwargs)
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
_bi100_model_trace("dense load_weights begin")
loaded_count = 0
stacked_params_mapping = [
# (param_name, weight_name, shard_id)
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
params_dict = dict(self.named_parameters())
for name, loaded_weight in weights:
loaded_count += 1
# Skip vision and MTP branches
if (name.startswith("model.visual")
or name.startswith("mtp.")
or name.startswith("model.mtp")):
continue
# Prefix remapping: checkpoint may wrap under language_model
if name.startswith("model.language_model."):
name = "model." + name[len("model.language_model."):]
# Skip positional embedding caches
if "rotary_emb.inv_freq" in name:
continue
if _load_full_attention_qgkv_weight(
params_dict, name, loaded_weight, self.text_cfg):
continue
if _load_gdn_projection_weight(
params_dict, name, loaded_weight, self.text_cfg):
continue
# Remap conv1d.weight → conv1d_weight
# The conv has depth (1) dim in the checkpoint that we handle separately
if ".linear_attn.conv1d.weight" in name:
name = name.replace(".linear_attn.conv1d.weight",
".linear_attn.conv1d_weight")
# Stacked param loading (gate_up_proj)
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name.endswith(".bias") and name not in params_dict:
break
if name not in params_dict:
break
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
if name.endswith(".bias") and name not in params_dict:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader",
default_weight_loader)
weight_loader(param, loaded_weight)
_bi100_model_trace(f"dense load_weights complete items={loaded_count}")
# ---------------------------------------------------------------------------
# Qwen3.6-35B-A3B (Qwen3_5-MoE architecture)
# ---------------------------------------------------------------------------
@MULTIMODAL_REGISTRY.register_image_input_mapper(qwen36_image_input_mapper)
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_qwen36_image_tokens)
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_qwen36)
@INPUT_REGISTRY.register_input_processor(input_processor_for_qwen36)
class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
"""Qwen3.6-35B-A3B: same hybrid-attention backbone as 27B, dense MLP
replaced by Qwen3_5MoeSparseBlock (256 routed experts + shared expert).
Only load_weights differs from the dense variant.
"""
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
_bi100_model_trace("MoE load_weights begin")
loaded_count = 0
vision_loaded_count = 0
# Checkpoint key format for this model (transformers Qwen3_5MoeExperts):
# mlp.experts.gate_up_proj shape (num_experts, 2*intermediate, hidden)
# mlp.experts.down_proj shape (num_experts, hidden, intermediate)
# mlp.gate.weight shape (num_experts, hidden) [router]
# mlp.shared_expert_gate.weight shape (1, hidden)
# mlp.shared_expert.{gate,up,down}_proj.weight [shared MLP]
# Our FusedMoE stores:
# mlp.experts.w13_weight shape (num_experts, 2*intermediate//tp, hidden)
# mlp.experts.w2_weight shape (num_experts, hidden, intermediate//tp)
# Our router/shared gate stores both tensors in one (num_experts+1, H)
# replicated weight. Our shared expert stores:
# mlp.shared_expert_gate_up.weight (merged gate+up)
# mlp.shared_expert_down.weight
stacked_params_mapping = [
# (param_name, weight_name, shard_id)
# shared expert
("shared_expert_gate_up", "shared_expert.gate_proj", 0),
("shared_expert_gate_up", "shared_expert.up_proj", 1),
# linear_attention dense proj (same as 27B)
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
params_dict = dict(self.named_parameters())
for name, loaded_weight in weights:
loaded_count += 1
if name.startswith("model.visual."):
name = "visual." + name[len("model.visual."):]
if "attn.qkv.weight" in name:
num_heads = self.config.vision_config.num_heads
hidden_size = self.config.vision_config.hidden_size
head_size = hidden_size // num_heads
loaded_weight = loaded_weight.view(
3, num_heads, head_size, hidden_size)
loaded_weight = loaded_weight.transpose(0, 1).reshape(
-1, hidden_size)
elif "attn.qkv.bias" in name:
num_heads = self.config.vision_config.num_heads
hidden_size = self.config.vision_config.hidden_size
head_size = hidden_size // num_heads
loaded_weight = loaded_weight.view(
3, num_heads, head_size)
loaded_weight = loaded_weight.transpose(0, 1).reshape(-1)
if name not in params_dict:
raise ValueError(f"unexpected Qwen3.6 vision weight: {name}")
param = params_dict[name]
weight_loader = getattr(param, "weight_loader",
default_weight_loader)
weight_loader(param, loaded_weight)
vision_loaded_count += 1
continue
# MTP is not used by the fixed evaluator command.
if (name.startswith("mtp.")
or name.startswith("model.mtp")):
continue
# Prefix remapping for VL checkpoint (Qwen3_5MoeForConditionalGeneration):
# model.language_model.model.{layers,embed_tokens,norm} -> model.{...}
# model.language_model.lm_head -> lm_head
# Prefix remapping: checkpoint may wrap under language_model
if name.startswith("model.language_model."):
name = "model." + name[len("model.language_model."):]
if "rotary_emb.inv_freq" in name:
continue
if _load_full_attention_qgkv_weight(
params_dict, name, loaded_weight, self.text_cfg):
continue
if _load_gdn_projection_weight(
params_dict, name, loaded_weight, self.text_cfg):
continue
if name.endswith(".mlp.gate.weight"):
fused_name = name[:-len("gate.weight")] \
+ "router_shared_gate.weight"
if fused_name not in params_dict:
raise ValueError(
f"missing fused router/shared gate: {fused_name}")
params_dict[fused_name].weight_loader(
params_dict[fused_name], loaded_weight, 0)
continue
if name.endswith(".mlp.shared_expert_gate.weight"):
fused_name = name[:-len("shared_expert_gate.weight")] \
+ "router_shared_gate.weight"
if fused_name not in params_dict:
raise ValueError(
f"missing fused router/shared gate: {fused_name}")
params_dict[fused_name].weight_loader(
params_dict[fused_name], loaded_weight, 1)
continue
if ".linear_attn.conv1d.weight" in name:
name = name.replace(".linear_attn.conv1d.weight",
".linear_attn.conv1d_weight")
# --- Fused routed-expert weights (all experts in one tensor) ---
if "mlp.experts.gate_up_proj" in name:
# loaded_weight: (num_experts, 2*intermediate, hidden)
w13_name = name.replace("mlp.experts.gate_up_proj",
"mlp.experts.w13_weight")
if w13_name not in params_dict:
continue
param = params_dict[w13_name]
n_exp = loaded_weight.shape[0]
inter = loaded_weight.shape[1] // 2
gate_w = loaded_weight[:, :inter, :].contiguous()
up_w = loaded_weight[:, inter:, :].contiguous()
for eid in range(n_exp):
param.weight_loader(param, gate_w[eid], "w1_weight", "w1", eid)
param.weight_loader(param, up_w[eid], "w3_weight", "w3", eid)
continue
if "mlp.experts.down_proj" in name:
# loaded_weight: (num_experts, hidden, intermediate)
w2_name = name.replace("mlp.experts.down_proj",
"mlp.experts.w2_weight")
if w2_name not in params_dict:
continue
param = params_dict[w2_name]
n_exp = loaded_weight.shape[0]
for eid in range(n_exp):
param.weight_loader(param, loaded_weight[eid], "w2_weight", "w2", eid)
continue
# --- Shared expert down_proj rename ---
if "mlp.shared_expert.down_proj" in name:
name = name.replace("mlp.shared_expert.down_proj",
"mlp.shared_expert_down")
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
continue
# --- Individual expert weights (FT checkpoint: experts.{i}.{proj}.weight) ---
# Standard transformers fine-tuning saves each expert separately instead of
# the pre-merged (num_experts, ...) tensors in the original checkpoint.
if ".mlp.experts." in name:
parts = name.split(".mlp.experts.", 1)
expert_rest = parts[1] # e.g. "0.gate_proj.weight"
dot_pos = expert_rest.find(".")
if dot_pos > 0 and expert_rest[:dot_pos].isdigit():
eid = int(expert_rest[:dot_pos])
proj_raw = expert_rest[dot_pos + 1:]
proj = proj_raw[:-7] if proj_raw.endswith(".weight") else proj_raw
prefix = parts[0] # e.g. "model.layers.0"
if proj == "gate_proj":
w13_name = f"{prefix}.mlp.experts.w13_weight"
if w13_name in params_dict:
param = params_dict[w13_name]
param.weight_loader(param, loaded_weight, "w1_weight", "w1", eid)
elif proj == "up_proj":
w13_name = f"{prefix}.mlp.experts.w13_weight"
if w13_name in params_dict:
param = params_dict[w13_name]
param.weight_loader(param, loaded_weight, "w3_weight", "w3", eid)
elif proj == "down_proj":
w2_name = f"{prefix}.mlp.experts.w2_weight"
if w2_name in params_dict:
param = params_dict[w2_name]
param.weight_loader(param, loaded_weight, "w2_weight", "w2", eid)
continue
# --- Stacked / standard weights ---
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
break
param = params_dict[name]
param.weight_loader(param, loaded_weight, shard_id)
break
else:
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
_bi100_model_trace(
f"MoE load_weights complete items={loaded_count} "
f"vision_items={vision_loaded_count}")
Reference in New Issue View Git Blame Copy Permalink