A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 13:45 UTC
[Bugfix ] fix bailing_moe_linear (#40859)
Signed-off-by: ghphotoframe <854746559@qq.com>
JiangWeixiang committed
Apr 28, 2026 at 13:39 UTC
ed57f771923703998a17ad656536ffb460447a2c
2 files changed
+15
-16
vllm/model_executor/layers/mamba/mamba_utils.py
-3
@@ -55,9 +55,6 @@ class MambaStateDtypeCalculator:
55
model_dtype: ModelDType | torch.dtype,
56
mamba_cache_dtype: MambaDType,
57
) -> tuple[torch.dtype, ...]:
58
- # TODO (tdoublep) requires testing
59
- if mamba_cache_dtype == "float32":
60
- raise ValueError("fp32 state for minimax is not yet supported")
58
state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
59
return (state_dtype,)
60
vllm/model_executor/models/bailing_moe_linear.py
+15
-13
@@ -17,6 +17,7 @@ from vllm.distributed import (
17
)
18
from vllm.forward_context import get_forward_context
19
from vllm.logger import init_logger
20
+from vllm.model_executor.custom_op import PluggableLayer
21
from vllm.model_executor.layers.fla.ops.layernorm_guard import (
22
RMSNormGated,
23
layernorm_fn,
@@ -211,7 +212,6 @@ class BailingMoeV25MLAAttention(nn.Module):
212
max_position=max_position,
213
is_neox_style=False,
214
rope_parameters=rope_parameters or None,
214
- dtype=torch.float32,
215
)
216
217
# Build MLAModules for MultiHeadLatentAttentionWrapper
@@ -425,14 +425,18 @@ class BailingGroupRMSNormGate(RMSNormGated):
425
param.data.copy_(loaded_weight[shard].contiguous())
426
427
428
-class BailingMoELinearAttention(nn.Module, MambaBase):
429
- """
430
- Bailing MoE Linear Attention implementation using minimax backend.
428
+# --8<-- [start:bailing_moe_linear_attention]
429
+@PluggableLayer.register("bailing_moe_linear_attention")
430
+class BailingMoELinearAttention(PluggableLayer, MambaBase):
431
+ """Pluggable Bailing MoE Linear Attention layer which allows OOT backends
432
+ to add custom implementations.
433
432
- This implements the linear attention mechanism from sglang, adapted for vLLM's
433
- v1 engine with MambaBase interface support.
434
+ This implements the linear attention mechanism from sglang, adapted for
435
+ vLLM's v1 engine with MambaBase interface support.
436
"""
437
438
+ # --8<-- [end:bailing_moe_linear_attention]
439
+
440
@property
441
def mamba_type(self) -> str:
442
return "linear_attention"
@@ -569,7 +573,6 @@ class BailingMoELinearAttention(nn.Module, MambaBase):
573
self.head_dim,
574
max_position=self.max_position_embeddings,
575
is_neox_style=True,
572
- dtype=torch.float32,
576
rope_parameters=rope_parameters or None,
577
)
578
@@ -754,8 +757,6 @@ class BailingMoELinearAttention(nn.Module, MambaBase):
757
758
def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):
759
"""Handle decode (single token per sequence)."""
757
- num_prefill_tokens = attn_metadata.num_prefill_tokens
758
- num_prefills = attn_metadata.num_prefills
760
hidden = linear_attention_decode(
761
q,
762
k,
@@ -763,10 +764,10 @@ class BailingMoELinearAttention(nn.Module, MambaBase):
764
kv_cache,
765
self.tp_slope,
766
state_indices_tensor,
766
- q_start=num_prefill_tokens,
767
- q_end=None,
768
- slot_start=num_prefills,
769
- slot_end=None,
767
+ q_start=0,
768
+ q_end=attn_metadata.num_decode_tokens,
769
+ slot_start=0,
770
+ slot_end=attn_metadata.num_decodes,
771
block_size=32,
772
)
773
return hidden
@@ -1149,6 +1150,7 @@ class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
1150
config.vocab_size,
1151
config.hidden_size,
1152
quant_config=quant_config,
1153
+ prefix=maybe_prefix(prefix, "lm_head"),
1154
)
1155
self.logits_processor = LogitsProcessor(config.vocab_size)
1156
else: