@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 83fe99399e

[Spec Decode] Add NVFP4 DSpark gathered top-k projection (#55713)

Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Andrii Skliar committed Sep 9, 2026 at 19:33 UTC 83fe99399ec0603b32393a14324b65c67ad04af2
4 files changed +491 -3
tests/v1/spec_decode/test_dspark_topk.py
+183
@@ -1,10 +1,20 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import pytest
5 import torch
6 from torch import nn
7
8 +from vllm.model_executor.layers.quantization.modelopt import (
9 + ModelOptLinearMethod,
10 + resolve,
11 +)
12 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
13 + dequantize_to_dtype,
14 + nvfp4_gathered_bias,
15 +)
16 from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
17 +from vllm.platforms import current_platform
18
19
20 def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead:
@@ -14,6 +24,8 @@ def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead:
24 weight.shape[1], weight.shape[0], bias=False, dtype=weight.dtype
25 )
26 head.markov_w2.weight.data.copy_(weight)
27 + head.markov_w2._retain_weight_for_gather = False
28 + head.markov_w2.is_w4a16_nvfp4 = False
29 return head
30
31
@@ -56,3 +68,174 @@ def test_gathered_markov_bias_matches_dense_at_full_vocab():
68
69 expected = original + markov_embed @ weight.T * scale
70 torch.testing.assert_close(result, expected)
71 +
72 +
73 +@pytest.mark.parametrize("cuda_alike_platform", [False, True])
74 +def test_gathered_markov_bias_dequantizes_selected_w4a16_rows(
75 + monkeypatch, cuda_alike_platform
76 +):
77 + packed_weight = torch.tensor(
78 + [[0x00] * 8, [0x22] * 8, [0xAA] * 8, [0x31] * 8], dtype=torch.uint8
79 + )
80 + weight_scale = torch.ones((4, 1), dtype=torch.float8_e4m3fn)
81 +
82 + layer = nn.Module()
83 + layer.weight = nn.Parameter(packed_weight, requires_grad=False)
84 + layer.weight_scale = nn.Parameter(weight_scale, requires_grad=False)
85 + layer.weight_scale_2 = nn.Parameter(torch.ones(1), requires_grad=False)
86 + layer._retain_weight_for_gather = True
87 +
88 + spec, ctx, fmt = resolve("W4A16_NVFP4", type("Config", (), {"group_size": 16}), "")
89 + method = ModelOptLinearMethod(spec, ctx, fmt)
90 +
91 + class RepackingKernel:
92 + @staticmethod
93 + def process_weights_after_loading(layer):
94 + layer.weight = nn.Parameter(
95 + torch.zeros((1, 1), dtype=torch.int32), requires_grad=False
96 + )
97 + layer.weight_scale = nn.Parameter(
98 + torch.zeros((1, 1), dtype=torch.float32), requires_grad=False
99 + )
100 + layer.weight_global_scale = nn.Parameter(
101 + torch.zeros(1, dtype=torch.float32), requires_grad=False
102 + )
103 +
104 + method.kernel = RepackingKernel()
105 + method.process_weights_after_loading(layer)
106 + assert layer.is_w4a16_nvfp4
107 +
108 + head = DSparkMarkovHead.__new__(DSparkMarkovHead)
109 + nn.Module.__init__(head)
110 + head.markov_w2 = layer
111 + markov_embed = torch.tensor([[1.0] * 16, [0.5, -0.5] * 8], dtype=torch.float32)
112 + index = torch.tensor([[1, 3], [2, 0]])
113 + values = torch.tensor([[0.25, -0.5], [1.0, 2.0]])
114 + original_values = values.clone()
115 + logits = torch.full((2, 4), float("-inf"))
116 +
117 + # CPU tensors must use the fallback even on CUDA/ROCm hosts.
118 + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: cuda_alike_platform)
119 + result = head.apply_bias_gathered(markov_embed, logits, values, index)
120 +
121 + dense_weight = torch.tensor([[0.0] * 16, [1.0] * 16, [-1.0] * 16, [0.5, 1.5] * 8])
122 + expected = original_values + torch.bmm(
123 + dense_weight[index], markov_embed.unsqueeze(-1)
124 + ).squeeze(-1)
125 + torch.testing.assert_close(result.gather(1, index), expected)
126 + torch.testing.assert_close(values, original_values)
127 +
128 +
129 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
130 +def test_nvfp4_gathered_bias_rejects_non_16_group_size():
131 + device = torch.device("cuda")
132 + rank = 32
133 + vocab_size = 2
134 +
135 + with pytest.raises(ValueError, match="group size of 16"):
136 + nvfp4_gathered_bias(
137 + torch.zeros((1, rank), dtype=torch.bfloat16, device=device),
138 + torch.zeros((vocab_size, rank // 2), dtype=torch.uint8, device=device),
139 + torch.ones(
140 + (vocab_size, rank // 32), dtype=torch.float8_e4m3fn, device=device
141 + ),
142 + torch.ones(1, dtype=torch.float32, device=device),
143 + torch.zeros((1, 1), dtype=torch.bfloat16, device=device),
144 + torch.zeros((1, 1), dtype=torch.int64, device=device),
145 + torch.full(
146 + (1, vocab_size), float("-inf"), dtype=torch.bfloat16, device=device
147 + ),
148 + 1.0,
149 + )
150 +
151 +
152 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
153 +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
154 +def test_nvfp4_gathered_markov_bias_kernel_matches_reference(batch_size: int):
155 + torch.manual_seed(0)
156 + device = torch.device("cuda")
157 + topk = 256
158 + rank = 512
159 + vocab_size = 512
160 + alpha = 0.75
161 +
162 + packed_weight = torch.randint(
163 + 0, 256, (vocab_size, rank // 2), dtype=torch.uint8, device=device
164 + )
165 + weight_scale = torch.full(
166 + (vocab_size, rank // 16),
167 + 0.5,
168 + dtype=torch.float8_e4m3fn,
169 + device=device,
170 + )
171 + global_scale = torch.tensor([0.25], dtype=torch.float32, device=device)
172 +
173 + markov_storage = torch.randn(
174 + batch_size, 2, rank, dtype=torch.bfloat16, device=device
175 + )
176 + values_storage = torch.randn(
177 + batch_size, 2, topk, dtype=torch.bfloat16, device=device
178 + )
179 + index_storage = torch.empty(batch_size, 2, topk, dtype=torch.int64, device=device)
180 + for batch_idx in range(batch_size):
181 + index_storage[batch_idx, 0] = (
182 + torch.arange(topk, device=device) + batch_idx
183 + ) % vocab_size
184 + logits_storage = torch.full(
185 + (batch_size, 2, vocab_size),
186 + float("-inf"),
187 + dtype=torch.bfloat16,
188 + device=device,
189 + )
190 +
191 + markov_embed = markov_storage[:, 0]
192 + values = values_storage[:, 0]
193 + index = index_storage[:, 0]
194 + logits = logits_storage[:, 0]
195 +
196 + layer = nn.Module()
197 + layer.register_buffer("_nvfp4_weight_for_gather", packed_weight)
198 + layer.register_buffer("_nvfp4_weight_scale_for_gather", weight_scale)
199 + layer.register_buffer("_nvfp4_weight_global_scale_for_gather", global_scale)
200 + layer._nvfp4_group_size_for_gather = 16
201 + layer._retain_weight_for_gather = True
202 + layer.is_w4a16_nvfp4 = True
203 + head = DSparkMarkovHead.__new__(DSparkMarkovHead)
204 + nn.Module.__init__(head)
205 + head.markov_w2 = layer
206 +
207 + flat_index = index.reshape(-1)
208 + selected_weight = dequantize_to_dtype(
209 + packed_weight.index_select(0, flat_index),
210 + weight_scale.index_select(0, flat_index),
211 + global_scale,
212 + dtype=torch.bfloat16,
213 + block_size=16,
214 + swizzle=False,
215 + ).view(batch_size, topk, rank)
216 + expected = (
217 + values
218 + + torch.bmm(selected_weight, markov_embed.unsqueeze(-1)).squeeze(-1) * alpha
219 + )
220 +
221 + result = head.apply_bias_gathered(markov_embed, logits, values, index, alpha)
222 +
223 + torch.testing.assert_close(
224 + result.gather(1, index), expected, rtol=1e-2, atol=6.25e-2
225 + )
226 + selected = torch.zeros_like(result, dtype=torch.bool).scatter_(1, index, True)
227 + assert torch.isneginf(result.masked_select(~selected)).all()
228 +
229 + if batch_size == 1:
230 + logits.fill_(float("-inf"))
231 + graph = torch.cuda.CUDAGraph()
232 + with torch.cuda.graph(graph):
233 + captured = head.apply_bias_gathered(
234 + markov_embed, logits, values, index, alpha
235 + )
236 + logits.fill_(float("-inf"))
237 + graph.replay()
238 + assert captured is logits
239 + torch.testing.assert_close(
240 + logits.gather(1, index), expected, rtol=1e-2, atol=6.25e-2
241 + )
vllm/model_executor/layers/quantization/modelopt.py
+24
@@ -2458,6 +2458,30 @@ class ModelOptLinearMethod(LinearMethodBase):
2458 self.akey.process(layer, ACT)
2459 maybe_fuse_global_scales(layer)
2460 self.fmt.post_process(layer)
2461 + layer.is_w4a16_nvfp4 = (
2462 + self.spec.weight == kNvfp4Static and self.spec.activation is None
2463 + )
2464 + if getattr(layer, "_retain_weight_for_gather", False):
2465 + if not layer.is_w4a16_nvfp4:
2466 + raise NotImplementedError(
2467 + "Gathered projection is only supported for ModelOpt W4A16 "
2468 + "NVFP4 weights."
2469 + )
2470 + assert self.ctx.group_size is not None
2471 + layer.register_buffer(
2472 + "_nvfp4_weight_for_gather", layer.weight.detach(), persistent=False
2473 + )
2474 + layer.register_buffer(
2475 + "_nvfp4_weight_scale_for_gather",
2476 + layer.weight_scale.detach(),
2477 + persistent=False,
2478 + )
2479 + layer.register_buffer(
2480 + "_nvfp4_weight_global_scale_for_gather",
2481 + layer.weight_global_scale.detach(),
2482 + persistent=False,
2483 + )
2484 + layer._nvfp4_group_size_for_gather = self.ctx.group_size
2485 self.kernel.process_weights_after_loading(layer)
2486
2487 def apply(self, layer, x, bias=None):
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+227 -1
@@ -7,10 +7,12 @@ import torch
7 from vllm.platforms import current_platform
8 from vllm.scalar_type import scalar_types
9 from vllm.triton_utils import tl, triton
10 +from vllm.utils.torch_utils import direct_register_custom_op
11
12 __all__ = [
13 "break_fp4_bytes",
14 "dequantize_to_dtype",
15 + "nvfp4_gathered_bias",
16 "ref_nvfp4_quant",
17 ]
18
@@ -43,6 +45,230 @@ def _e2m1_inline(nibble):
45 return tl.where(sign == 1, -val, val)
46
47
48 +@triton.jit
49 +def _nvfp4_gathered_bias_kernel(
50 + markov_ptr,
51 + packed_weight_ptr,
52 + weight_scale_ptr,
53 + global_scale_ptr,
54 + values_ptr,
55 + index_ptr,
56 + logits_ptr,
57 + stride_markov_b,
58 + stride_markov_k,
59 + stride_weight_v,
60 + stride_weight_k,
61 + stride_scale_v,
62 + stride_scale_k,
63 + stride_values_b,
64 + stride_values_j,
65 + stride_index_b,
66 + stride_index_j,
67 + stride_logits_b,
68 + stride_logits_v,
69 + alpha,
70 + TOPK: tl.constexpr,
71 + K: tl.constexpr,
72 + NUM_BLOCKS: tl.constexpr,
73 +):
74 + """Apply selected NVFP4 Markov-projection rows to dense logits.
75 +
76 + BF16: ``weight[index] -> baddbmm_ -> scatter_``.
77 + NVFP4: ``gather packed W2 row -> dequantize 16-element blocks -> dot with
78 + Markov embedding -> add base value -> scatter corrected logit``.
79 +
80 + The fused NVFP4 path is needed because packed weights cannot be indexed as
81 + BF16 rows, while dequantizing the full vocabulary would defeat top-k.
82 + """
83 + candidate_id = tl.program_id(0)
84 + request_id = candidate_id // TOPK
85 + candidate = candidate_id - request_id * TOPK
86 +
87 + vocab_id = tl.load(
88 + index_ptr + request_id * stride_index_b + candidate * stride_index_j
89 + ).to(tl.int64)
90 +
91 + block_ids = tl.arange(0, NUM_BLOCKS)
92 + byte_ids = tl.arange(0, 8)
93 + block_mask = block_ids * 16 < K
94 +
95 + packed_offsets = block_ids[:, None] * 8 + byte_ids[None, :]
96 + packed = tl.load(
97 + packed_weight_ptr
98 + + vocab_id * stride_weight_v
99 + + packed_offsets * stride_weight_k,
100 + mask=block_mask[:, None],
101 + other=0,
102 + )
103 + low = _e2m1_inline(packed & 0x0F)
104 + high = _e2m1_inline((packed >> 4) & 0x0F)
105 + decoded = tl.interleave(low, high)
106 +
107 + raw_weight_scale = tl.load(
108 + weight_scale_ptr + vocab_id * stride_scale_v + block_ids * stride_scale_k,
109 + mask=block_mask,
110 + other=0,
111 + )
112 + weight_scale = tl.cast(
113 + raw_weight_scale,
114 + tl.float8e4nv,
115 + bitcast=True,
116 + ).to(tl.float32)
117 + global_scale = tl.load(global_scale_ptr).to(tl.float32)
118 +
119 + element_ids = tl.arange(0, 16)
120 + k_offsets = block_ids[:, None] * 16 + element_ids[None, :]
121 + k_mask = k_offsets < K
122 + markov = tl.load(
123 + markov_ptr + request_id * stride_markov_b + k_offsets * stride_markov_k,
124 + mask=k_mask,
125 + other=0.0,
126 + )
127 +
128 + dequant_weight = (decoded * weight_scale[:, None] * global_scale).to(tl.bfloat16)
129 + products = dequant_weight.to(tl.float32) * markov.to(tl.float32)
130 + block_sums = tl.sum(products, axis=1)
131 + dot = tl.sum(block_sums, axis=0)
132 +
133 + base_value = tl.load(
134 + values_ptr + request_id * stride_values_b + candidate * stride_values_j
135 + ).to(tl.float32)
136 + corrected = base_value + alpha * dot
137 + tl.store(
138 + logits_ptr + request_id * stride_logits_b + vocab_id * stride_logits_v,
139 + corrected,
140 + )
141 +
142 +
143 +def _nvfp4_gathered_bias_impl(
144 + markov_embed: torch.Tensor,
145 + packed_weight: torch.Tensor,
146 + weight_scale: torch.Tensor,
147 + global_scale: torch.Tensor,
148 + values: torch.Tensor,
149 + index: torch.Tensor,
150 + logits: torch.Tensor,
151 + alpha: float,
152 +) -> None:
153 + if index.numel() == 0:
154 + return
155 +
156 + assert markov_embed.dtype == torch.bfloat16
157 + assert values.dtype == torch.bfloat16
158 + assert logits.dtype == torch.bfloat16
159 + assert packed_weight.dtype == torch.uint8
160 + assert weight_scale.dtype == torch.float8_e4m3fn
161 + assert index.dtype == torch.int64
162 + assert markov_embed.is_cuda
163 + assert all(
164 + tensor.device == markov_embed.device
165 + for tensor in (
166 + packed_weight,
167 + weight_scale,
168 + global_scale,
169 + values,
170 + index,
171 + logits,
172 + )
173 + )
174 + assert markov_embed.ndim == 2
175 + assert packed_weight.ndim == 2
176 + assert weight_scale.ndim == 2
177 + assert values.ndim == 2
178 + assert index.ndim == 2
179 + assert logits.ndim == 2
180 + assert global_scale.numel() == 1
181 + assert values.shape == index.shape
182 + assert logits.shape[0] == index.shape[0] == markov_embed.shape[0]
183 +
184 + k = markov_embed.shape[-1]
185 + if k % 16 != 0 or packed_weight.shape[-1] * 2 != k:
186 + raise ValueError(
187 + "nvfp4_gathered_bias requires k to be divisible by 16 and "
188 + f"packed_weight width to equal k / 2; got k={k} and "
189 + f"packed_weight width={packed_weight.shape[-1]}"
190 + )
191 + if weight_scale.shape[-1] != k // 16:
192 + raise ValueError(
193 + "nvfp4_gathered_bias only supports an NVFP4 group size of 16; "
194 + f"got {weight_scale.shape[-1]} scale blocks for k={k}"
195 + )
196 +
197 + topk = index.shape[1]
198 + num_blocks = triton.next_power_of_2(k // 16)
199 + scale_bytes = weight_scale.view(torch.uint8)
200 + _nvfp4_gathered_bias_kernel[(index.numel(),)](
201 + markov_embed,
202 + packed_weight,
203 + scale_bytes,
204 + global_scale,
205 + values,
206 + index,
207 + logits,
208 + markov_embed.stride(0),
209 + markov_embed.stride(1),
210 + packed_weight.stride(0),
211 + packed_weight.stride(1),
212 + scale_bytes.stride(0),
213 + scale_bytes.stride(1),
214 + values.stride(0),
215 + values.stride(1),
216 + index.stride(0),
217 + index.stride(1),
218 + logits.stride(0),
219 + logits.stride(1),
220 + alpha,
221 + TOPK=topk,
222 + K=k,
223 + NUM_BLOCKS=num_blocks,
224 + num_warps=4,
225 + num_stages=2,
226 + )
227 +
228 +
229 +def _nvfp4_gathered_bias_fake(
230 + markov_embed: torch.Tensor,
231 + packed_weight: torch.Tensor,
232 + weight_scale: torch.Tensor,
233 + global_scale: torch.Tensor,
234 + values: torch.Tensor,
235 + index: torch.Tensor,
236 + logits: torch.Tensor,
237 + alpha: float,
238 +) -> None:
239 + return None
240 +
241 +
242 +direct_register_custom_op(
243 + op_name="nvfp4_gathered_bias",
244 + op_func=_nvfp4_gathered_bias_impl,
245 + mutates_args=["logits"],
246 + fake_impl=_nvfp4_gathered_bias_fake,
247 +)
248 +
249 +
250 +def nvfp4_gathered_bias(
251 + markov_embed: torch.Tensor,
252 + packed_weight: torch.Tensor,
253 + weight_scale: torch.Tensor,
254 + global_scale: torch.Tensor,
255 + values: torch.Tensor,
256 + index: torch.Tensor,
257 + logits: torch.Tensor,
258 + alpha: float,
259 +) -> None:
260 + torch.ops.vllm.nvfp4_gathered_bias(
261 + markov_embed,
262 + packed_weight,
263 + weight_scale,
264 + global_scale,
265 + values,
266 + index,
267 + logits,
268 + alpha,
269 + )
270 +
271 +
272 @triton.jit
273 def _dequantize_nvfp4_kernel(
274 fp4_ptr,
@@ -347,7 +573,7 @@ def dequantize_to_dtype(
573 # Two fp4 values are packed into one uint8.
574 assert tensor_fp4.dtype == torch.uint8
575
350 - if not swizzle and current_platform.is_cuda_alike():
576 + if not swizzle and tensor_fp4.is_cuda and current_platform.is_cuda_alike():
577 return _triton_dequantize_nvfp4(
578 tensor_fp4, tensor_sf, global_scale, dtype, block_size
579 )
vllm/model_executor/models/qwen3_dspark.py
+57 -2
@@ -25,9 +25,14 @@ from vllm.logger import init_logger
25 from vllm.model_executor.layers.linear import ReplicatedLinear
26 from vllm.model_executor.layers.logits_processor import LogitsProcessor
27 from vllm.model_executor.layers.quantization import QuantizationConfig
28 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
29 + dequantize_to_dtype,
30 + nvfp4_gathered_bias,
31 +)
32 from vllm.model_executor.layers.vocab_parallel_embedding import (
33 ParallelLMHead,
34 )
35 +from vllm.platforms import current_platform
36
37 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
38 from .utils import (
@@ -60,6 +65,7 @@ class DSparkMarkovHead(nn.Module):
65 markov_rank: int,
66 prefix: str,
67 quant_config: QuantizationConfig | None = None,
68 + retain_weight_for_gather: bool = False,
69 ) -> None:
70 super().__init__()
71 self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
@@ -71,6 +77,8 @@ class DSparkMarkovHead(nn.Module):
77 prefix=maybe_prefix(prefix, "markov_w2"),
78 disable_tp=True,
79 )
80 + self.markov_w2._retain_weight_for_gather = retain_weight_for_gather
81 + self.markov_w2.is_w4a16_nvfp4 = False
82
83 def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
84 """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
@@ -98,8 +106,52 @@ class DSparkMarkovHead(nn.Module):
106 positions. This method scatters the corrected candidate values into
107 that dense buffer so the normal sampler sees the truncated proposal.
108 """
101 - weight = self.markov_w2.weight[index]
102 - corrected = values.unsqueeze(-1)
109 + if self.markov_w2._retain_weight_for_gather and self.markov_w2.is_w4a16_nvfp4:
110 + # Use the fast Triton kernel for W4A16 NVFP4 weights with BF16
111 + # activations on CUDA. FP16 activations, non-CUDA accelerators, and
112 + # CPU use the slower dequantize/scatter fallback below.
113 + if (
114 + current_platform.is_cuda()
115 + and markov_embed.is_cuda
116 + and markov_embed.dtype == torch.bfloat16
117 + and values.dtype == torch.bfloat16
118 + and logits.dtype == torch.bfloat16
119 + ):
120 + nvfp4_gathered_bias(
121 + markov_embed,
122 + self.markov_w2._nvfp4_weight_for_gather,
123 + self.markov_w2._nvfp4_weight_scale_for_gather,
124 + self.markov_w2._nvfp4_weight_global_scale_for_gather,
125 + values,
126 + index,
127 + logits,
128 + scale,
129 + )
130 + return logits
131 +
132 + flat_index = index.reshape(-1)
133 + packed_weight = self.markov_w2._nvfp4_weight_for_gather.index_select(
134 + 0, flat_index
135 + )
136 + weight_scale = self.markov_w2._nvfp4_weight_scale_for_gather.index_select(
137 + 0, flat_index
138 + )
139 + weight = dequantize_to_dtype(
140 + packed_weight,
141 + weight_scale,
142 + self.markov_w2._nvfp4_weight_global_scale_for_gather,
143 + dtype=markov_embed.dtype,
144 + block_size=self.markov_w2._nvfp4_group_size_for_gather,
145 + swizzle=False,
146 + ).view(*index.shape, -1)
147 + else:
148 + weight = self.markov_w2.weight[index]
149 + if weight.dtype != markov_embed.dtype:
150 + raise RuntimeError(
151 + "Gathered Markov projection requires unquantized weights or "
152 + "retained ModelOpt W4A16 NVFP4 weights."
153 + )
154 + corrected = values.unsqueeze(-1).clone()
155 corrected.baddbmm_(
156 weight,
157 markov_embed.unsqueeze(-1),
@@ -160,6 +212,9 @@ class Qwen3DSparkModel(DFlashQwen3Model):
212 config.markov_rank,
213 prefix=maybe_prefix(prefix, "markov_head"),
214 quant_config=self.quant_config,
215 + retain_weight_for_gather=(
216 + getattr(config, "dspark_draft_topk", None) is not None
217 + ),
218 )
219 self.confidence_head: DSparkConfidenceHead | None = None
220 if getattr(config, "enable_confidence_head", False):