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 8, 06:15 UTC
[BUGFIX] Fix accuracy regression for NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 with TP>1 (#34476)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Feb 15, 2026 at 11:25 UTC
604b9eaec53a11ae5193348f6a623dd7cdef48bf
1 file changed
+117
-83
vllm/model_executor/layers/mamba/mamba_mixer2.py
+117
-83
@@ -17,6 +17,7 @@ from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.model_executor.custom_op import CustomOp, PluggableLayer
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
+ MergedColumnParallelLinear,
21
RowParallelLinear,
22
)
23
from vllm.model_executor.layers.mamba.abstract import MambaBase
@@ -301,94 +302,127 @@ class MambaMixer2(MambaBase, PluggableLayer):
302
self.groups_ssm_state_size = self.n_groups * self.ssm_state_size
303
self.conv_dim = intermediate_size + 2 * self.groups_ssm_state_size
304
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
- )
305
+ if n_groups % self.tp_size == 0:
306
+ self.conv1d = MergedColumnParallelLinear(
307
+ input_size=conv_kernel_size,
308
+ output_sizes=[
309
+ intermediate_size,
310
+ self.groups_ssm_state_size,
311
+ self.groups_ssm_state_size,
312
+ ],
313
+ bias=use_conv_bias,
314
+ quant_config=None,
315
+ prefix=f"{prefix}.conv1d",
316
+ )
317
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
- )
318
+ self.in_proj = MergedColumnParallelLinear(
319
+ input_size=hidden_size,
320
+ output_sizes=[
321
+ intermediate_size,
322
+ intermediate_size,
323
+ self.groups_ssm_state_size,
324
+ self.groups_ssm_state_size,
325
+ self.num_heads,
326
+ ],
327
+ bias=use_bias,
328
+ quant_config=quant_config,
329
+ prefix=f"{prefix}.in_proj",
330
+ )
331
+ else:
332
+ # This is the n_groups == 1 case,
333
+ # where we need to duplicate groups if TP>1.
334
+
335
+ self.conv1d = ColumnParallelLinear(
336
+ input_size=conv_kernel_size,
337
+ output_size=self.conv_dim,
338
+ bias=use_conv_bias,
339
+ quant_config=None,
340
+ prefix=f"{prefix}.conv1d",
341
+ )
342
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
- )
343
+ self.in_proj = ColumnParallelLinear(
344
+ input_size=hidden_size,
345
+ output_size=intermediate_size + self.conv_dim + self.num_heads,
346
+ bias=use_bias,
347
+ quant_config=quant_config,
348
+ prefix=f"{prefix}.in_proj",
349
+ )
350
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
- )
351
+ # - because in_proj is a concatenation of 3 weights, we
352
+ # need to interleave them before sharding
353
+ # - use the custom weight loader mamba_v2_sharded_weight_loader
354
+ # for conv1d.bias, covn1d.weight and in_proj.weight
355
+ # - need to set these settings, to assign the groups
356
+ # to the head shards
357
+ group_shard_settings = (
358
+ self.groups_ssm_state_size, # expected model size
359
+ (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned
360
+ n_groups == 1, # if there was only one group
361
+ )
362
+ intermediate_settings = (intermediate_size, 0, False)
363
+ head_settings = (self.num_heads, 0, False)
364
+
365
+ # - the weight already has a "weight_loader" attribute
366
+ # which set_weight_attrs will raise if we do not
367
+ # delete before trying to override it
368
+ # - ditto for the other two weights below
369
+ delattr(self.conv1d.bias, "weight_loader")
370
+ set_weight_attrs(
371
+ self.conv1d.bias,
372
+ {
373
+ "weight_loader": mamba_v2_sharded_weight_loader(
374
+ [
375
+ intermediate_settings,
376
+ group_shard_settings,
377
+ group_shard_settings,
378
+ ],
379
+ self.tp_size,
380
+ tp_rank,
381
+ )
382
+ },
383
+ )
384
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
- )
385
+ delattr(self.conv1d.weight, "weight_loader")
386
+ set_weight_attrs(
387
+ self.conv1d.weight,
388
+ {
389
+ "weight_loader": mamba_v2_sharded_weight_loader(
390
+ [
391
+ intermediate_settings,
392
+ group_shard_settings,
393
+ group_shard_settings,
394
+ ],
395
+ self.tp_size,
396
+ tp_rank,
397
+ )
398
+ },
399
+ )
400
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})
401
+ # Create the custom weight loader for Mamba sharding with group
402
+ # replication. This handles the interleaved projections correctly.
403
+ mamba_loader = mamba_v2_sharded_weight_loader(
404
+ [
405
+ intermediate_settings, # for gate
406
+ intermediate_settings,
407
+ group_shard_settings,
408
+ group_shard_settings,
409
+ head_settings, # for dt
410
+ ],
411
+ self.tp_size,
412
+ tp_rank,
413
+ )
414
+
415
+ # Apply the custom weight loader to in_proj.weight
416
+ # Works for both non-quantized (Parameter) and quantized
417
+ # (ModelWeightParameter which extends BasevLLMParameter)
418
+ if isinstance(self.in_proj.weight, BasevLLMParameter):
419
+ # For BasevLLMParameter subclasses (quantized layers like FP8)
420
+ # These have a weight_loader property that can be directly set
421
+ self.in_proj.weight.weight_loader = mamba_loader
422
+ else:
423
+ # For standard Parameter (non-quantized layers)
424
+ delattr(self.in_proj.weight, "weight_loader")
425
+ set_weight_attrs(self.in_proj.weight, {"weight_loader": mamba_loader})
426
427
# unsqueeze to fit conv1d weights shape into the linear weights shape.
428
# Can't do this in `weight_loader` since it already exists in