[Bugfix] Preserve Marlin runtime tensor storage across weight reload (#48438)
Signed-off-by: Ryan Clark <ryanclark2k@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Ryan Clark committed
Jul 30, 2026 at 14:50 UTC
68fb303a4177f52b2a2795abe602031a6d330b83
6 files changed
+330
-17
tests/model_executor/model_loader/test_reload.py
+277
@@ -167,6 +167,283 @@ def test_materialize_layer_preserves_non_meta_tensors():
167
assert torch.equal(layer.bias.data, bias_values)
168
169
170
+_MARLIN_SIZE_K, _MARLIN_SIZE_N, _MARLIN_GROUP_SIZE = 128, 64, 64
171
+
172
+
173
+def _stub_marlin_ops(monkeypatch):
174
+ from vllm import _custom_ops as ops
175
+ from vllm.model_executor.layers.quantization.utils import marlin_utils
176
+
177
+ monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
178
+ monkeypatch.setattr(
179
+ ops,
180
+ "gptq_marlin_repack",
181
+ lambda w, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
182
+ size_k // 16, size_n * 2, dtype=torch.int32
183
+ ),
184
+ )
185
+
186
+
187
+def _make_act_order_marlin_kernel():
188
+ from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
189
+ MarlinLinearKernel,
190
+ )
191
+ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
192
+ MPLinearLayerConfig,
193
+ )
194
+ from vllm.scalar_type import scalar_types
195
+
196
+ kernel = object.__new__(MarlinLinearKernel)
197
+ kernel.config = MPLinearLayerConfig(
198
+ full_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N),
199
+ partition_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N),
200
+ weight_type=scalar_types.uint4b8,
201
+ act_type=torch.float16,
202
+ group_size=_MARLIN_GROUP_SIZE,
203
+ zero_points=False,
204
+ has_g_idx=True,
205
+ )
206
+ kernel.w_q_name = "qweight"
207
+ kernel.w_s_name = "scales"
208
+ kernel.w_zp_name = None
209
+ kernel.w_gidx_name = "g_idx"
210
+ return kernel
211
+
212
+
213
+def _load_marlin_checkpoint_format_weights(layer, g_idx):
214
+ from vllm.model_executor.parameter import (
215
+ GroupQuantScaleParameter,
216
+ PackedvLLMParameter,
217
+ RowvLLMParameter,
218
+ )
219
+
220
+ layer.qweight = PackedvLLMParameter(
221
+ data=torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32),
222
+ input_dim=0,
223
+ output_dim=1,
224
+ packed_dim=0,
225
+ packed_factor=8,
226
+ weight_loader=default_weight_loader,
227
+ )
228
+ layer.scales = GroupQuantScaleParameter(
229
+ data=torch.ones(
230
+ _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16
231
+ ),
232
+ input_dim=0,
233
+ output_dim=1,
234
+ weight_loader=default_weight_loader,
235
+ )
236
+ layer.g_idx = RowvLLMParameter(
237
+ data=g_idx.clone(),
238
+ input_dim=0,
239
+ weight_loader=default_weight_loader,
240
+ )
241
+
242
+
243
+def _random_g_idx(generator):
244
+ return torch.randint(
245
+ 0,
246
+ _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE,
247
+ (_MARLIN_SIZE_K,),
248
+ dtype=torch.int32,
249
+ generator=generator,
250
+ )
251
+
252
+
253
+def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init):
254
+ """Marlin workspace and act-order sort indices must be recomputed into
255
+ the same storage when weights are reloaded (RL weight sync), so device
256
+ addresses captured by CUDA graphs remain valid."""
257
+ from vllm.model_executor.layers.quantization.utils import marlin_utils
258
+
259
+ _stub_marlin_ops(monkeypatch)
260
+ kernel = _make_act_order_marlin_kernel()
261
+
262
+ generator = torch.Generator().manual_seed(0)
263
+ first_g_idx = _random_g_idx(generator)
264
+ second_g_idx = _random_g_idx(generator)
265
+
266
+ layer = torch.nn.Module()
267
+ _load_marlin_checkpoint_format_weights(layer, first_g_idx)
268
+ kernel.process_weights_after_loading(layer)
269
+
270
+ workspace_ptr = kernel.workspace.data_ptr()
271
+ sort_indices_ptr = layer.g_idx_sort_indices.data_ptr()
272
+
273
+ # Reload: fresh checkpoint-format tensors with a different act-order
274
+ _load_marlin_checkpoint_format_weights(layer, second_g_idx)
275
+ kernel.process_weights_after_loading(layer)
276
+
277
+ assert kernel.workspace.data_ptr() == workspace_ptr
278
+ assert torch.all(kernel.workspace == 0)
279
+ assert layer.g_idx_sort_indices.data_ptr() == sort_indices_ptr
280
+ expected_sort_indices = marlin_utils.marlin_sort_g_idx(second_g_idx)[1]
281
+ assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices)
282
+ # registered as a Parameter so layerwise reload copy-back preserves it
283
+ assert isinstance(layer.g_idx_sort_indices, torch.nn.Parameter)
284
+
285
+
286
+@pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"])
287
+def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):
288
+ """The Marlin fallback prepare_* functions rerun on weight reload and must
289
+ reuse the workspace storage whose address captured CUDA graphs hold."""
290
+ from vllm import _custom_ops as ops
291
+ from vllm.model_executor.layers.quantization.utils import (
292
+ marlin_utils,
293
+ marlin_utils_fp4,
294
+ marlin_utils_fp8,
295
+ )
296
+
297
+ size_k, size_n = 128, 64
298
+
299
+ monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
300
+ monkeypatch.setattr(
301
+ ops,
302
+ "gptq_marlin_repack",
303
+ lambda b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros(
304
+ size_k // 16, size_n * 2, dtype=torch.int32
305
+ ),
306
+ )
307
+
308
+ layer = torch.nn.Module()
309
+ layer.output_size_per_partition = size_n
310
+ layer.input_size_per_partition = size_k
311
+ layer.orig_dtype = torch.float16
312
+ layer.params_dtype = torch.float16
313
+
314
+ if variant == "fp8":
315
+ prepare = marlin_utils_fp8.prepare_fp8_layer_for_marlin
316
+
317
+ def load_checkpoint_format_weights():
318
+ layer.weight = torch.nn.Parameter(
319
+ torch.zeros(size_k, size_n, dtype=torch.float8_e4m3fn),
320
+ requires_grad=False,
321
+ )
322
+ layer.weight_scale = torch.nn.Parameter(
323
+ torch.ones(1, dtype=torch.float32), requires_grad=False
324
+ )
325
+ elif variant == "mxfp8":
326
+ prepare = marlin_utils_fp8.prepare_mxfp8_layer_for_marlin
327
+
328
+ def load_checkpoint_format_weights():
329
+ layer.weight = torch.nn.Parameter(
330
+ torch.zeros(size_n, size_k, dtype=torch.float8_e4m3fn),
331
+ requires_grad=False,
332
+ )
333
+ layer.weight_scale = torch.nn.Parameter(
334
+ torch.full((size_n, size_k // 32), 127, dtype=torch.uint8),
335
+ requires_grad=False,
336
+ )
337
+ else:
338
+ prepare = marlin_utils_fp4.prepare_fp4_layer_for_marlin
339
+
340
+ def load_checkpoint_format_weights():
341
+ layer.weight = torch.nn.Parameter(
342
+ torch.zeros(size_n, size_k // 2, dtype=torch.uint8),
343
+ requires_grad=False,
344
+ )
345
+ layer.weight_scale = torch.nn.Parameter(
346
+ torch.ones(size_n, size_k // 16, dtype=torch.float8_e4m3fn),
347
+ requires_grad=False,
348
+ )
349
+ layer.weight_global_scale = torch.nn.Parameter(
350
+ torch.ones(1, dtype=torch.float32), requires_grad=False
351
+ )
352
+
353
+ load_checkpoint_format_weights()
354
+ prepare(layer)
355
+ workspace_ptr = layer.workspace.data_ptr()
356
+
357
+ # Reload: fresh checkpoint-format tensors, prepare runs again
358
+ load_checkpoint_format_weights()
359
+ prepare(layer)
360
+
361
+ assert layer.workspace.data_ptr() == workspace_ptr
362
+ assert torch.all(layer.workspace == 0)
363
+
364
+
365
+def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch):
366
+ """An incompatible existing workspace means the address captured by CUDA
367
+ graphs is already unusable; allocating a replacement would hide that."""
368
+ from vllm.model_executor.layers.quantization.utils import marlin_utils
369
+
370
+ monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
371
+ device = torch.device("cpu")
372
+
373
+ workspace = marlin_utils.marlin_make_workspace_new(device)
374
+ reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace)
375
+ assert reused is workspace
376
+
377
+ with pytest.raises(ValueError, match="incompatible"):
378
+ marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace)
379
+ with pytest.raises(ValueError, match="incompatible"):
380
+ marlin_utils.marlin_make_workspace_new(
381
+ device, existing=workspace.to(torch.int64)
382
+ )
383
+
384
+
385
+def test_marlin_act_order_layerwise_reload_accounting(monkeypatch, dist_init):
386
+ """`g_idx_sort_indices` is generated during weight processing and never
387
+ loaded from checkpoints. Registering it as a Parameter must not count it
388
+ toward `load_numel_total`: reload restores the construction-time tensor
389
+ set before sizing, so act-order layers still process during streaming
390
+ instead of deferring (and buffering weights) until finalization."""
391
+ from vllm.model_executor.layers.quantization.base_config import (
392
+ QuantizeMethodBase,
393
+ )
394
+ from vllm.model_executor.layers.quantization.utils import marlin_utils
395
+ from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info
396
+
397
+ _stub_marlin_ops(monkeypatch)
398
+ kernel = _make_act_order_marlin_kernel()
399
+
400
+ class _KernelQuantMethod(QuantizeMethodBase):
401
+ def create_weights(self, layer, *args, **kwargs):
402
+ raise NotImplementedError
403
+
404
+ def apply(self, layer, *args, **kwargs):
405
+ raise NotImplementedError
406
+
407
+ def process_weights_after_loading(self, layer):
408
+ kernel.process_weights_after_loading(layer)
409
+
410
+ generator = torch.Generator().manual_seed(0)
411
+ layer = torch.nn.Module()
412
+ layer.quant_method = _KernelQuantMethod()
413
+ _load_marlin_checkpoint_format_weights(layer, _random_g_idx(generator))
414
+
415
+ # Metadata is recorded at model construction, before any processing
416
+ record_metadata_for_reloading(layer)
417
+ checkpoint_numel = sum(t.numel() for t in get_layer_tensors(layer).values())
418
+
419
+ kernel.process_weights_after_loading(layer)
420
+ sort_indices = layer.g_idx_sort_indices
421
+
422
+ initialize_layerwise_reload(layer)
423
+ info = get_layerwise_info(layer)
424
+ assert info.load_numel_total == checkpoint_numel
425
+
426
+ # Stream a new checkpoint; the layer must process as soon as its last
427
+ # tensor arrives
428
+ new_g_idx = _random_g_idx(generator)
429
+ checkpoint = {
430
+ "qweight": torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32),
431
+ "scales": torch.ones(
432
+ _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16
433
+ ),
434
+ "g_idx": new_g_idx,
435
+ }
436
+ for name, weight in checkpoint.items():
437
+ param = getattr(layer, name)
438
+ param.weight_loader(param, weight)
439
+
440
+ assert not info.can_load()
441
+ assert not info.loaded_weights
442
+ assert layer.g_idx_sort_indices is sort_indices
443
+ expected_sort_indices = marlin_utils.marlin_sort_g_idx(new_g_idx)[1]
444
+ assert torch.equal(layer.g_idx_sort_indices.data, expected_sort_indices)
445
+
446
+
447
def test_model_cleanup(dist_init, default_vllm_config):
448
layer = QKVParallelLinear(2, 3, 4)
449
assert layer.weight.weight_loader.__self__ is layer
vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+8
-3
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
25
unpack_cols,
26
)
27
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
28
+from vllm.model_executor.utils import replace_parameter
29
from vllm.platforms import current_platform
30
from vllm.scalar_type import scalar_types
31
@@ -111,8 +112,10 @@ class MarlinLinearKernel(MPLinearKernel):
112
else:
113
padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
114
114
- # Allocate marlin workspace.
115
- self.workspace = marlin_make_workspace_new(device)
115
+ # Allocate marlin workspace, reusing existing storage on reload.
116
+ self.workspace = marlin_make_workspace_new(
117
+ device, existing=getattr(self, "workspace", None)
118
+ )
119
120
# Default names since marlin requires empty parameters for these,
121
# TODO: remove this requirement from marlin (allow optional tensors)
@@ -174,7 +177,9 @@ class MarlinLinearKernel(MPLinearKernel):
177
getattr(layer, self.w_gidx_name)
178
)
179
self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
177
- layer.g_idx_sort_indices = g_idx_sort_indices
180
+ replace_parameter(
181
+ layer, "g_idx_sort_indices", g_idx_sort_indices, prefer_copy=True
182
+ )
183
else:
184
setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
185
layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+3
-1
@@ -504,7 +504,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
504
and self.wna16_backend != WNA16MoEBackend.EMULATION
505
):
506
layer.workspace = marlin_make_workspace_new(
507
- layer.w13_weight_g_idx.device, 4
507
+ layer.w13_weight_g_idx.device,
508
+ 4,
509
+ existing=getattr(layer, "workspace", None),
510
)
511
512
# Alias packed weights to w13_weight/w2_weight for the modular kernel interface
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+21
-4
@@ -397,14 +397,31 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso
397
398
399
def marlin_make_workspace_new(
400
- device: torch.device, max_blocks_per_sm: int = 1
400
+ device: torch.device,
401
+ max_blocks_per_sm: int = 1,
402
+ existing: torch.Tensor | None = None,
403
) -> torch.Tensor:
404
# In the new marlin kernel, we use the num of threadblocks as workspace
405
# size. The num of threadblocks is sms_count * max_blocks_per_sm.
406
sms = num_compute_units(device.index)
405
- return torch.zeros(
406
- sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False
407
- )
407
+ size = sms * max_blocks_per_sm
408
+ # On weight reload, reuse the existing storage so the workspace address
409
+ # captured by CUDA graphs stays valid.
410
+ if existing is not None:
411
+ if (
412
+ existing.device != device
413
+ or existing.dtype != torch.int
414
+ or existing.numel() != size
415
+ ):
416
+ raise ValueError(
417
+ f"Existing Marlin workspace is incompatible "
418
+ f"(device={existing.device}, dtype={existing.dtype}, "
419
+ f"numel={existing.numel()}; expected device={device}, "
420
+ f"dtype={torch.int}, numel={size}). Reload must reuse the "
421
+ f"workspace storage captured by CUDA graphs."
422
+ )
423
+ return existing.zero_()
424
+ return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False)
425
426
427
def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool:
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+9
-3
@@ -240,7 +240,9 @@ def prepare_fp4_layer_for_marlin(
240
device = layer.weight.device
241
242
# WORKSPACE
243
- layer.workspace = marlin_make_workspace_new(device)
243
+ layer.workspace = marlin_make_workspace_new(
244
+ device, existing=getattr(layer, "workspace", None)
245
+ )
246
247
# WEIGHT
248
# Repack weights to marlin format
@@ -396,7 +398,9 @@ def prepare_nvfp4_moe_layer_for_marlin(
398
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
399
400
# WORKSPACE
399
- layer.workspace = marlin_make_workspace_new(device, 4)
401
+ layer.workspace = marlin_make_workspace_new(
402
+ device, 4, existing=getattr(layer, "workspace", None)
403
+ )
404
perm = torch.empty(0, dtype=torch.int, device=device)
405
406
# WEIGHT
@@ -485,7 +489,9 @@ def prepare_moe_fp4_layer_for_marlin(
489
# WORKSPACE
490
device = layer.w13_weight.device
491
param_dtype = layer.params_dtype
488
- layer.workspace = marlin_make_workspace_new(device, 4)
492
+ layer.workspace = marlin_make_workspace_new(
493
+ device, 4, existing=getattr(layer, "workspace", None)
494
+ )
495
perm = torch.empty(0, dtype=torch.int, device=device)
496
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
497
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+12
-6
@@ -132,7 +132,9 @@ def prepare_fp8_layer_for_marlin(
132
device = layer.weight.device
133
134
# WORKSPACE
135
- layer.workspace = marlin_make_workspace_new(device)
135
+ layer.workspace = marlin_make_workspace_new(
136
+ device, existing=getattr(layer, "workspace", None)
137
+ )
138
139
# WEIGHT
140
# Repack weights to marlin format
@@ -278,9 +280,9 @@ def prepare_fp8_moe_layer_for_marlin(
280
281
# WORKSPACE
282
device = layer.w13_weight.device
281
- # NOTE(rob): we do not need to register the workspace as a param
282
- # because it is not used as part of the weight reloading process.
283
- layer.workspace = marlin_make_workspace_new(device, 4)
283
+ layer.workspace = marlin_make_workspace_new(
284
+ device, 4, existing=getattr(layer, "workspace", None)
285
+ )
286
perm = torch.empty(0, dtype=torch.int, device=device)
287
288
# WEIGHT
@@ -463,7 +465,9 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
465
device = layer.weight.device
466
467
# WORKSPACE
466
- layer.workspace = marlin_make_workspace_new(device)
468
+ layer.workspace = marlin_make_workspace_new(
469
+ device, existing=getattr(layer, "workspace", None)
470
+ )
471
472
# WEIGHT - repack FP8 weights to Marlin format
473
perm = torch.empty(0, dtype=torch.int, device=device)
@@ -549,7 +553,9 @@ def prepare_mxfp8_moe_layer_for_marlin(
553
param_dtype = torch.get_default_dtype()
554
perm = torch.empty(0, dtype=torch.int, device=device)
555
552
- layer.workspace = marlin_make_workspace_new(device, 4)
556
+ layer.workspace = marlin_make_workspace_new(
557
+ device, 4, existing=getattr(layer, "workspace", None)
558
+ )
559
560
def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
561
if "w13" in name: