@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ff08e51940

Fork 0

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 14, 22:35 UTC

[NVIDIA] Fix Llama4 Scout FP4 functionality issues (#21499)

Signed-off-by: Po-Han Huang <pohanh@nvidia.com>

Po-Han Huang (NVIDIA) committed Jul 30, 2025 at 22:33 UTC ff08e51940a77d2dd14a6c512bec4613d060b4fa

3 files changed +219 -70

vllm/model_executor/layers/fused_moe/layer.py

+14 -1

@@ -874,6 +874,14 @@ class FusedMoE(torch.nn.Module):

874 elif shard_id == "w2":

875 param_data[expert_id] = loaded_weight

876

877 + def _load_w13_weight_scale(self, shard_dim: int,

878 + loaded_weight: torch.Tensor,

879 + param: torch.Tensor, tp_rank: int):

880 + shard_size = param.shape[shard_dim]

881 + loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,

882 + shard_size)

883 + param.copy_(loaded_weight)

884 +

885 def _load_model_weight_or_group_weight_scale(self,

886 shard_dim: int,

887 expert_data: torch.Tensor,

@@ -1123,7 +1131,12 @@ class FusedMoE(torch.nn.Module):

1131 "weight_scale_2" in weight_name if uses_weight_scale_2 else

1132 "weight_scale" in weight_name) or "input_scale" in weight_name

1133

1126 - if per_tensor_conditions:

1134 + if "w13_weight_scale" in weight_name:

1135 + self._load_w13_weight_scale(shard_dim=shard_dim,

1136 + loaded_weight=loaded_weight,

1137 + param=param,

1138 + tp_rank=self.tp_rank)

1139 + elif per_tensor_conditions:

1140 self._load_per_tensor_weight_scale(

1141 shard_id=shard_id,

1142 param=param,

vllm/model_executor/layers/quantization/modelopt.py

-2

@@ -778,8 +778,6 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

778 # Swizzle the weight blockscale.

779 # contracting dimension is input dimension

780 # block_size = 16;

781 - assert (layer.weight_scale.shape[1] % 16 == 0), (

782 - "Expected weight_scale.dim(1) to be divisible by 16")

781 assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (

782 "Weight Block scale must be represented as FP8-E4M3")

783 swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)

vllm/model_executor/models/llama4.py

+205 -67

@@ -342,34 +342,94 @@ class Llama4Model(LlamaModel):

342 expert_params_mapping: list[tuple[str, str, int, str]],

343 fused: bool = True,

344 ) -> bool:

345 + """

346 + Load MoE expert weights.

347 +

348 + Args:

349 + name: The name of the weight to load.

350 + loaded_weight: The weight to load.

351 + params_dict: The dictionary of module parameters.

352 + loaded_params: The set of already loaded parameters.

353 + expert_params_mapping: The mapping of expert parameters. Must be

354 + generated by FusedMoE.make_expert_params_mapping().

355 + fused: Whether the expert weights are fused into a single weight

356 + tensor or are separate weight tensors for each expert.

357 + When fused is True, loaded_weight should have shape of:

358 + [num_experts, hidden_in, hidden_out] for gate/up/down proj and

359 + [hidden_out, hidden_in] for the others like router.

360 + When fused is False, loaded_weight should have shape of:

361 + [hidden_out, hidden_in].

362 +

363 + Returns:

364 + True if loaded_weight is one of MoE weights and the MoE expert

365 + weights are loaded successfully, False otherwise.

366 + """

367 +

368 + # Whether the MoE expert weights are loaded successfully.

369 expert_param_loaded = False

346 - if "experts.gate_up_proj" in name:

347 - loaded_weight = loaded_weight.chunk(2, dim=-1)

370 +

371 + # If fused is True, the loaded weight is in the layout of:

372 + # [num_experts, hidden_in, hidden_out], so we must transpose the last

373 + # two dimensions to match the expected layout of the parameters.

374 + if fused and loaded_weight.ndim == 3:

375 + loaded_weight = loaded_weight.transpose(-1, -2)

376 +

377 + # If the gate_proj and up_proj weights are fused into a single

378 + # weight tensor, we need to split the weight tensor into a tuple

379 + # of two weight tensors along the hidden_out dimension.

380 + if "experts.gate_up_proj" in name:

381 + loaded_weight = loaded_weight.chunk(2, dim=-2)

382 +

383 + # Iterate over all the expert parameters and load the weights if we find

384 + # a match in weight name.

385 for (param_name, weight_name, expert_id,

386 shard_id) in expert_params_mapping:

387 +

388 + # Get a view of the loaded_weight to avoid modifying the original

389 + # one across iterations.

390 new_loaded_weight = loaded_weight

391 +

392 + # If expert weights are fused into a single weight tensor, remove

393 + # the expert index from the expected weight name.

394 if fused:

395 + # The string between e_str and proj_str is the expert index.

396 e_str, _, proj_str, _ = weight_name.split('.')

397 weight_name = f"{e_str}.{proj_str}"

398 param_name = f"{param_name}weight"

399 +

400 + # Skip if the current weight is not one of the MoE weights.

401 if weight_name not in name:

402 continue

403 +

404 + # Replace the weight name with the parameter name.

405 full_param_name = name.replace(weight_name, param_name)

358 - # Skip layers on other devices.

406 +

407 + # Skip if the current weight corresponds to a parameter that

408 + # does not exist on the current PP (pipeline parallel) rank.

409 if is_pp_missing_parameter(name, self):

410 continue

411 +

412 + # Skip if the current weight is for the bias.

413 if ((name.endswith(".bias") or name.endswith("_bias"))

414 and name not in params_dict):

415 continue

416 +

417 param = params_dict[full_param_name]

418 weight_loader = param.weight_loader

419 +

420 if fused:

421 + # If the parameter is for w13 together, the corresponding weight

422 + # will be a tuple, so we must select the correct weight

423 + # depending on the shard id, which is either "w1" or "w3".

424 if "w13" in full_param_name:

425 + assert shard_id in ["w1", "w3"]

426 shard_idx = 0 if shard_id == "w1" else 1

427 new_loaded_weight = new_loaded_weight[shard_idx]

370 - new_loaded_weight = new_loaded_weight.transpose(-1, -2)

428 +

429 + # If EP (expert parallel) is enabled, update expert_id to the

430 + # starting expert index for the current EP rank and extract the

431 + # corresponding expert weights.

432 layer_idx = extract_layer_index(name)

372 - # EP mapping

433 expert_map = self.layers[

434 layer_idx].feed_forward.experts.expert_map

435 if expert_map is not None:

@@ -382,6 +442,9 @@ class Llama4Model(LlamaModel):

442 else:

443 # TODO: add EP support for non fused weights

444 pass

445 +

446 + # Load the weight into the module parameter with corresponding

447 + # shard id and expert id.

448 weight_loader(param,

449 new_loaded_weight,

450 full_param_name,

@@ -390,10 +453,13 @@ class Llama4Model(LlamaModel):

453

454 loaded_params.add(full_param_name)

455 expert_param_loaded = True

456 +

457 return expert_param_loaded

458

459 def load_weights(self, weights: Iterable[tuple[str,

460 torch.Tensor]]) -> set[str]:

461 + # Name mapping from the parameter name to the shard name and

462 + # corresponding shard id.

463 stacked_params_mapping = [

464 # (param_name, shard_name, shard_id)

465 (".qkv_proj", ".q_proj", "q"),

@@ -402,26 +468,43 @@ class Llama4Model(LlamaModel):

468 (".gate_up_proj", ".gate_proj", 0),

469 (".gate_up_proj", ".up_proj", 1),

470 ]

471 + # Indicate whether the expert weights are fused into a single weight

472 + # tensor.

473 fused_experts_params = False

474 + # Expert parameter mapping for the case where the expert weights are

475 + # not fused into a single weight tensor.

476 expert_params_mapping = FusedMoE.make_expert_params_mapping(

477 ckpt_gate_proj_name="gate_proj",

478 ckpt_down_proj_name="down_proj",

479 ckpt_up_proj_name="up_proj",

480 num_experts=self.num_experts)

481 + # Expert parameter mapping for the case where the expert weights are

482 + # fused into a single weight tensor.

483 expert_params_mapping_fused = FusedMoE.make_expert_params_mapping(

484 ckpt_gate_proj_name="gate_up_proj",

485 ckpt_down_proj_name="down_proj",

486 ckpt_up_proj_name="gate_up_proj",

487 num_experts=1)

488 + # All the module parameters.

489 params_dict = dict(self.named_parameters())

490 + # The module parameters that have been loaded.

491 loaded_params: set[str] = set()

492 +

493 + # Iterate over all the weights and load them into module parameters.

494 for name, loaded_weight in weights:

495 +

496 + # If the name contains "experts.gate_up_proj" or "experts.down_proj"

497 + # without the expert indices, it means the expert weights are fused

498 + # into a single weight tensor across all experts.

499 if "experts.gate_up_proj" in name or "experts.down_proj" in name:

500 fused_experts_params = True

501 expert_params_mapping = expert_params_mapping_fused

502 +

503 + # If kv cache quantization scales exist and the weight name

504 + # corresponds to one of the kv cache quantization scales, load

505 + # them.

506 if (self.quant_config is not None and

507 (scale_name := self.quant_config.get_cache_scale(name))):

424 - # Loading kv cache quantization scales

508 param = params_dict[scale_name]

509 weight_loader = getattr(param, "weight_loader",

510 default_weight_loader)

@@ -430,84 +513,119 @@ class Llama4Model(LlamaModel):

513 weight_loader(param, loaded_weight)

514 loaded_params.add(scale_name)

515 continue

516 +

517 + # Iterate over stacked_params_mapping to check if the current weight

518 + # is one of the stacked parameters. If so, load the weight with the

519 + # corresponding shard id. Note that MoE weights are handled

520 + # separately in the else block.

521 for param_name, weight_name, shard_id in stacked_params_mapping:

522 + # Skip if the current weight is not one of the stacked

523 + # parameters or if the current weight is a MoE weight.

524 if weight_name not in name or "experts" in name:

525 continue

436 - # This check is for ModelOpt ckpts with kv cache quant enabled

526 +

527 + # For ModelOpt checkpoints, we need to rename the self_attn

528 + # weight/weight_scale names except for kv cache scales.

529 if not (name.endswith(

530 (".k_scale", ".v_scale")) and "self_attn" in name):

531 name = name.replace(weight_name, param_name)

532 +

533 + # Skip if the current weight corresponds to a parameter that

534 + # does not exist on the current PP (pipeline parallel) rank.

535 if is_pp_missing_parameter(name, self):

536 continue

442 - if name.endswith("scale") and "expert" not in name:

443 - # Remapping the name of FP8 kv-scale.

537 +

538 + # Remap kv cache scale names for ModelOpt checkpoints.

539 + # TODO: ModelOpt should implement get_cache_scale() such that

540 + # kv cache scale name remapping can be done there.

541 + if name.endswith("scale"):

542 name = maybe_remap_kv_scale_name(name, params_dict)

543 if name is None:

544 continue

545 +

546 + # Load the weight into the module parameter with corresponding

547 + # shard id and exit the for loop and the else block.

548 param = params_dict[name]

549 weight_loader = getattr(param, "weight_loader",

550 default_weight_loader)

551 +

552 if weight_loader == default_weight_loader:

553 weight_loader(param, loaded_weight)

554 else:

555 weight_loader(param, loaded_weight, shard_id)

556 +

557 loaded_params.add(name)

558 break

559 +

560 + # Handle normal (non-stacked) weights and MoE weights.

561 else:

457 - moe_loaded = self.load_moe_expert_weights(

458 - name,

459 - loaded_weight,

460 - params_dict,

461 - loaded_params,

462 - expert_params_mapping,

463 - fused=fused_experts_params)

464 -

465 - if not moe_loaded:

466 - if is_pp_missing_parameter(name, self):

467 - continue

562 + # First, try to load MoE weights using load_moe_expert_weights.

563 + # If successful, move on to next loaded weight.

564 + if self.load_moe_expert_weights(name,

565 + loaded_weight,

566 + params_dict,

567 + loaded_params,

568 + expert_params_mapping,

569 + fused=fused_experts_params):

570 + continue

571

469 - # Handle flat expert scale parameters that

470 - # don't match per-expert patterns

471 - if ("experts." in name and ("w13_input_scale" in name

472 - or "w13_weight_scale" in name

473 - or "w2_input_scale" in name

474 - or "w2_weight_scale" in name)):

475 - # These are flat expert scales that apply to all experts

476 - param = params_dict[name]

477 - weight_loader = getattr(param, "weight_loader",

478 - default_weight_loader)

479 -

480 - # Check for MoE-specific loading support via

481 - # attribute instead of expensive runtime reflection

482 - supports_moe = getattr(weight_loader,

483 - 'supports_moe_loading', False)

484 -

485 - if supports_moe:

486 - # This is a MoE weight loader

487 - if "w13_" in name:

488 - shard_id = "w1"

489 - elif "w2_" in name:

490 - shard_id = "w2"

491 - else:

492 - shard_id = "w1"

493 -

494 - weight_loader(param,

495 - loaded_weight,

496 - name,

497 - shard_id=shard_id,

498 - expert_id=0)

499 - else:

500 - # Regular weight loader (handles both

501 - # param.weight_loader and default_weight_loader)

502 - weight_loader(param, loaded_weight)

503 - loaded_params.add(name)

504 - continue

572 + # Skip if the current weight corresponds to a parameter that

573 + # does not exist on the current PP (pipeline parallel) rank.

574 + if is_pp_missing_parameter(name, self):

575 + continue

576 +

577 + # Handle flat expert scale parameters that don't match

578 + # per-expert patterns, i.e. one weight scale tensor for all

579 + # experts.

580 + scale_names = [

581 + "w13_input_scale", "w13_weight_scale", "w2_input_scale",

582 + "w2_weight_scale"

583 + ]

584 + if ("experts." in name and any(scale_name in name

585 + for scale_name in scale_names)):

586

587 param = params_dict[name]

588 weight_loader = getattr(param, "weight_loader",

589 default_weight_loader)

509 - weight_loader(param, loaded_weight)

590 +

591 + # If weight loader supports special moe loading, use it to

592 + # avoid expensive runtime reflection

593 + if getattr(weight_loader, 'supports_moe_loading', False):

594 + # Map the weight name to the corresponding shard id.

595 + shard_id = "w2" if "w2_" in name else "w1"

596 +

597 + # Transpose if weight scales are FP8 block scales with

598 + # three dimensions:

599 + # [num_experts, hidden_in, hidden_out].

600 + if name.endswith("weight_scale") \

601 + and loaded_weight.dtype == torch.float8_e4m3fn \

602 + and loaded_weight.ndim == 3:

603 + loaded_weight = loaded_weight.transpose(-1, -2)

604 +

605 + # Load the weight into the module parameter with

606 + # corresponding shard id and expert id.

607 + weight_loader(param,

608 + loaded_weight,

609 + name,

610 + shard_id=shard_id,

611 + expert_id=0)

612 +

613 + else:

614 + # Regular weight loader (handles both

615 + # param.weight_loader and default_weight_loader)

616 + weight_loader(param, loaded_weight)

617 +

618 loaded_params.add(name)

619 + continue

620 +

621 + # Handle normal (non-stacked, non-MoE) weights.

622 + param = params_dict[name]

623 + weight_loader = getattr(param, "weight_loader",

624 + default_weight_loader)

625 + weight_loader(param, loaded_weight)

626 + loaded_params.add(name)

627 +

628 + # Finally, return the set of loaded parameters.

629 return loaded_params

630

631

@@ -560,23 +678,43 @@ class Llama4ForCausalLM(LlamaForCausalLM):

678 loaded_weight: torch.Tensor,

679 ) -> tuple[str, torch.Tensor]:

680

563 - def permute(w: torch.Tensor, n_heads: int):

681 + # Helper function to permute the weight's channels

682 + def permute(w: torch.Tensor, n_heads: int, is_weight_scale: bool):

683 +

684 + # Calculate the expected shape of the weight.

685 + # Do not rely on w's shape, as it may be in another layout.

686 attn_in = self.config.head_dim * n_heads

687 attn_out = self.config.hidden_size

688

689 + # If the weight is FP4 packed as uint8, we need to divide attn_out

690 + # by 2.

691 + if w.dtype == torch.uint8 and w.shape[1] * 2 == attn_out:

692 + attn_out = attn_out // 2

693 +

694 + # If the weight is a weight scale, we need to divide attn_out by

695 + # block size, which is currently 16.

696 + elif w.dtype == torch.float8_e4m3fn and is_weight_scale \

697 + and w.shape[1] * 16 == attn_out:

698 + attn_out = attn_out // 16

699 +

700 return w.view(n_heads, attn_in // n_heads // 2, 2,

701 attn_out).transpose(1, 2).reshape(attn_in, attn_out)

702

703 modules = name.split(".")

704

572 - # rotary embeds should be sliced

573 - if ("wk" in modules or "k_proj" in modules) \

574 - and modules[-1] == "weight":

575 - loaded_weight = permute(loaded_weight,

576 - self.config.num_key_value_heads)

577 - elif ("wq" in modules or "q_proj" in modules) \

578 - and modules[-1] == "weight":

579 - loaded_weight = permute(loaded_weight,

580 - self.config.num_attention_heads)

705 + # Permute Q/K weights and weight block scales for rotary embedding

706 + is_weight = modules[-1] == "weight"

707 + is_nvfp4_weight_scale = (modules[-1] == "weight_scale" and

708 + loaded_weight.dtype == torch.float8_e4m3fn)

709 +

710 + if is_weight or is_nvfp4_weight_scale:

711 + if ("wk" in modules or "k_proj" in modules):

712 + loaded_weight = permute(loaded_weight,

713 + self.config.num_key_value_heads,

714 + is_nvfp4_weight_scale)

715 + elif ("wq" in modules or "q_proj" in modules):

716 + loaded_weight = permute(loaded_weight,

717 + self.config.num_attention_heads,

718 + is_nvfp4_weight_scale)

719

720 return name, loaded_weight