diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 472c19e74a..ba054b3e30 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -727,7 +727,7 @@ steps: - tests/quantization/reference_mxfp4.py - tests/quantization/test_auto_round.py commands: - - uv pip install --system b12x==1.3.0 + - uv pip install --system b12x==1.5.1 - pytest -v -s model_executor/kernels/test_b12x_linear.py model_executor/test_b12x_warmup.py kernels/quantization/test_block_fp8.py diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index b5f56d71cb..e9c2dec5d8 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -17,7 +17,7 @@ vllm serve \ --moe-backend b12x ``` -Only pass `--moe-backend b12x` for a compatible NVFP4 or MXFP4 MoE model. The +Only pass `--moe-backend b12x` for a compatible packed-weight MoE model. The linear and MoE backends can be selected independently. b12x uses MXFP8 activations by default for MXFP4 MoE and the checkpoint's @@ -29,9 +29,34 @@ force BF16 activations for either FP4 weight format. | Backend | Supported configurations | | ------- | ------------------------ | -| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 | -| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations | +| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4 (including W4A16), MXFP4, IQ2_XS, IQ2_XXS, and Q8_0 | +| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations; IQ2_XS, IQ2_XXS, and Q8_0 weights with BF16 activations | +| Embedding | Packed Q8_0 lookup with BF16 output | -Dense W4A16 layers are not handled by b12x and continue to use another -compatible backend such as Marlin. The b12x MoE backend does not support expert -parallelism, expert maps, EXL3, or NF3. +The b12x MoE backend does not support expert parallelism, expert maps, EXL3, or NF3. + +## ModelOpt block codecs + +ModelOpt mixed-precision checkpoints select block codecs through the +`quantization_config.quantized_layers` entries in `config.json`. Each entry +declares `quant_algo`, `group_size`, `block_payload_bytes`, and `packing="ggml"`: + +| `quant_algo` | `group_size` | `block_payload_bytes` | +| --- | --- | --- | +| `IQ2_XS` | 256 | 74 | +| `IQ2_XXS` | 256 | 66 | +| `Q8_0` | 32 | 34 | + +Serialized matrices have uint8 shape `[N, K/group_size, block_payload_bytes]`. +Dense preparation rearranges the encoded blocks without expanding the matrix +to BF16. Routed experts use bias-free SiLU or non-gated ReLU²; tensor-parallel +shards must preserve whole blocks and satisfy the kernel's alignment constraints. +Only Q8_0 supports packed embedding lookup. + +Quantized normalization weights and other packed vectors are not supported. +Store normalization weights as ordinary floating-point tensors and omit their +entries from `quantization_config.quantized_layers`. + +Dense, embedding, and expert matrices remain packed. Their execution uses BF16 +activations without activation quantization. Prepare the kernels and workspaces +before CUDA graph capture. diff --git a/setup.py b/setup.py index b173b9dca8..2a5b54f410 100644 --- a/setup.py +++ b/setup.py @@ -1706,7 +1706,7 @@ setup( # only; also needs system GStreamer + libv4l (see docs). "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], "flashinfer": [], # Kept for backwards compatibility - "b12x": ["b12x==1.3.0"], + "b12x": ["b12x==1.5.1"], # Optional deps for Helion kernel development # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml diff --git a/tests/kernels/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py index d0b1ce7cae..7fb83d5605 100644 --- a/tests/kernels/moe/test_b12x.py +++ b/tests/kernels/moe/test_b12x.py @@ -179,7 +179,14 @@ def _make_b12x_moe_kernel( num_experts=num_experts, experts_per_token=topk, hidden_dim=hidden_states.shape[1], - intermediate_size=w2.shape[2] * 2, + intermediate_size=w2.shape[2] + * ( + 32 + if quant_config.weight_quant_dtype == "q8_0" + else 256 + if quant_config.weight_quant_dtype in ("iq2_xs", "iq2_xxs") + else 2 + ), in_dtype=hidden_states.dtype, activation=activation, ) @@ -1272,3 +1279,78 @@ def test_b12x_moe_cuda_graph_replay( assert torch.isfinite(expected).all() assert torch.isfinite(actual).all() torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120") +@pytest.mark.parametrize("codec", ["iq2_xxs", "q8_0"]) +@pytest.mark.parametrize("tokens", [1, 2, 4]) +@torch.inference_mode() +def test_b12x_block_quant_relu2_moe_graph(codec, tokens, workspace_init): + from b12x.testing.iq2_xs_reference import moe_reference_iq2_xs + + from tests.model_executor.kernels.test_b12x_linear import _block_weights + + with set_current_vllm_config(VllmConfig()): + torch.manual_seed(71) + e, h, i, topk = 4, 256, 256, 2 + raw1, raw2 = _block_weights(e * i, h, codec), _block_weights(e * h, i, codec) + w1 = raw1.view(e, i, *raw1.shape[1:]).cuda() + w2 = raw2.view(e, h, *raw2.shape[1:]).cuda() + x = torch.randn(tokens, h, device="cuda", dtype=torch.bfloat16) * 0.125 + scores = torch.randn(tokens, e, device="cuda", dtype=torch.bfloat16) + from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantDesc + + quant = FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(dtype=codec), + _w2=FusedMoEQuantDesc(dtype=codec), + ) + activation = MoEActivation.RELU2_NO_MUL + kernel = _make_b12x_moe_kernel(x, w1, w2, topk, activation, quant) + weights, ids, _ = fused_topk(x, scores, topk, renormalize=False) + + def apply(): + return kernel.apply( + hidden_states=x, + w1=w1, + w2=w2, + topk_weights=weights, + topk_ids=ids, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ) + + apply() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = apply() + for _ in range(2): + x.normal_(std=0.125) + scores.normal_() + new_weights, new_ids, _ = fused_topk(x, scores, topk, renormalize=False) + weights.copy_(new_weights) + ids.copy_(new_ids) + allocated = torch.accelerator.memory_allocated() + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() == allocated + expected = moe_reference_iq2_xs( + x, + raw1.view(e, i, *raw1.shape[1:]), + raw2.view(e, h, *raw2.shape[1:]), + ids, + weights, + activation="relu2", + ) + assert torch.isfinite(output).all() + relative_l2 = ( + output.float() - expected.float() + ).norm() / expected.float().norm() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), expected.float().flatten(), dim=0 + ) + assert cosine >= 0.999 and relative_l2 <= 0.01 + graph.reset() diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index ac72cfb845..a63b340e62 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -21,6 +21,7 @@ from vllm.model_executor.kernels.linear import ( B12xMxFp4LinearKernel, B12xMxfp8LinearKernel, B12xNvFp4LinearKernel, + B12xNvFp4W4A16LinearKernel, B12xTensorFP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig, Mxfp8LinearLayerConfig, @@ -782,7 +783,7 @@ def test_b12x_nvfp4_can_implement_supported_config() -> None: assert reason is None -def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: +def test_b12x_backend_selects_w4a16(monkeypatch) -> None: import vllm.model_executor.kernels.linear as linear_mod monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) @@ -792,17 +793,17 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: monkeypatch.setitem( linear_mod._POSSIBLE_NVFP4_KERNELS, PlatformEnum.CUDA, - [B12xNvFp4LinearKernel, MarlinNvFp4LinearKernel], + [B12xNvFp4LinearKernel, B12xNvFp4W4A16LinearKernel, MarlinNvFp4LinearKernel], ) monkeypatch.setattr( - MarlinNvFp4LinearKernel, + B12xNvFp4W4A16LinearKernel, "is_supported", classmethod(lambda cls, compute_capability=None: (True, None)), ) kernel = init_nvfp4_linear_kernel(use_a16=True) - assert isinstance(kernel, MarlinNvFp4LinearKernel) + assert isinstance(kernel, B12xNvFp4W4A16LinearKernel) def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: @@ -859,3 +860,227 @@ def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: layer.alpha, ) assert kwargs == {"out_dtype": torch.bfloat16} + + +def test_b12x_nvfp4_w4a16_uses_weight_only_scale_and_a16_dispatch(monkeypatch): + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + calls = [] + packed_weight = object() + + def pack_weight(weight, scale, *, recipe, global_scale): + assert weight is layer.weight + assert scale is layer.weight_scale + assert recipe == "nvfp4" + assert global_scale is layer.weight_global_scale + return packed_weight + + def mm(source, weight, *, required_mode): + assert weight is packed_weight + assert required_mode == "a16" + assert source.dtype == torch.bfloat16 + assert source.is_contiguous() + calls.append(source) + return torch.full((source.shape[0], 48), 3.0, dtype=source.dtype) + + def reject_activation_quantization(*args, **kwargs): + pytest.fail("W4A16 must not quantize activations") + + monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", reject_activation_quantization) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(pack_weight=pack_weight, mm=mm), + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_intrinsics", + lambda: types.SimpleNamespace(swizzle_block_scale=lambda scale: scale), + ) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.empty((48, 64), dtype=torch.uint8), requires_grad=False + ) + layer.weight_scale = torch.nn.Parameter( + torch.empty((48, 8), dtype=torch.float8_e4m3fn), requires_grad=False + ) + layer.weight_global_scale = torch.tensor([0.125]) + kernel = object.__new__(B12xNvFp4W4A16LinearKernel) + kernel.process_weights_after_loading(layer) + assert layer.b12x_warmup_provider is kernel + kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16).compile() + assert [source.shape[0] for source in calls] == [1, 8] + + x = torch.randn((2, 3, 256), dtype=torch.bfloat16)[..., ::2] + output = kernel.apply_weights(layer, x, torch.ones(48, dtype=torch.bfloat16)) + torch.testing.assert_close(calls[-1], x.reshape(6, 128)) + torch.testing.assert_close(output, torch.full((2, 3, 48), 4.0, dtype=x.dtype)) + with pytest.raises(ValueError, match="requires BF16 input"): + kernel.apply_weights(layer, x.float()) + + +@pytest.mark.parametrize("m,n,k", [(2, 136, 96), (17, 256, 256)]) +@torch.inference_mode() +def test_b12x_nvfp4_w4a16_preserves_bf16_activations_under_graph_replay( + monkeypatch, default_vllm_config, m, n, k +) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + supported, reason = B12xNvFp4W4A16LinearKernel.is_supported() + if not supported: + pytest.skip(reason) + + def reject_activation_quantization(*args, **kwargs): + pytest.fail("W4A16 must not quantize activations") + + monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", reject_activation_quantization) + torch.manual_seed(17) + codes = torch.randint(0, 16, (n, k), device="cuda") + scales = (2.0 ** torch.randint(-2, 2, (n, k // 16), device="cuda")).to( + torch.float8_e4m3fn + ) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + (codes[:, ::2] | codes[:, 1::2] << 4).to(torch.uint8), requires_grad=False + ) + layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False) + layer.weight_global_scale = torch.nn.Parameter( + torch.tensor([0.125], device="cuda"), requires_grad=False + ) + lut = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6], + device="cuda", + ) + decoded = lut[codes] * scales.float().repeat_interleave(16, dim=1) * 0.125 + x_storage = torch.randn(1, m, k * 2, device="cuda", dtype=torch.bfloat16) + x = x_storage[..., ::2] + bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) + kernel = B12xNvFp4W4A16LinearKernel(None) + kernel.process_weights_after_loading(layer) + kernel.get_b12x_warmup_unit(layer, (1, 2, 4, 32), x.dtype).compile() + run = torch.compile( + lambda inputs: kernel.apply_weights(layer, inputs, bias), + fullgraph=True, + dynamic=True, + ) + run(x) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run(x) + for _ in range(2): + x_storage.normal_() + graph.replay() + expected = (x.double() @ decoded.double().T).bfloat16() + bias + torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.01) + + +def _block_weights(n, k, codec): + block_size, block_bytes = ( + (32, 34) if codec == "q8_0" else (256, 66 if codec == "iq2_xxs" else 74) + ) + raw = torch.randint(0, 256, (n, k // block_size, block_bytes), dtype=torch.uint8) + raw[..., :2] = torch.full( + (n, k // block_size, 1), 1 / 128, dtype=torch.float16 + ).view(torch.uint8) + return raw + + +@pytest.mark.parametrize("codec", ["iq2_xs", "iq2_xxs", "q8_0"]) +@pytest.mark.parametrize("profile_first", [False, True]) +@torch.inference_mode() +def test_b12x_block_quant_dense_compiled_graph( + default_vllm_config, codec, profile_first +): + supported, reason = B12xNvFp4W4A16LinearKernel.is_supported() + if not supported: + pytest.skip(reason) + from b12x.testing.iq2_xs_reference import dequantize_blocks + + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantLinearMethod, + ) + + default_vllm_config.scheduler_config.max_num_batched_tokens = 32 + default_vllm_config.scheduler_config.max_num_scheduled_tokens = 32 + default_vllm_config.compilation_config.cudagraph_capture_sizes = [1, 4] + method = ModelOptBlockQuantLinearMethod(codec) + raw = _block_weights(128, 256, codec) + decoded = dequantize_blocks(raw).bfloat16().cuda() + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(raw.cuda(), requires_grad=False) + method.process_weights_after_loading(layer) + assert layer.b12x_block_plan.prepared is None + if profile_first: + method.apply(layer, torch.zeros(32, 256, device="cuda", dtype=torch.bfloat16)) + layer.b12x_warmup_provider.get_b12x_warmup_unit( + layer, (1, 4, 32), torch.bfloat16 + ).compile() + assert layer.b12x_block_plan.prepared is not None + run = torch.compile(lambda x: method.apply(layer, x), fullgraph=True, dynamic=True) + for rows in (1, 4, 17): + x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16) + run(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run(x) + for _ in range(2): + x.normal_() + allocated = torch.accelerator.memory_allocated() + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() == allocated + expected = (x.float() @ decoded.float().T).bfloat16() + torch.testing.assert_close(output, expected, rtol=0.015, atol=0.03125) + graph.reset() + + +@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("profile_first", [False, True]) +@torch.inference_mode() +def test_b12x_q8_embedding_graph(default_vllm_config, id_dtype, profile_first): + supported, reason = B12xNvFp4W4A16LinearKernel.is_supported() + if not supported: + pytest.skip(reason) + from b12x.testing.q8_0_reference import dequantize_blocks + + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantLinearMethod, + ) + + default_vllm_config.scheduler_config.max_num_batched_tokens = 32 + default_vllm_config.scheduler_config.max_num_scheduled_tokens = 32 + default_vllm_config.compilation_config.cudagraph_capture_sizes = [1, 4] + method = ModelOptBlockQuantLinearMethod("q8_0") + method.is_embedding = True + raw = _block_weights(128, 256, "q8_0") + decoded = dequantize_blocks(raw).bfloat16().cuda() + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(raw.cuda(), requires_grad=False) + method.process_weights_after_loading(layer) + assert all(plan.prepared is None for plan in layer.b12x_embedding_plans.values()) + if profile_first: + method.embedding(layer, torch.zeros(32, device="cuda", dtype=id_dtype)) + layer.b12x_warmup_provider.get_b12x_warmup_unit( + layer, (1, 4, 32), torch.bfloat16 + ).compile() + assert all( + plan.prepared is not None for plan in layer.b12x_embedding_plans.values() + ) + run = torch.compile( + lambda x: method.embedding(layer, x), fullgraph=True, dynamic=True + ) + for count in (7, 9, 70): + ids = torch.arange(count, device="cuda", dtype=id_dtype) + torch._dynamo.mark_dynamic(ids, 0) + torch.testing.assert_close(run(ids), decoded[ids.long()], rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run(ids) + ids.add_(20) + allocated = torch.accelerator.memory_allocated() + graph.replay() + torch.accelerator.synchronize() + assert torch.accelerator.memory_allocated() == allocated + torch.testing.assert_close(output, decoded[ids.long()], rtol=0, atol=0) + graph.reset() diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index 71e22e3023..2313fa4304 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -25,6 +25,60 @@ def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: ) == (1, 2, 8, 128, 2048) +@pytest.mark.parametrize("is_embedding", [False, True]) +def test_b12x_block_quant_warmup_covers_each_layers_plans(monkeypatch, is_embedding): + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantLinearMethod, + ) + from vllm.model_executor.warmup.b12x_warmup import _collect_warmup_units + + method = object.__new__(ModelOptBlockQuantLinearMethod) + method.codec = "q8_0" + method.is_embedding = is_embedding + model = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()]) + calls = [] + monkeypatch.setattr( + method, "apply", lambda layer, x: calls.append((layer, x.shape, x.dtype)) + ) + monkeypatch.setattr( + method, + "embedding", + lambda layer, ids: calls.append((layer, ids.shape, ids.dtype)), + ) + for index, layer in enumerate(model): + layer.b12x_warmup_provider = method + if is_embedding: + layer.weight = torch.nn.Parameter( + torch.empty((8, 8, 34), dtype=torch.uint8), requires_grad=False + ) + layer.b12x_embedding_plans = { + dtype: SimpleNamespace( + handle=2 * index + offset, query=SimpleNamespace(max_rows=32) + ) + for offset, dtype in enumerate((torch.int32, torch.int64)) + } + else: + layer.b12x_block_weight = SimpleNamespace( + in_features=256, values=torch.empty(1) + ) + layer.b12x_block_plan = SimpleNamespace( + handle=index, token_counts=(1, 4, 32) + ) + + units = list(_collect_warmup_units(model, (1, 17, 32), torch.bfloat16)) + assert len(units) == 2 + for unit in units: + unit.compile() + expected = ( + [((32,), torch.int32), ((32,), torch.int64)] + if is_embedding + else [((rows, 256), torch.bfloat16) for rows in (1, 4, 17, 32)] + ) + assert calls == [ + (layer, shape, dtype) for layer in model for shape, dtype in expected + ] + + @pytest.mark.parametrize( ("kernel_cls", "module_name", "call_name", "layer", "name"), [ diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 33c798cba0..3ecff626c1 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -7,6 +7,7 @@ Run `pytest tests/quantization/test_modelopt.py`. import os from contextlib import nullcontext +from types import SimpleNamespace from typing import Any, NoReturn from unittest.mock import MagicMock, Mock, patch @@ -1027,3 +1028,248 @@ def test_modelopt_fp8_pb_wo_rejects_non_128_input(): scheme.create_weights( torch.nn.Module(), mo.WEIGHT, mo.CkptCtx(), shapes, Mock() ) + + +def _iq2_xs_moe_config(tp_size=1, tp_rank=0): + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + return SimpleNamespace( + moe_backend="b12x", + in_dtype=torch.bfloat16, + activation=MoEActivation.SILU, + has_bias=False, + swiglu_limit=None, + swiglu_alpha=None, + swiglu_beta=None, + hidden_dim=256, + intermediate_size_per_partition=512 // tp_size, + tp_rank=tp_rank, + moe_parallel_config=SimpleNamespace( + use_ep=False, + ep_size=1, + use_all2all_kernels=False, + enable_eplb=False, + tp_size=tp_size, + ), + ) + + +@pytest.mark.parametrize( + "codec,block_size,block_bytes", + [("iq2_xs", 256, 74), ("iq2_xxs", 256, 66), ("q8_0", 32, 34)], +) +def test_modelopt_mixed_precision_dispatches_iq2_xs_experts( + monkeypatch, codec, block_size, block_bytes +): + from vllm.model_executor.layers.fused_moe.b12x import B12xExperts + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantMoEMethod, + ) + + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = _mixed_precision_config( + { + "model.language_model.layers.0.mlp.experts": { + "quant_algo": codec.upper(), + "group_size": block_size, + "block_payload_bytes": block_bytes, + "packing": "ggml", + }, + } + ) + layer = MagicMock(spec=RoutedExperts) + layer.moe_config = _iq2_xs_moe_config() + method = config.get_quant_method( + layer, + "language_model.model.layers.0.mlp.experts", + ) + assert isinstance(method, ModelOptBlockQuantMoEMethod) + quant = method.get_fused_moe_quant_config(layer) + assert quant.weight_quant_dtype == codec + assert quant.quant_dtype is None + + +@pytest.mark.parametrize( + "field,value", + [ + ("group_size", 0), + ("block_payload_bytes", 75), + ("packing", "unknown"), + ], +) +@pytest.mark.parametrize( + "codec,block_size,block_bytes", + [("iq2_xs", 256, 74), ("iq2_xxs", 256, 66), ("q8_0", 32, 34)], +) +def test_modelopt_block_quant_rejects_incompatible_block_metadata( + field, value, codec, block_size, block_bytes +): + recipe = { + "quant_algo": codec.upper(), + "group_size": block_size, + "block_payload_bytes": block_bytes, + "packing": "ggml", + field: value, + } + with pytest.raises(ValueError, match=f"unsupported {codec.upper()} block contract"): + _mixed_precision_config({"model.layers.0.mlp.experts": recipe}) + + +@pytest.mark.parametrize("parallel", ["row", "column"]) +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize( + "codec,block_size,block_bytes", + [("iq2_xs", 256, 74), ("iq2_xxs", 256, 66), ("q8_0", 32, 34)], +) +def test_modelopt_block_quant_dense_loading_preserves_whole_blocks( + monkeypatch, parallel, tp_rank, codec, block_size, block_bytes +): + import vllm.model_executor.parameter as parameter_module + import vllm.utils.b12x as b12x_utils + from vllm.model_executor.layers.linear import LinearBase + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantLinearMethod, + ) + + monkeypatch.setattr( + b12x_utils, + "get_b12x_blockscaled", + lambda: SimpleNamespace( + BlockQuantLinearWeight=object, + is_supported=lambda: True, + ), + ) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + prefix = "model.layers.1.mixer.shared_experts.up_proj" + config = _mixed_precision_config( + { + prefix: { + "quant_algo": codec.upper(), + "group_size": block_size, + "block_payload_bytes": block_bytes, + "packing": "ggml", + } + } + ) + method = config.get_quant_method(MagicMock(spec=LinearBase), prefix) + assert isinstance(method, ModelOptBlockQuantLinearMethod) + layer = torch.nn.Module() + n, k = (128, 256) if parallel == "row" else (64, 512) + method.create_weights( + layer, k, [n], 512, 128, torch.bfloat16, weight_loader=lambda *args: None + ) + source = torch.randint( + 0, 256, (128, 512 // block_size, block_bytes), dtype=torch.uint8 + ) + if parallel == "row": + layer.weight.load_row_parallel_weight(source) + expected = source[ + :, tp_rank * (256 // block_size) : (tp_rank + 1) * (256 // block_size) + ] + else: + layer.weight.load_column_parallel_weight(source) + expected = source[tp_rank * n : (tp_rank + 1) * n] + torch.testing.assert_close(layer.weight, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("tp_size,tp_rank", [(1, 0), (2, 0), (2, 1)]) +@pytest.mark.parametrize("gated", [True, False]) +@pytest.mark.parametrize( + "codec,block_size,block_bytes", + [("iq2_xs", 256, 74), ("iq2_xxs", 256, 66), ("q8_0", 32, 34)], +) +def test_modelopt_block_quant_loads_expert_blocks_with_aligned_tp_slices( + monkeypatch, + tp_size, + tp_rank, + gated, + codec, + block_size, + block_bytes, +): + import vllm.model_executor.parameter as parameter_module + from vllm.model_executor.layers.fused_moe.b12x import B12xExperts + from vllm.model_executor.layers.quantization.modelopt_block_quant import ( + ModelOptBlockQuantMoEMethod, + ) + + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + config = _iq2_xs_moe_config(tp_size, tp_rank) + if not gated: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + config.activation = MoEActivation.RELU2_NO_MUL + method = ModelOptBlockQuantMoEMethod(config, codec=codec) + layer = torch.nn.Module() + layer.layer_name = "model.layers.0.mlp.experts" + layer.is_fused_checkpoint_transposed = False + layer._orient_fused_weight = RoutedExperts._orient_fused_weight + layer.get_expert_mapping = lambda **kwargs: ( + RoutedExperts.build_expert_params_mapping( + "gate_proj", + "down_proj", + "up_proj", + num_experts=2, + routed_experts_prefix="", + **kwargs, + ) + ) + local_i = 512 // tp_size + method.create_weights(layer, 2, 256, local_i, torch.bfloat16) + assert layer.w13_weight.shape[1] == local_i * (2 if gated else 1) + layer.w13_weight.data.zero_() + layer.w2_weight.data.zero_() + for shard, name in ( + ("w1", "w13_weight"), + ("w3", "w13_weight"), + ("w2", "w2_weight"), + ): + if shard == "w3" and not gated: + continue + shape = ( + (256, 512 // block_size, block_bytes) + if shard == "w2" + else (512, 256 // block_size, block_bytes) + ) + source = torch.randint(1, 256, shape, dtype=torch.uint8) + param = getattr(layer, name) + projection = {"w1": "gate", "w3": "up", "w2": "down"}[shard] + loaded = list( + RoutedExperts.load_weights( + layer, + [(f"1.{projection}_proj.weight", source)], + ) + ) + assert loaded == [name] + if shard == "w2": + expected = source[ + :, + tp_rank * (local_i // block_size) : (tp_rank + 1) + * (local_i // block_size), + ] + actual = param[1] + else: + expected = source[tp_rank * local_i : (tp_rank + 1) * local_i] + offset = 0 if shard == "w1" else local_i + actual = param[1, offset : offset + local_i] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert param[0].count_nonzero() == 0 + with pytest.raises(ValueError, match="invalid packed tensor"): + layer.w2_weight.weight_loader( + layer.w2_weight, + torch.zeros(256, 512), + "w2_weight", + "w2", + 0, + ) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 0ce22f477a..bb7adad164 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -135,6 +135,7 @@ from vllm.model_executor.kernels.linear.nvfp4 import ( ) from vllm.model_executor.kernels.linear.nvfp4.b12x import ( B12xNvFp4LinearKernel, + B12xNvFp4W4A16LinearKernel, ) from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( CutlassNvFp4LinearKernel, @@ -275,6 +276,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { B12xMxFp4LinearKernel, B12xMxfp8LinearKernel, B12xNvFp4LinearKernel, + B12xNvFp4W4A16LinearKernel, B12xTensorFP8ScaledMMLinearKernel, }, "cutlass": { @@ -582,6 +584,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { FlashInferCudnnNvFp4LinearKernel, FbgemmNvFp4LinearKernel, B12xNvFp4LinearKernel, + B12xNvFp4W4A16LinearKernel, TorchNvFp4LinearKernel, EmulationNvFp4LinearKernel, HummingNvFp4LinearKernel, @@ -1127,6 +1130,7 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: current platform.""" config = NvFp4LinearLayerConfig() a16_kernels = ( + B12xNvFp4W4A16LinearKernel, FlashInferCuteDslNvFp4W4A16LinearKernel, MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel, @@ -1341,6 +1345,7 @@ __all__ = [ "B12xMxfp8LinearKernel", "B12xMxFp4LinearKernel", "B12xNvFp4LinearKernel", + "B12xNvFp4W4A16LinearKernel", "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", diff --git a/vllm/model_executor/kernels/linear/nvfp4/b12x.py b/vllm/model_executor/kernels/linear/nvfp4/b12x.py index 4ff531dd48..6cbd0e3a82 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/b12x.py +++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py @@ -138,4 +138,64 @@ class B12xNvFp4LinearKernel(NvFp4LinearKernel): ) -__all__ = ["B12xNvFp4LinearKernel"] +class B12xNvFp4W4A16LinearKernel(B12xNvFp4LinearKernel): + """BF16 x NVFP4 GEMM without activation quantization.""" + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + layer.b12x_nvfp4_packed_weight = blockscaled.pack_weight( + layer.weight, + layer.weight_scale, + recipe="nvfp4", + global_scale=layer.weight_global_scale, + ) + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, layer.weight.shape[1] * 2), + dtype=output_dtype, + device=layer.weight.device, + ) + self.apply_weights(layer, source) + + return B12xWarmupUnit( + name="NVFP4 W4A16", + key=( + type(self), + layer.weight.device, + tuple(layer.weight.shape), + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if x.dtype != torch.bfloat16: + raise ValueError(f"b12x NVFP4 W4A16 requires BF16 input, got {x.dtype}") + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + output = blockscaled.mm( + x.reshape(-1, x.shape[-1]).contiguous(), + layer.b12x_nvfp4_packed_weight, + required_mode="a16", + ) + if bias is not None: + output = output + bias + return output.view(*x.shape[:-1], layer.weight.shape[0]) + + +__all__ = ["B12xNvFp4LinearKernel", "B12xNvFp4W4A16LinearKernel"] diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index d31ba1489d..6ba4d24072 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.utils.b12x import ( + B12X_BLOCK_CODECS, B12xWarmupUnit, get_b12x_fused_moe, reuse_packed_weight_storage, @@ -38,6 +39,7 @@ _B12X_MOE_MODES: dict[ tuple[torch.dtype | str | None, torch.dtype | str | None], tuple[str, str, str], ] = { + **{(codec, None): ("w4a16", codec, "w31") for codec in B12X_BLOCK_CODECS}, ("mxfp4", "mxfp8"): ("w4a8_mx", "fp4_e8m0_k32", "w31"), ("mxfp4", None): ("w4a16", "fp4_e8m0_k32", "w31"), ("nvfp4", "nvfp4"): ("nvfp4", "modelopt_nvfp4", "w31"), @@ -206,10 +208,13 @@ class B12xExperts(mk.FusedMoEExpertsModular): quant_config: FusedMoEQuantConfig, ): super().__init__(moe_config, quant_config) - if quant_config.weight_quant_dtype not in ("mxfp4", "nvfp4"): + if quant_config.weight_quant_dtype not in ( + "mxfp4", + "nvfp4", + *B12X_BLOCK_CODECS, + ): raise ValueError( - "b12x MoE requires MXFP4 or NVFP4 weights, got " - f"{quant_config.weight_quant_dtype}" + f"unsupported b12x MoE weight format: {quant_config.weight_quant_dtype}" ) scheme = ( quant_config.weight_quant_dtype, @@ -293,6 +298,31 @@ class B12xExperts(mk.FusedMoEExpertsModular): raise RuntimeError( "b12x MoE weights must be prepared before CUDA graph capture" ) + if self._source_format in B12X_BLOCK_CODECS: + fused_moe = _require_b12x_fused_moe() + block_size, _ = B12X_BLOCK_CODECS[self._source_format] + plan = fused_moe.plan_weights( + source=fused_moe.PackedSource( + format=self._source_format, w13_layout=self._w13_layout + ), + activation=fused_moe.ActivationSpec( + mode="a16", + nonlinearity=_b12x_activation_name(activation), + io_dtype=params_dtype, + ), + geometry=fused_moe.MoEGeometry( + num_experts=w1.shape[0], + hidden_size=w2.shape[1], + intermediate_size=w2.shape[2] * block_size, + ), + ) + prepared = fused_moe.prepare_weights( + plan=plan, + weights=fused_moe.BlockQuantWeights( + w13=w1, w2=w2, codec=self._source_format + ), + ) + return prepared._impl if self.w1_scale is None or self.w2_scale is None: raise ValueError("b12x MoE requires w1 and w2 block scales") @@ -345,6 +375,8 @@ class B12xExperts(mk.FusedMoEExpertsModular): ) def _refresh_quant_config(self, layer: torch.nn.Module) -> None: + if self._source_format in B12X_BLOCK_CODECS: + return self.quant_config._w1.scale = layer.w13_weight_scale self.quant_config._w2.scale = layer.w2_weight_scale if self._source_format != "modelopt_nvfp4": @@ -446,8 +478,10 @@ class B12xExperts(mk.FusedMoEExpertsModular): ): return ( False, - "MXFP4 W4A8 requires hidden size divisible by 256 and " - "per-rank intermediate size divisible by 32", + ( + "MXFP4 W4A8 requires hidden size divisible by 256 and " + "per-rank intermediate size divisible by 32" + ), ) return mk.FusedMoEExperts.is_supported_config( cls, moe_config, weight_key, activation_key, activation_format @@ -530,14 +564,15 @@ class B12xExperts(mk.FusedMoEExpertsModular): w2: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple[int, int, int, int, int]: - if w1.numel() and w2.numel(): + if w1.numel() and w2.numel() and self._source_format not in B12X_BLOCK_CODECS: return super().moe_problem_size(a1, w1, w2, topk_ids) prepared = self._prepared() tokens = int(a1.shape[0] if a1.ndim == 2 else a1.shape[1]) return ( int(prepared.num_experts), tokens, - int(prepared.intermediate_size) * 2, + int(prepared.intermediate_size) + * (2 if self.moe_config.activation.is_gated else 1), int(a1.shape[-1]), int(topk_ids.shape[1]), ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index ae5829a311..71f4206008 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -986,12 +986,20 @@ class RoutedExperts(PluggableLayer): for expert_name, loaded_weight in weights: qual_name = f"{self.layer_name}.{expert_name}" - is_fused = loaded_weight.dim() == 3 + expert_key = qual_name.partition("experts.")[2].partition(".")[0] + named_expert = ( + mapping_by_expert.get(expert_key) + if qual_name.count("experts.") == 1 + else None + ) + # Encoded per-expert [N, K / block_size, bytes_per_block] weights + # have the same rank as fused [E, N, K] weights. An explicit expert + # name identifies the per-expert layout without treating N as E. + is_fused = loaded_weight.dim() == 3 and named_expert is None # Fused tensors and ambiguous names keep the full mapping. candidates = expert_mapping if not is_fused and qual_name.count("experts.") == 1: - expert_key = qual_name.partition("experts.")[2].partition(".")[0] - candidates = mapping_by_expert.get(expert_key, expert_mapping) + candidates = named_expert or expert_mapping matched = False for param_name, weight_name, expert_id, shard_id in candidates: diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 98af93c30f..8354f885d9 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -99,7 +99,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( requantize_with_max_scale, ) -from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) from vllm.model_executor.parameter import ( BlockQuantScaleParameter, ChannelQuantScaleParameter, @@ -111,6 +114,7 @@ from vllm.model_executor.utils import ( replace_parameter, set_weight_attrs, ) +from vllm.utils.b12x import B12X_BLOCK_CODECS from vllm.utils.math_utils import cdiv if TYPE_CHECKING: @@ -1559,6 +1563,19 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): super().__init__(exclude_modules) self.kv_cache_quant_method = kv_cache_quant_method self.quantized_layers = quantized_layers + for prefix, recipe in quantized_layers.items(): + codec = recipe.get("quant_algo", "").lower() + if codec in B12X_BLOCK_CODECS and any( + recipe.get(key) != value + for key, value in { + "group_size": B12X_BLOCK_CODECS[codec][0], + "block_payload_bytes": B12X_BLOCK_CODECS[codec][1], + "packing": "ggml", + }.items() + ): + raise ValueError( + f"unsupported {codec.upper()} block contract for {prefix}: {recipe}" + ) self.fp8_config = fp8_config self.nvfp4_config = nvfp4_config self.w4a16_nvfp4_config = w4a16_nvfp4_config @@ -1811,6 +1828,22 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): quant_algo = self._resolve_quant_algo(prefix) + if quant_algo is not None and quant_algo.lower() in B12X_BLOCK_CODECS: + from .modelopt_block_quant import ( + ModelOptBlockQuantLinearMethod, + ModelOptBlockQuantMoEMethod, + ) + + if isinstance(layer, RoutedExperts): + return ModelOptBlockQuantMoEMethod( + layer.moe_config, codec=quant_algo.lower() + ) + if isinstance(layer, (LinearBase, VocabParallelEmbedding)): + return ModelOptBlockQuantLinearMethod(codec=quant_algo.lower()) + raise ValueError( + f"{quant_algo} requires a linear, embedding or routed experts" + ) + if isinstance(layer, (LinearBase, ParallelLMHead)): # Per-prefix algo -> its sub-config, then the generic linear method. # resolve() reads group_size off whichever sub-config it is handed. diff --git a/vllm/model_executor/layers/quantization/modelopt_block_quant.py b/vllm/model_executor/layers/quantization/modelopt_block_quant.py new file mode 100644 index 0000000000..d0933563ef --- /dev/null +++ b/vllm/model_executor/layers/quantization/modelopt_block_quant.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ModelOpt safetensors block codecs with BF16 activations.""" + +import torch + +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + FusedMoEQuantDesc, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.linear import ( + LinearMethodBase, + register_weight_loader_v2_supported_method, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.utils.b12x import B12X_BLOCK_CODECS, B12xWarmupUnit, get_b12x_blockscaled + + +@torch.library.custom_op( + "vllm::modelopt_block_embedding_chunked", mutates_args=("out",) +) +def _modelopt_block_embedding_chunked( + weight: torch.Tensor, ids: torch.Tensor, out: torch.Tensor, plan_handle: int +) -> None: + from b12x.preparation.types import plan_from_handle + from b12x.sequence import embedding + + plan = plan_from_handle(plan_handle) + capacity = plan.query.max_rows + for start in range(0, ids.numel(), capacity): + embedding.run( + weight, + ids[start : start + capacity], + out=out[start : start + capacity], + plan=plan, + ) + + +@_modelopt_block_embedding_chunked.register_fake +def _modelopt_block_embedding_chunked_fake( + weight: torch.Tensor, ids: torch.Tensor, out: torch.Tensor, plan_handle: int +) -> None: + return None + + +class ModelOptBlockQuantMoEMethod(FusedMoEMethodBase): + """Load native blocks and retain compact b12x prepared expert storage.""" + + def __init__(self, moe_config: FusedMoEConfig, codec: str = "iq2_xs"): + super().__init__(moe_config) + self.codec = codec.lower() + self.block_size, self.block_bytes = B12X_BLOCK_CODECS[self.codec] + if moe_config.moe_backend not in ("auto", "b12x"): + raise ValueError( + "Block-quantized routed experts require the b12x MoE backend" + ) + if not B12xExperts._supports_current_device(): + raise ValueError("Block-quantized routed experts require b12x on SM12x") + if not B12xExperts._supports_parallel_config(moe_config.moe_parallel_config): + raise ValueError( + "Block-quantized supports tensor parallelism without EP or EPLB" + ) + if moe_config.in_dtype != torch.bfloat16: + raise ValueError("Block-quantized routed experts require BF16 activations") + if ( + moe_config.activation + not in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) + or moe_config.has_bias + ): + raise ValueError( + "Block-quantized routed experts require bias-free SiLU or ReLU2" + ) + self.gated = moe_config.activation == MoEActivation.SILU + if any( + value is not None + for value in ( + moe_config.swiglu_limit, + moe_config.swiglu_alpha, + moe_config.swiglu_beta, + ) + ): + raise ValueError( + "Block-quantized routed experts require standard SiLU parameters" + ) + + def maybe_roundup_sizes( + self, + hidden_size, + intermediate_size_per_partition, + act_dtype, + moe_parallel_config, + ): + alignment = max(128, self.block_size) + if hidden_size % alignment or intermediate_size_per_partition % alignment: + raise ValueError( + f"{self.codec.upper()} hidden and per-rank intermediate sizes " + f"must align to {alignment}" + ) + return hidden_size, intermediate_size_per_partition + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + self.maybe_roundup_sizes( + hidden_size, + intermediate_size_per_partition, + params_dtype, + self.moe.moe_parallel_config, + ) + for name, shape in ( + ( + "w13_weight", + ( + num_experts, + (2 if self.gated else 1) * intermediate_size_per_partition, + hidden_size // self.block_size, + self.block_bytes, + ), + ), + ( + "w2_weight", + ( + num_experts, + hidden_size, + intermediate_size_per_partition // self.block_size, + self.block_bytes, + ), + ), + ): + layer.register_parameter( + name, + ModelWeightParameter( + data=torch.empty(shape, dtype=torch.uint8), + input_dim=2, + output_dim=1, + weight_loader=self.weight_loader, + ), + ) + + def weight_loader( + self, + param, + loaded_weight, + weight_name, + shard_id, + expert_id, + return_success=False, + ): + if shard_id not in ("w1", "w2", "w3"): + raise ValueError(f"invalid Block-quantized expert projection: {shard_id}") + if not self.gated and shard_id == "w3": + raise ValueError("non-gated Block-quantized experts have no w3 projection") + local_i = self.moe.intermediate_size_per_partition + global_i = local_i * self.moe.moe_parallel_config.tp_size + hidden = self.moe.hidden_dim + expected = ( + (hidden, global_i // self.block_size, self.block_bytes) + if shard_id == "w2" + else (global_i, hidden // self.block_size, self.block_bytes) + ) + if loaded_weight.dtype != torch.uint8 or tuple(loaded_weight.shape) != expected: + raise ValueError( + f"invalid packed tensor {weight_name}: expected uint8{expected}, " + f"got {loaded_weight.dtype}{tuple(loaded_weight.shape)}" + ) + if not 0 <= expert_id < param.shape[0]: + raise ValueError(f"invalid Block-quantized expert index: {expert_id}") + destination = param.data[expert_id] + if shard_id == "w2": + source = loaded_weight.narrow( + 1, + self.moe.tp_rank * (local_i // self.block_size), + local_i // self.block_size, + ) + else: + source = loaded_weight.narrow(0, self.moe.tp_rank * local_i, local_i) + destination = destination.narrow( + 0, 0 if shard_id == "w1" else local_i, local_i + ) + destination.copy_(source) + return True if return_success else None + + def get_fused_moe_quant_config(self, layer): + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(dtype=self.codec), + _w2=FusedMoEQuantDesc(dtype=self.codec), + ) + + def process_weights_after_loading(self, layer): + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + prepare_finalize = maybe_make_prepare_finalize( + moe=self.moe, + quant_config=self.moe_quant_config, + routing_tables=layer._expert_routing_tables(), + allow_new_interface=True, + ) + assert prepare_finalize is not None + experts = B12xExperts(self.moe, self.moe_quant_config) + self.moe_kernel = mk.FusedMoEKernel(prepare_finalize, experts) + experts.process_weights_after_loading(layer) + + def apply( + self, + layer, + x, + topk_weights, + topk_ids, + shared_experts, + shared_experts_input, + ): + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + +ModelOptBlockQuantMoEMethod.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +@register_weight_loader_v2_supported_method +class ModelOptBlockQuantLinearMethod(LinearMethodBase): + """Load native block payloads and execute the prepared b12x dense API.""" + + def __init__(self, codec: str = "iq2_xs"): + self.codec = codec.lower() + self.block_size, self.block_bytes = B12X_BLOCK_CODECS[self.codec] + self.is_embedding = False + from vllm.model_executor.kernels.linear import _get_linear_backend + from vllm.utils.b12x import get_b12x_blockscaled + + if _get_linear_backend(quantization=self.codec) not in ("auto", "b12x"): + raise ValueError( + "Block-quantized dense weights require the b12x linear backend" + ) + api = get_b12x_blockscaled() + if ( + api is None + or not hasattr(api, "BlockQuantLinearWeight") + or not api.is_supported() + ): + raise ValueError("Block-quantized dense weights require b12x on SM12x") + + def create_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): + n = sum(output_partition_sizes) + self.is_embedding = isinstance( + layer, VocabParallelEmbedding + ) and not isinstance(layer, ParallelLMHead) + if self.is_embedding and self.codec != "q8_0": + raise ValueError("packed embedding currently requires Q8_0") + if ( + params_dtype != torch.bfloat16 + or input_size_per_partition % self.block_size + or n % 8 + ): + raise ValueError( + f"{self.codec.upper()} dense weights require BF16, " + f"K{self.block_size} and N8" + ) + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty( + (n, input_size_per_partition // self.block_size, self.block_bytes), + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=extra_weight_attrs["weight_loader"], + ), + ) + + def process_weights_after_loading(self, layer): + from b12x.sequence import embedding + + from vllm.model_executor.utils import replace_parameter + + api = get_b12x_blockscaled() + assert api is not None + config = get_current_vllm_config() + scheduler = config.scheduler_config + capture_sizes = config.compilation_config.cudagraph_capture_sizes or [] + capacity = max( + scheduler.max_num_batched_tokens, + scheduler.max_num_scheduled_tokens or 0, + *capture_sizes, + ) + weight = layer.weight.data + if self.is_embedding: + bases = weight[..., :2].contiguous().view(torch.float16) + if not torch.isfinite(bases).all().item(): + raise ValueError("Q8_0 embedding scales must be finite") + layer.b12x_embedding_plans = {} + width = weight.shape[1] * self.block_size + if config.model_config is not None: + capacity = max(capacity, config.model_config.max_model_len) + + for dtype in (torch.int32, torch.int64): + plan = embedding.plan( + embedding.EmbeddingQuery( + max_rows=capacity, + table_rows=weight.shape[0], + width=width, + row_stride=weight.stride(0), + weight_dtype="uint8", + id_dtype=str(dtype).removeprefix("torch."), + ), + device=weight.device, + ) + layer.b12x_embedding_plans[dtype] = plan + else: + packed = api.pack_weight(weight, recipe=self.codec) + query = api.BlockscaledQuery( + recipe=self.codec, + num_tokens=capacity, + in_features=packed.in_features, + padded_in_features=packed.in_features, + out_features=packed.out_features, + activation_mode="a16", + ) + plan = api.plan_regimes( + query, + exact_m=tuple(sorted({m for m in capture_sizes if 0 < m < capacity})), + ) + layer.b12x_block_weight = packed + layer.b12x_block_plan = plan + if not self.is_embedding: + replace_parameter( + layer, "weight", torch.empty(0, dtype=torch.uint8, device=weight.device) + ) + + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + if self.is_embedding: + plans = tuple(layer.b12x_embedding_plans.values()) + else: + plans = (layer.b12x_block_plan,) + + def compile() -> None: + if self.is_embedding: + for dtype, plan in layer.b12x_embedding_plans.items(): + ids = torch.zeros( + plan.query.max_rows, dtype=dtype, device=layer.weight.device + ) + self.embedding(layer, ids) + else: + packed = layer.b12x_block_weight + for tokens in sorted(set(token_counts) | set(plans[0].token_counts)): + source = torch.zeros( + (tokens, packed.in_features), + dtype=output_dtype, + device=packed.values.device, + ) + self.apply(layer, source) + + return B12xWarmupUnit( + name=( + self.codec.upper() + (" embedding" if self.is_embedding else " linear") + ), + # Every layer owns a plan that must be prepared before capture. + key=(type(self), *(plan.handle for plan in plans)), + compile=compile, + ) + + def apply(self, layer, x, bias=None): + if x.dtype != torch.bfloat16: + raise ValueError( + "Block-quantized dense execution requires BF16 activations" + ) + api = get_b12x_blockscaled() + assert api is not None + output = api.mm( + x.reshape(-1, x.shape[-1]).contiguous(), + layer.b12x_block_weight, + plan=layer.b12x_block_plan, + bias=bias, + ) + return output.view(*x.shape[:-1], layer.b12x_block_weight.out_features) + + def embedding(self, layer, input_): + plan = layer.b12x_embedding_plans[input_.dtype] + width = layer.weight.shape[1] * self.block_size + out = torch.empty( + (*input_.shape, width), dtype=torch.bfloat16, device=layer.weight.device + ) + ids, rows = input_.reshape(-1), out.view(-1, width) + _modelopt_block_embedding_chunked(layer.weight, ids, rows, plan.handle) + return out diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 050a13431d..4d6b5fb2cf 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -584,6 +584,8 @@ class NemotronHModel(nn.Module, EagleModelMixin): self.embed_tokens = VocabParallelEmbedding( self.vocab_size, config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", ) self.has_moe = "E" in config.hybrid_override_pattern diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 1ed0febc85..144895ab1b 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -11,6 +11,9 @@ from typing import Any import torch +# Logical elements and encoded bytes per block. +B12X_BLOCK_CODECS = {"iq2_xs": (256, 74), "iq2_xxs": (256, 66), "q8_0": (32, 34)} + @dataclass(frozen=True) class B12xWarmupUnit: