Dflash2 load fix (#53435)
Karan Shelar committed
Aug 25, 2026 at 10:02 UTC
a9a17e7095a66ef6c6685a1c7ddd657781a78d3c
2 files changed
+115
-1
tests/v1/spec_decode/test_dflash2.py
+112
@@ -116,3 +116,115 @@ def test_selector_asks_for_fp32_proposal_logits():
116
117
assert dtype is torch.float32
118
assert fill == float("-inf")
119
+
120
+
121
+@pytest.mark.skip_global_cleanup
122
+def test_dflash2_model_decoder_layer_cls(monkeypatch):
123
+ from types import SimpleNamespace
124
+
125
+ from vllm.config import set_current_vllm_config
126
+ from vllm.model_executor.models.qwen3_dflash2 import (
127
+ DFlash2Qwen3DecoderLayer,
128
+ DFlash2Qwen3Model,
129
+ )
130
+
131
+ # 1. Mock get_current_vllm_config and TP groups
132
+ mock_current_vllm_config = SimpleNamespace(
133
+ cache_config=SimpleNamespace(
134
+ block_size=16,
135
+ user_specified_block_size=False,
136
+ kv_cache_dtype_skip_layers=[],
137
+ cache_dtype="auto",
138
+ sliding_window=None,
139
+ enable_prefix_caching=False,
140
+ ),
141
+ kv_transfer_config=None,
142
+ speculative_config=None,
143
+ attention_config=SimpleNamespace(
144
+ use_non_causal=False,
145
+ backend=None,
146
+ backend_per_kind={},
147
+ ),
148
+ parallel_config=SimpleNamespace(
149
+ prefill_context_parallel_size=1,
150
+ decode_context_parallel_size=1,
151
+ ),
152
+ compilation_config=SimpleNamespace(
153
+ compile_custom_ops=False,
154
+ custom_ops="all",
155
+ enabled_custom_ops=set(),
156
+ static_forward_context={},
157
+ mode=0, # CompilationMode.NONE is 0
158
+ ),
159
+ model_config=SimpleNamespace(
160
+ dtype=torch.float32,
161
+ is_mm_prefix_lm=False,
162
+ ),
163
+ )
164
+ from vllm.platforms import current_platform
165
+
166
+ monkeypatch.setattr(
167
+ current_platform,
168
+ "get_attn_backend_cls",
169
+ lambda *args, **kwargs: (
170
+ "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
171
+ ),
172
+ )
173
+
174
+ class MockGroup:
175
+ rank_in_group = 0
176
+ world_size = 1
177
+
178
+ monkeypatch.setattr(
179
+ "vllm.distributed.parallel_state._TP",
180
+ MockGroup(),
181
+ )
182
+
183
+ # 2. Mock vllm_config
184
+ hf_config = SimpleNamespace(
185
+ vocab_size=1000,
186
+ hidden_size=256,
187
+ num_hidden_layers=2,
188
+ num_attention_heads=8,
189
+ num_key_value_heads=2,
190
+ max_position_embeddings=2048,
191
+ rms_norm_eps=1e-6,
192
+ rope_parameters={},
193
+ intermediate_size=512,
194
+ hidden_act="silu",
195
+ dflash_config={
196
+ "selector_rank": 4,
197
+ "selector_top_k": 3,
198
+ "conv_kernel_size": 3,
199
+ "conv_group_size": 2,
200
+ "use_aux_hidden_state": False,
201
+ },
202
+ )
203
+ vllm_config = SimpleNamespace(
204
+ speculative_config=SimpleNamespace(
205
+ draft_model_config=SimpleNamespace(
206
+ hf_config=hf_config,
207
+ quantization=None,
208
+ ),
209
+ num_speculative_tokens=4,
210
+ enable_adaptive_verification=False,
211
+ ),
212
+ model_config=SimpleNamespace(
213
+ dtype=torch.float32,
214
+ is_mm_prefix_lm=False,
215
+ ),
216
+ load_config=SimpleNamespace(
217
+ quantization=None,
218
+ quantization_param_path=None,
219
+ ),
220
+ )
221
+ mock_current_vllm_config.speculative_config = vllm_config.speculative_config
222
+ vllm_config.compilation_config = mock_current_vllm_config.compilation_config
223
+
224
+ # 3. Instantiate the model under meta device to avoid parameter allocation issues
225
+ with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"):
226
+ model = DFlash2Qwen3Model(vllm_config=vllm_config)
227
+
228
+ # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass)
229
+ assert len(model.layers) == 2
230
+ assert isinstance(model.layers[0], DFlash2Qwen3DecoderLayer)
vllm/model_executor/models/qwen3_dflash.py
+3
-1
@@ -377,6 +377,8 @@ class DFlashQwen3DecoderLayer(nn.Module):
377
378
@support_torch_compile
379
class DFlashQwen3Model(nn.Module):
380
+ decoder_layer_cls = DFlashQwen3DecoderLayer
381
+
382
hf_to_vllm_mapper = WeightsMapper(
383
orig_to_new_substr={
384
"midlayer.": "layers.0.",
@@ -439,7 +441,7 @@ class DFlashQwen3Model(nn.Module):
441
442
self.layers = nn.ModuleList(
443
[
442
- DFlashQwen3DecoderLayer(
444
+ self.decoder_layer_cls(
445
current_vllm_config,
446
config=self.config,
447
layer_idx=layer_idx,