chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

B300 NVFP4 port: advance layers/quantization/utils.py to HEAD (piecemeal-gap fix)

Browse Source 
Ported flashinfer_trtllm.py (Stage 1) calls prepare_static_weights_for_trtllm_fp4_moe(is_gated=...),
but utils.py (which DEFINES it) was not ported → old signature → TypeError at process_weights.
utils.py has 0 local edits; wholesale-replace to HEAD. New fn is gated-aware (gemm1_rows), passes
is_gated_act_gemm to flashinfer _maybe_get_cached_w3_w1_permute_indices (verified present on 0.6.12).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-21 16:08:46 +00:00
parent 2aca99d677
commit ae5e98cf86
1 changed files with 93 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

146
python/sglang/srt/layers/quantization/utils.py
View File

@@ -1,3 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/quant_utils.py
from __future__ import annotations
@@ -43,6 +45,28 @@ def get_scalar_types():
ScalarType, scalar_types = get_scalar_types()
def _module_path_match(ignored: str, prefix: str) -> bool:
# Match on dotted module-path boundaries so that `mlp.gate` does NOT
# match `mlp.gate_up_proj`. Needed for quant configs (e.g. Qwen3.6-FP8)
# whose `modules_to_not_convert` lists MoE-template names like `mlp.gate`
# that collide with fused dense MLP names by plain substring.
ignored = ignored.rstrip(".")
prefix = prefix.rstrip(".")
if ignored == prefix:
return True
if prefix.startswith(ignored + "."):
return True
return ("." + ignored + ".") in ("." + prefix + ".")
# Known fused-linear -> shard names. Used as a fallback when the quant
# config doesn't ship packed_modules_mapping (typical for HF FP8 configs).
_FALLBACK_FUSED_SHARDS: Mapping[str, List[str]] = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
def is_layer_skipped(
prefix: str,
ignored_layers: List[str],
@@ -56,16 +80,19 @@ def is_layer_skipped(
# in the safetensors checkpoint. So, we convert the name
# from the fused version to unfused + check to make sure that
# each shard of the fused layer has the same scheme.
if proj_name in fused_mapping:
effective_fused = (
fused_mapping if proj_name in fused_mapping else _FALLBACK_FUSED_SHARDS
)
if proj_name in effective_fused:
shard_prefixes = [
prefix.replace(proj_name, shard_proj_name)
for shard_proj_name in fused_mapping[proj_name]
for shard_proj_name in effective_fused[proj_name]
]
is_skipped = None
for shard_prefix in shard_prefixes:
is_shard_skipped = any(
ignored in shard_prefix for ignored in ignored_layers
_module_path_match(ignored, shard_prefix) for ignored in ignored_layers
)
if is_skipped is None:
@@ -77,7 +104,9 @@ def is_layer_skipped(
"to have the same precision."
)
else:
is_skipped = any(ignored in prefix for ignored in ignored_layers)
is_skipped = any(
_module_path_match(ignored, prefix) for ignored in ignored_layers
)
if "gate_up_proj" in prefix:
prefix_gate = prefix.replace("gate_up_proj", "gate_proj")
prefix_up = prefix.replace("gate_up_proj", "up_proj")
@@ -625,6 +654,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(
hidden_size,
intermediate_size,
num_experts,
is_gated: bool = True,
):
from flashinfer import nvfp4_block_scale_interleave
from flashinfer.fused_moe.core import (
@@ -636,14 +666,16 @@ def prepare_static_weights_for_trtllm_fp4_moe(
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
gemm1_rows = (2 if is_gated else 1) * intermediate_size
# Convert quantized weights to proper formats
gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
num_experts, 2 * intermediate_size, hidden_size // 2
num_experts, gemm1_rows, hidden_size // 2
) # packed fp4
gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
torch.float8_e4m3fn
).reshape(
num_experts, 2 * intermediate_size, hidden_size // 16
num_experts, gemm1_rows, hidden_size // 16
) # fp8 scaling factors
gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
@@ -655,52 +687,66 @@ def prepare_static_weights_for_trtllm_fp4_moe(
num_experts, hidden_size, intermediate_size // 16
) # fp8 scaling factors
gemm1_weights_fp4_shuffled = []
gemm1_scales_fp4_shuffled = []
gemm2_weights_fp4_shuffled = []
gemm2_scales_fp4_shuffled = []
# Pre-allocate output tensors so per-expert shuffles write directly into
# contiguous slices instead of building lists + torch.stack(). This avoids
# O(num_experts) transient GPU allocations whose freed blocks fragment the
# CUDA address space
gemm1_weights_fp4_shuffled = torch.empty_like(gemm1_weights_fp4.view(torch.uint8))
gemm2_weights_fp4_shuffled = torch.empty_like(gemm2_weights_fp4.view(torch.uint8))
# Pre-allocate scale output tensors and a reusable scratch buffer for
# the permuted input to nvfp4_block_scale_interleave.
# nvfp4_block_scale_interleave flattens its input to 1-D, so the
# per-expert output size equals the per-expert input numel.
def _alloc_scale_buffers(scales):
per_expert_shape = scales[0].view(torch.uint8).shape
per_expert_numel = scales[0].numel()
output = scales.new_empty((num_experts, per_expert_numel), dtype=torch.uint8)
scratch = torch.empty(per_expert_shape, dtype=torch.uint8, device=scales.device)
return output, scratch
gemm1_scales_fp4_shuffled, g1s_scratch = _alloc_scale_buffers(
gemm1_scales_linear_fp4
)
gemm2_scales_fp4_shuffled, g2s_scratch = _alloc_scale_buffers(
gemm2_scales_linear_fp4
)
for i in range(num_experts):
# Calculate the permute indices for the following:
# 1. Reorder rows of W1 and scales for fused gated activation
# 2. Shuffle weights and scaling factors for transposed mma output
# for both w3_w1 and w2 weights and scale factors
permute_indices = _maybe_get_cached_w3_w1_permute_indices(
_cache_permute_indices,
gemm1_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
is_gated_act_gemm=is_gated,
)
gemm1_weights_fp4_shuffled.append(
gemm1_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
.contiguous()
)
gemm1_weights_fp4_shuffled[i] = gemm1_weights_fp4[i].view(torch.uint8)[
permute_indices.to(gemm1_weights_fp4.device)
]
permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
_cache_permute_indices,
gemm1_scales_linear_fp4[i].view(torch.uint8),
epilogue_tile_m,
num_elts_per_sf=16,
is_gated_act_gemm=is_gated,
)
gemm1_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm1_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm1_scales_linear_fp4.device)
]
.contiguous()
)
# Reuse scratch buffer for the permuted scale input
torch.index_select(
gemm1_scales_linear_fp4[i].view(torch.uint8),
0,
permute_sf_indices.to(gemm1_scales_linear_fp4.device),
out=g1s_scratch,
)
gemm1_scales_fp4_shuffled[i] = nvfp4_block_scale_interleave(g1s_scratch)
permute_indices = get_w2_permute_indices_with_cache(
_cache_permute_indices,
gemm2_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
)
gemm2_weights_fp4_shuffled.append(
gemm2_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
.contiguous()
)
gemm2_weights_fp4_shuffled[i] = gemm2_weights_fp4[i].view(torch.uint8)[
permute_indices.to(gemm2_weights_fp4.device)
]
permute_sf_indices = get_w2_permute_indices_with_cache(
_cache_permute_indices,
@@ -708,30 +754,24 @@ def prepare_static_weights_for_trtllm_fp4_moe(
epilogue_tile_m,
num_elts_per_sf=16,
)
gemm2_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm2_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm2_scales_linear_fp4.device)
]
.contiguous()
)
torch.index_select(
gemm2_scales_linear_fp4[i].view(torch.uint8),
0,
permute_sf_indices.to(gemm2_scales_linear_fp4.device),
out=g2s_scratch,
)
gemm2_scales_fp4_shuffled[i] = nvfp4_block_scale_interleave(g2s_scratch)
# Stack weights for all experts
gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
gemm1_scales_fp4_shuffled = (
torch.stack(gemm1_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
)
del g1s_scratch, g2s_scratch
gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
gemm2_scales_fp4_shuffled = (
torch.stack(gemm2_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, hidden_size, intermediate_size // 16)
)
# Weight outputs stay as uint8 (FP4 packed) — the TRTLLM kernel expects this.
gemm1_scales_fp4_shuffled = gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn
).reshape(num_experts, gemm1_rows, hidden_size // 16)
gemm2_scales_fp4_shuffled = gemm2_scales_fp4_shuffled.view(
torch.float8_e4m3fn
).reshape(num_experts, hidden_size, intermediate_size // 16)
return (
gemm1_weights_fp4_shuffled,
gemm1_scales_fp4_shuffled,