diff --git a/vllm/envs.py b/vllm/envs.py index 5950757e166b..2669ff61ab91 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -118,6 +118,7 @@ VLLM_MOE_AWQ_GEMV_HIP: bool = False VLLM_MOE_GPTQ_EXLLAMA: bool = False VLLM_MOE_HYBRID_W4A16: bool = False + VLLM_W4A16_PREFILL_DEQUANT: bool = False VLLM_ROCM_USE_MOE_WNA16_CUDA_KERNEL: bool = False VLLM_ROCM_USE_AITER: bool = False VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False @@ -1112,6 +1113,21 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_MOE_HYBRID_W4A16": lambda: ( os.getenv("VLLM_MOE_HYBRID_W4A16", "true").lower() in ("true", "1") ), + # ROCm W4A16 linear: cache a dequantized copy of each weight (in the model's + # activation dtype, fp16 or bf16) at load time so the prefill GEMM (batch + # M > skinny threshold) runs a dense hipBLASLt GEMM instead of the fused + # in-kernel int4 unpack. Decode still uses the int4 weights. Trades VRAM + # (~4x the int4 weight) for prefill compute. The copy is only kept while + # there is room left in the gpu_memory_utilization budget (budget = total * + # util): a weight is cached if (free_mem - total * (1 - util)) >= + # dequant_bytes at load time, otherwise it keeps the int4 prefill path. The + # check is self-limiting -- each copy shrinks the spendable budget, so later + # layers stop being cached once the budget is exhausted (a message is + # emitted). Raise gpu_memory_utilization to cache more weights at the cost of + # KV-cache memory. + "VLLM_W4A16_PREFILL_DEQUANT": lambda: ( + os.getenv("VLLM_W4A16_PREFILL_DEQUANT", "false").lower() in ("true", "1") + ), # Use exllama 4-bit kernel for MoE GPTQ instead of Triton. # Requires exllama-native weight format [E, K/8, N] int32. "VLLM_MOE_GPTQ_EXLLAMA": lambda: ( diff --git a/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py index 987ba349750f..e34059b9d1e5 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py @@ -16,6 +16,8 @@ import torch +import vllm.envs as envs +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( unpack_quantized_values_into_int32, ) @@ -30,6 +32,8 @@ from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig +logger = init_logger(__name__) + SUPPORTED_GROUP_SIZES = [32, 64, 128] # Maximum batch size M for the HIP skinny kernel path (C++ supports N_in @@ -606,6 +610,7 @@ def _hybrid_w4a16_apply_impl( cu_count: int, group_size: int, packed_scale_zp: torch.Tensor | None = None, + w_dequant: torch.Tensor | None = None, ) -> torch.Tensor: """Dispatch between skinny GEMM and Triton based on batch size M. @@ -618,6 +623,11 @@ def _hybrid_w4a16_apply_impl( single format: dequant = (nibble - zp_raw) * scale. packed_scale_zp: [N, K//G] fp32 carrier packing scale + zero-point per group (Triton prefill, asymmetric only), or None for symmetric. + w_dequant: [N, K] dense dequantized weight (VLLM_W4A16_PREFILL_DEQUANT), or None. + When present, prefill (M > MAX_SKINNY_BATCH_SIZE) runs a dense GEMM + on it. Passed as an op arg (not branched on in apply_weights) so the + M-branch stays out of the compiled graph -- decode (small M) replays + the same int4 cudagraph whether or not the dequantized copy exists. Registered as a custom op so torch.compile treats it as opaque. """ @@ -638,6 +648,16 @@ def _hybrid_w4a16_apply_impl( with ctx: return ops.wvSplitK_int4_g(w_q, x_2d, w_s, cu_count, group_size, w_zp, bias) + # Prefill with the pre-dequantized dense copy (load-time cached), if present. + if w_dequant is not None: + ctx = ( + nullcontext() + if torch.compiler.is_compiling() + else torch.profiler.record_function(f"hybrid_dequant_w4a16 {M}x{N}x{K}") + ) + with ctx: + return torch.nn.functional.linear(x_2d, w_dequant, bias) + ctx = ( nullcontext() if torch.compiler.is_compiling() @@ -669,6 +689,7 @@ def _hybrid_w4a16_apply_fake( cu_count: int, group_size: int, packed_scale_zp: torch.Tensor | None = None, + w_dequant: torch.Tensor | None = None, ) -> torch.Tensor: M = x_2d.size(0) N = w_q.size(0) @@ -769,6 +790,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w_s_skinny = w_s_raw.data.contiguous() # ---- Process zero-points for asymmetric quantization ---- + w_zp = None if c.zero_points: assert self.w_zp_name is not None w_zp_raw = getattr(layer, self.w_zp_name) @@ -807,6 +829,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # plain integer. Consumed by the int-domain subtract (RDNA3 has no # v_pk_fma_bf16). Bit-identical to the separate scale+zp loads. if c.zero_points and c.act_type in (torch.float16, torch.bfloat16): + assert w_zp is not None # set above whenever c.zero_points is True scale_u16 = w_s_skinny.view(torch.uint16).to(torch.int32) & 0xFFFF if c.act_type == torch.float16: w_s_f32 = w_s_skinny.to(torch.float32) @@ -824,6 +847,83 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(packed_scale_zp, requires_grad=False), ) + # ---- Optional: cache a dequantized copy of the weight (in the model's + # activation dtype, fp16 or bf16) so the + # prefill path (M > MAX_SKINNY_BATCH_SIZE) can run a dense hipBLASLt GEMM + # and skip the in-kernel int4 unpack. Decode still uses the int4 weights + # (bandwidth-bound at small M). The copy costs ~4x the int4 weight, so it + # is gated per-weight on the gpu_memory_utilization budget: weights are + # cached greedily until the budget is exhausted, after which the rest keep + # the int4 prefill path. Allocating here (before the memory profiler runs) + # lets the KV-cache sizing account for the copies. + if envs.VLLM_W4A16_PREFILL_DEQUANT and unpacked.device.type == "cuda": + self._maybe_cache_dequant_prefill_weight(layer, unpacked, w_s_skinny, w_zp) + + def _maybe_cache_dequant_prefill_weight( + self, + layer: torch.nn.Module, + unpacked: torch.Tensor, + w_s_skinny: torch.Tensor, + w_zp: torch.Tensor | None, + ) -> None: + """Cache a dense dequantized copy of the weight (in the activation dtype, + fp16 or bf16) for the prefill GEMM, when it still fits in vLLM's budget. + + The copy is registered as ``_hybrid_w_dequant`` and consumed by the prefill + branch of ``apply_weights``. The gate is anchored to the configured + ``gpu_memory_utilization`` budget rather than a raw free-VRAM number: + + budget = total * gpu_memory_utilization (weights + acts + KV) + spendable = free_now - total * (1 - util) (room left in budget) + + We cache the copy only while ``spendable`` clears it. The check uses live + free memory, so it is self-limiting: as copies are allocated, + ``spendable`` shrinks and later layers stop being cached once the budget + is exhausted. Skipped weights keep the int4 prefill path (warns once). + """ + from vllm.config import get_current_vllm_config + from vllm.utils.mem_utils import MemorySnapshot + + c = self.config + N, K = unpacked.shape + dequant_bytes = N * K * torch.empty((), dtype=c.act_type).element_size() + + util = get_current_vllm_config().cache_config.gpu_memory_utilization + # MemorySnapshot mirrors vLLM's own KV-cache profiler: on integrated/UMA + # GPUs (e.g. gfx1151 Strix Halo) cudaMemGetInfo underreports free memory, + # so it falls back to psutil there -- keeping this gate consistent with + # how the KV-cache budget is actually sized. + snapshot = MemorySnapshot(device=unpacked.device) + free_bytes, total_bytes = snapshot.free_memory, snapshot.total_memory + # Memory vLLM deliberately leaves untouched outside its budget. + out_of_budget = total_bytes * (1.0 - util) + # Room still free within the budget right now. + spendable = free_bytes - out_of_budget + if spendable < dequant_bytes: + logger.warning_once( + "VLLM_W4A16_PREFILL_DEQUANT: no room left in the " + "gpu_memory_utilization=%.2f budget to cache a dequantized " + "prefill copy; affected W4A16 weights use the int4 prefill path. Raise " + "gpu_memory_utilization to cache more.", + util, + ) + return + + G = c.group_size + u = unpacked.to(torch.float32) # [N, K] natural order, nibble 0..15 + scale_exp = w_s_skinny.to(torch.float32).repeat_interleave(G, dim=1) + if w_zp is not None: + zp_exp = w_zp.to(torch.float32).repeat_interleave(G, dim=1) + w_dequant = ((u - zp_exp) * scale_exp).to(c.act_type) + else: + w_dequant = ((u - 8.0) * scale_exp).to(c.act_type) + # Buffer (not a Parameter): this is a derived, recomputable cache, so it + # should move with the module but stay out of .parameters() and, with + # persistent=False, out of state_dict(). + layer.register_buffer( + "_hybrid_w_dequant", w_dequant.contiguous(), persistent=False + ) + def apply_weights( self, layer: torch.nn.Module, @@ -842,6 +942,14 @@ def apply_weights( N = w_q.shape[0] out_shape = x.shape[:-1] + (N,) + # Dequantized copy cached at load time (opt-in via + # VLLM_W4A16_PREFILL_DEQUANT, subject to the free-VRAM gate), or None when + # absent/skipped. Passed straight to the op; the M-branch (prefill -> + # dense dequant, decode -> int4) lives INSIDE the opaque op so it never + # enters the compiled graph -- decode replays the same int4 cudagraph + # whether or not the dequantized copy exists. + w_dequant = getattr(layer, "_hybrid_w_dequant", None) + cu_count = num_compute_units() output = torch.ops.vllm.hybrid_w4a16_apply( x_2d, @@ -853,5 +961,6 @@ def apply_weights( cu_count, c.group_size, packed_scale_zp, + w_dequant, ) return output.reshape(out_shape)