[Perf][DSpark] Collapse DeepSeek-V4.1 draft states before SP all-gather (#56903)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Juntian Liu committed
Sep 15, 2026 at 02:58 UTC
073f883b56a5b585366da874473db100bc6fff7e
2 files changed
+23
-3
tests/kernels/test_mhc_kernels.py
+21
@@ -147,6 +147,27 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
147
)
148
149
150
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
151
+@pytest.mark.parametrize("num_tokens", [1, 7, 32])
152
+@pytest.mark.parametrize("tp_size", [4, 8])
153
+def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size):
154
+ """Simulate SP on one GPU, including padding and ranks without valid tokens."""
155
+ set_random_seed(0)
156
+ padded_tokens = num_tokens + (-num_tokens) % tp_size
157
+ streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
158
+ pre_mix = torch.rand(padded_tokens, 4, device=DEVICE)
159
+ stream_shards = streams.chunk(tp_size)
160
+ mix_shards = pre_mix.chunk(tp_size)
161
+ before = hc_collapse_triton(
162
+ torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens]
163
+ )
164
+ after = torch.cat(
165
+ [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)]
166
+ )[:num_tokens]
167
+ assert after.shape == (num_tokens, 5120)
168
+ torch.testing.assert_close(after, before, atol=0, rtol=0)
169
+
170
+
171
@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
172
@pytest.mark.parametrize(
173
"num_tokens,sinkhorn_repeat,norm_eps",
vllm/models/deepseek_v41/nvidia/dspark.py
+2
-3
@@ -217,15 +217,14 @@ class DSparkDeepseekV4Model(nn.Module):
217
residual,
218
)
219
hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
220
- if self.use_sequence_parallel:
221
- hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
222
- pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
220
# Collapse the hc copies with the pre-mix from the last layer's FFN
221
# mixes — the mix the reference forward_head applies (hc_pre with the
222
# last block's ffn pre-mix). Return the PRE-norm head hidden;
223
# compute_logits applies self.norm.
224
assert pre_mix is not None
225
hidden_states = hc_collapse_triton(hidden_states, pre_mix)
226
+ if self.use_sequence_parallel:
227
+ hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
228
return hidden_states
229
230