From ae5e98cf86ee339bda494151621beb53ae3b85b7 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 21 Jun 2026 16:08:46 +0000 Subject: [PATCH] B300 NVFP4 port: advance layers/quantization/utils.py to HEAD (piecemeal-gap fix) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../sglang/srt/layers/quantization/utils.py | 146 +++++++++++------- 1 file changed, 93 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/layers/quantization/utils.py b/python/sglang/srt/layers/quantization/utils.py index 198b201de..6c5aed185 100644 --- a/python/sglang/srt/layers/quantization/utils.py +++ b/python/sglang/srt/layers/quantization/utils.py @@ -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,