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 8, 16:30 UTC
[Model] Enable Step3p5ForCausalLM testing (#33755)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Feb 7, 2026 at 21:25 UTC
db4ede97434300e2dea211f7234024dd7fe34bfc
3 files changed
+28
-32
docs/models/supported_models.md
+1
-1
@@ -471,7 +471,7 @@ th {
471
| `StableLMEpochForCausalLM` | StableLM Epoch | `stabilityai/stablelm-zephyr-3b`, etc. | | ✅︎ |
472
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
473
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
474
-| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/step-3.5-flash`, etc. | | ✅︎ |
474
+| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ |
475
| `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
476
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
477
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+15
-6
@@ -481,16 +481,21 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
481
"ByteDance-Seed/Seed-OSS-36B-Instruct",
482
trust_remote_code=True,
483
),
484
+ "SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
485
+ "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),
486
+ "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
487
+ "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),
488
"Step1ForCausalLM": _HfExamplesInfo(
489
"stepfun-ai/Step-Audio-EditX", trust_remote_code=True
490
),
491
"Step3p5ForCausalLM": _HfExamplesInfo(
488
- "stepfun-ai/step-3.5-flash", is_available_online=False
492
+ "stepfun-ai/Step-3.5-Flash",
493
+ use_original_num_layers=True,
494
+ # Initialize at least one MoE layer
495
+ hf_overrides={
496
+ "num_hidden_layers": 4,
497
+ },
498
),
490
- "SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
491
- "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),
492
- "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
493
- "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),
499
"Step3TextForCausalLM": _HfExamplesInfo("stepfun-ai/step3", trust_remote_code=True),
500
"SolarForCausalLM": _HfExamplesInfo(
501
"upstage/solar-pro-preview-instruct", trust_remote_code=True
@@ -1129,8 +1134,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1134
),
1135
"Step3p5MTP": _HfExamplesInfo(
1136
"stepfun-ai/Step-3.5-Flash",
1132
- trust_remote_code=True,
1137
speculative_model="stepfun-ai/Step-3.5-Flash",
1138
+ use_original_num_layers=True,
1139
+ # Initialize at least one MoE layer
1140
+ hf_overrides={
1141
+ "num_hidden_layers": 4,
1142
+ },
1143
is_available_online=False,
1144
),
1145
}
vllm/model_executor/models/step3p5.py
+12
-25
@@ -36,7 +36,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
37
from vllm.model_executor.layers.rotary_embedding import get_rope
38
from vllm.model_executor.layers.vocab_parallel_embedding import (
39
- DEFAULT_VOCAB_PADDING_SIZE,
39
ParallelLMHead,
40
VocabParallelEmbedding,
41
)
@@ -770,37 +769,17 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
769
):
770
super().__init__()
771
config = vllm_config.model_config.hf_config
773
- lora_config = vllm_config.lora_config
774
- self.config = config
775
- self.vllm_config = vllm_config
776
-
772
self.model = Step3p5Model(
773
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
774
)
780
-
781
- self.moe_layers: list[FusedMoEBlock] = []
782
- for layer in self.model.layers:
783
- if isinstance(layer, PPMissingLayer):
784
- continue
785
- assert isinstance(layer, Step3p5DecoderLayer)
786
- if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock):
787
- self.moe_layers.append(layer.moe)
788
-
775
if get_pp_group().is_last_rank:
790
- self.unpadded_vocab_size = config.vocab_size
791
- if lora_config:
792
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
776
self.lm_head = ParallelLMHead(
794
- self.unpadded_vocab_size,
777
+ config.vocab_size,
778
config.hidden_size,
796
- org_num_embeddings=config.vocab_size,
797
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
798
- if not lora_config
799
- else lora_config.lora_vocab_padding_size,
800
- )
801
- self.logits_processor = LogitsProcessor(
802
- self.unpadded_vocab_size, config.vocab_size
779
+ quant_config=vllm_config.quant_config,
780
+ prefix=maybe_prefix(prefix, "lm_head"),
781
)
782
+ self.logits_processor = LogitsProcessor(config.vocab_size)
783
else:
784
self.lm_head = PPMissingLayer()
785
@@ -809,6 +788,14 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
788
)
789
790
# Set MoE hyperparameters
791
+ self.moe_layers: list[FusedMoEBlock] = []
792
+ for layer in self.model.layers:
793
+ if isinstance(layer, PPMissingLayer):
794
+ continue
795
+ assert isinstance(layer, Step3p5DecoderLayer)
796
+ if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock):
797
+ self.moe_layers.append(layer.moe)
798
+
799
self.expert_weights = []
800
assert len(self.moe_layers) > 0, "No MoE layers found in the model."
801
example_layer = self.moe_layers[0]