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 9, 00:00 UTC
[MISC] Fix Tensor Parallelism for Quantized Mamba Models with n_groups=1 (#33257)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Feb 4, 2026 at 00:10 UTC
a372f3f40afd0aed802242ce59b6a2640d4ef59e
1 file changed
+84
-118
vllm/model_executor/layers/mamba/mamba_mixer2.py
+84
-118
@@ -17,7 +17,6 @@ from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
- MergedColumnParallelLinear,
20
RowParallelLinear,
21
)
22
from vllm.model_executor.layers.mamba.abstract import MambaBase
@@ -40,6 +39,7 @@ from vllm.model_executor.model_loader.weight_utils import (
39
composed_weight_loader,
40
sharded_weight_loader,
41
)
42
+from vllm.model_executor.parameter import BasevLLMParameter
43
from vllm.model_executor.utils import set_weight_attrs
44
from vllm.platforms import current_platform
45
from vllm.utils.torch_utils import direct_register_custom_op
@@ -280,13 +280,6 @@ class MambaMixer2(MambaBase, CustomOp):
280
"then num_groups must equal 1."
281
)
282
283
- assert (
284
- (n_groups % self.tp_size == 0) or self.tp_size == 1 or quant_config is None
285
- ), (
286
- "Tensor parallel currently supported for quantized models only "
287
- "if tensor parallel world size divides num groups."
288
- )
289
-
283
self.ssm_state_size = ssm_state_size
284
self.conv_kernel_size = conv_kernel_size
285
self.activation = activation
@@ -308,121 +301,94 @@ class MambaMixer2(MambaBase, CustomOp):
301
self.groups_ssm_state_size = self.n_groups * self.ssm_state_size
302
self.conv_dim = intermediate_size + 2 * self.groups_ssm_state_size
303
311
- if n_groups % self.tp_size == 0:
312
- self.conv1d = MergedColumnParallelLinear(
313
- input_size=conv_kernel_size,
314
- output_sizes=[
315
- intermediate_size,
316
- self.groups_ssm_state_size,
317
- self.groups_ssm_state_size,
318
- ],
319
- bias=use_conv_bias,
320
- quant_config=None,
321
- prefix=f"{prefix}.conv1d",
322
- )
304
+ # Use ColumnParallelLinear with custom weight loaders for both cases:
305
+ # - When n_groups % tp_size == 0: standard sharding without duplication
306
+ # - When n_groups == 1: groups are duplicated across TP ranks
307
+ # The custom weight loader handles both cases correctly.
308
+
309
+ self.conv1d = ColumnParallelLinear(
310
+ input_size=conv_kernel_size,
311
+ output_size=self.conv_dim,
312
+ bias=use_conv_bias,
313
+ quant_config=None,
314
+ prefix=f"{prefix}.conv1d",
315
+ )
316
324
- self.in_proj = MergedColumnParallelLinear(
325
- input_size=hidden_size,
326
- output_sizes=[
327
- intermediate_size,
328
- intermediate_size,
329
- self.groups_ssm_state_size,
330
- self.groups_ssm_state_size,
331
- self.num_heads,
332
- ],
333
- bias=use_bias,
334
- quant_config=quant_config,
335
- prefix=f"{prefix}.in_proj",
336
- )
337
- else:
338
- # This is the n_groups == 1 case,
339
- # where we need to duplicate groups if TP>1.
340
-
341
- self.conv1d = ColumnParallelLinear(
342
- input_size=conv_kernel_size,
343
- output_size=self.conv_dim,
344
- bias=use_conv_bias,
345
- quant_config=None,
346
- prefix=f"{prefix}.conv1d",
347
- )
317
+ self.in_proj = ColumnParallelLinear(
318
+ input_size=hidden_size,
319
+ output_size=intermediate_size + self.conv_dim + self.num_heads,
320
+ bias=use_bias,
321
+ quant_config=quant_config,
322
+ prefix=f"{prefix}.in_proj",
323
+ )
324
349
- self.in_proj = ColumnParallelLinear(
350
- input_size=hidden_size,
351
- output_size=intermediate_size + self.conv_dim + self.num_heads,
352
- bias=use_bias,
353
- quant_config=quant_config,
354
- prefix=f"{prefix}.in_proj",
355
- )
325
+ # Configure shard settings for the custom weight loader:
326
+ # - group_shard_settings handles group duplication when n_groups == 1
327
+ # - When n_groups % tp_size == 0, extra=0 and duplicate_groups=False
328
+ group_shard_settings = (
329
+ self.groups_ssm_state_size, # expected model size
330
+ (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned
331
+ n_groups == 1, # duplicate groups when n_groups == 1
332
+ )
333
+ intermediate_settings = (intermediate_size, 0, False)
334
+ head_settings = (self.num_heads, 0, False)
335
+
336
+ # Apply custom weight loaders for conv1d (bias and weight)
337
+ delattr(self.conv1d.bias, "weight_loader")
338
+ set_weight_attrs(
339
+ self.conv1d.bias,
340
+ {
341
+ "weight_loader": mamba_v2_sharded_weight_loader(
342
+ [
343
+ intermediate_settings,
344
+ group_shard_settings,
345
+ group_shard_settings,
346
+ ],
347
+ self.tp_size,
348
+ tp_rank,
349
+ )
350
+ },
351
+ )
352
357
- # - because in_proj is a concatenation of 3 weights, we
358
- # need to interleave them before sharding
359
- # - use the custom weight loader mamba_v2_sharded_weight_loader
360
- # for conv1d.bias, covn1d.weight and in_proj.weight
361
- # - need to set these settings, to assign the groups
362
- # to the head shards
363
- group_shard_settings = (
364
- self.groups_ssm_state_size, # expected model size
365
- (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned
366
- n_groups == 1, # if there was only one group
367
- )
368
- intermediate_settings = (intermediate_size, 0, False)
369
- head_settings = (self.num_heads, 0, False)
370
-
371
- # - the weight already has a "weight_loader" attribute
372
- # which set_weight_attrs will raise if we do not
373
- # delete before trying to override it
374
- # - ditto for the other two weights below
375
- delattr(self.conv1d.bias, "weight_loader")
376
- set_weight_attrs(
377
- self.conv1d.bias,
378
- {
379
- "weight_loader": mamba_v2_sharded_weight_loader(
380
- [
381
- intermediate_settings,
382
- group_shard_settings,
383
- group_shard_settings,
384
- ],
385
- self.tp_size,
386
- tp_rank,
387
- )
388
- },
389
- )
353
+ delattr(self.conv1d.weight, "weight_loader")
354
+ set_weight_attrs(
355
+ self.conv1d.weight,
356
+ {
357
+ "weight_loader": mamba_v2_sharded_weight_loader(
358
+ [
359
+ intermediate_settings,
360
+ group_shard_settings,
361
+ group_shard_settings,
362
+ ],
363
+ self.tp_size,
364
+ tp_rank,
365
+ )
366
+ },
367
+ )
368
391
- delattr(self.conv1d.weight, "weight_loader")
392
- set_weight_attrs(
393
- self.conv1d.weight,
394
- {
395
- "weight_loader": mamba_v2_sharded_weight_loader(
396
- [
397
- intermediate_settings,
398
- group_shard_settings,
399
- group_shard_settings,
400
- ],
401
- self.tp_size,
402
- tp_rank,
403
- )
404
- },
405
- )
369
+ # Create the custom weight loader for in_proj
370
+ mamba_loader = mamba_v2_sharded_weight_loader(
371
+ [
372
+ intermediate_settings, # for gate
373
+ intermediate_settings,
374
+ group_shard_settings,
375
+ group_shard_settings,
376
+ head_settings, # for dt
377
+ ],
378
+ self.tp_size,
379
+ tp_rank,
380
+ )
381
407
- if quant_config is None:
408
- # - quant layers do not have a weight loader
409
- delattr(self.in_proj.weight, "weight_loader")
410
- set_weight_attrs(
411
- self.in_proj.weight,
412
- {
413
- "weight_loader": mamba_v2_sharded_weight_loader(
414
- [
415
- intermediate_settings, # for gate
416
- intermediate_settings,
417
- group_shard_settings,
418
- group_shard_settings,
419
- head_settings, # for dt
420
- ],
421
- self.tp_size,
422
- tp_rank,
423
- )
424
- },
425
- )
382
+ # Apply the custom weight loader to in_proj.weight
383
+ # Works for both non-quantized (Parameter) and quantized
384
+ # (ModelWeightParameter which extends BasevLLMParameter)
385
+ if isinstance(self.in_proj.weight, BasevLLMParameter):
386
+ # For BasevLLMParameter subclasses (quantized layers like FP8)
387
+ self.in_proj.weight.weight_loader = mamba_loader
388
+ else:
389
+ # For standard Parameter (non-quantized layers)
390
+ delattr(self.in_proj.weight, "weight_loader")
391
+ set_weight_attrs(self.in_proj.weight, {"weight_loader": mamba_loader})
392
393
# unsqueeze to fit conv1d weights shape into the linear weights shape.
394
# Can't do this in `weight_loader` since it already exists in