|
53 | 53 | from vllm.utils.import_utils import has_deep_gemm |
54 | 54 |
|
55 | 55 | import modelopt.torch.quantization as mtq |
| 56 | +from modelopt.torch.opt.config_loader import load_config |
56 | 57 | from modelopt.torch.quantization.config import QuantizerAttributeConfig |
| 58 | +from modelopt.torch.quantization.conversion import set_quantizer_by_cfg |
57 | 59 | from modelopt.torch.quantization.nn import TensorQuantizer |
58 | 60 | from modelopt.torch.quantization.plugins import vllm as vllm_plugin |
59 | 61 | from modelopt.torch.quantization.plugins.vllm import ( |
@@ -841,6 +843,86 @@ def test_update_kv_cfg_for_mla_skips_non_mla_and_warns_on_affine(): |
841 | 843 | } |
842 | 844 |
|
843 | 845 |
|
| 846 | +class _NativeMLAAttention(torch.nn.Module): |
| 847 | + def forward(self, query, kv_c, k_pe, *args, **kwargs): |
| 848 | + return query, kv_c, k_pe |
| 849 | + |
| 850 | + |
| 851 | +@pytest.mark.skipif(VllmMLAAttention is None, reason="this vLLM has no MLAAttention") |
| 852 | +@pytest.mark.parametrize("rope_dim", [0, 64], ids=("nope", "rope")) # GLM-5.3-Flash, DeepSeek-V3 |
| 853 | +def test_kv_nvfp4_mla_unit_quantizes_the_mla_kv_cache(monkeypatch, rope_dim): |
| 854 | + """Like vLLM's ``nvfp4_ds_mla`` cache: an NVFP4 latent and an unscaled FP8 RoPE key ``k_pe``. |
| 855 | + An empty NoPE ``k_pe`` passes through.""" |
| 856 | + monkeypatch.setattr( |
| 857 | + vllm_plugin, |
| 858 | + "create_parallel_state", |
| 859 | + lambda: vllm_plugin.ParallelState(data_parallel_group=None), |
| 860 | + ) |
| 861 | + |
| 862 | + class _TestQuantVLLMMLAAttention(vllm_plugin._QuantVLLMMLAAttention, _NativeMLAAttention): |
| 863 | + pass |
| 864 | + |
| 865 | + mla = _new_attention(_TestQuantVLLMMLAAttention) |
| 866 | + mla._setup() |
| 867 | + set_quantizer_by_cfg( |
| 868 | + mla, |
| 869 | + [{"quantizer_name": "*", "enable": False}, *load_config("configs/ptq/units/kv_nvfp4_mla")], |
| 870 | + ) |
| 871 | + |
| 872 | + assert mla.kv_c_bmm_quantizer.is_enabled |
| 873 | + assert mla.kv_c_bmm_quantizer.amax == 6.0 * 448.0 |
| 874 | + assert mla.k_pe_bmm_quantizer.is_enabled |
| 875 | + assert not mla.q_bmm_quantizer.is_enabled |
| 876 | + |
| 877 | + mla.to("cuda") |
| 878 | + query = torch.randn(5, 4, 512, device="cuda", dtype=torch.bfloat16) |
| 879 | + kv_c = torch.randn(5, 512, device="cuda", dtype=torch.bfloat16) |
| 880 | + k_pe = torch.randn(5, 1, rope_dim, device="cuda", dtype=torch.bfloat16) |
| 881 | + out_query, out_kv_c, out_k_pe = mla(query, kv_c, k_pe) |
| 882 | + assert out_query is query |
| 883 | + assert not torch.equal(out_kv_c, kv_c) |
| 884 | + # NVFP4 with a fixed global scale is idempotent: the output is already on the grid. |
| 885 | + assert torch.equal(mla.kv_c_bmm_quantizer(out_kv_c), out_kv_c) |
| 886 | + if rope_dim: |
| 887 | + assert not torch.equal(out_k_pe, k_pe) |
| 888 | + assert torch.equal(out_k_pe, k_pe.to(torch.float8_e4m3fn).to(torch.bfloat16)) |
| 889 | + else: |
| 890 | + assert out_k_pe is k_pe |
| 891 | + |
| 892 | + |
| 893 | +@pytest.mark.parametrize( |
| 894 | + ("name", "shape"), |
| 895 | + [("kv_c_bmm_quantizer", (8, 512)), ("k_pe_bmm_quantizer", (8, 1, 64))], |
| 896 | + ids=("kv_c", "k_pe"), |
| 897 | +) |
| 898 | +def test_kv_nvfp4_mla_quantizer_replays_in_cuda_graph(name, shape): |
| 899 | + """The latent's constant amax is created on the CPU; on the GPU (FakeQuantWorker moves every |
| 900 | + quantizer there before CUDA graph capture) each fake quant of the unit captures and replays |
| 901 | + like eager.""" |
| 902 | + layer = torch.nn.Module() |
| 903 | + setattr(layer, name, TensorQuantizer()) |
| 904 | + set_quantizer_by_cfg(layer, load_config("configs/ptq/units/kv_nvfp4_mla")) |
| 905 | + quantizer = getattr(layer, name) |
| 906 | + if name == "kv_c_bmm_quantizer": |
| 907 | + assert quantizer._amax.device.type == "cpu" |
| 908 | + quantizer.to("cuda") |
| 909 | + |
| 910 | + static_in = torch.randn(*shape, device="cuda", dtype=torch.bfloat16) |
| 911 | + stream = torch.cuda.Stream() |
| 912 | + stream.wait_stream(torch.cuda.current_stream()) |
| 913 | + with torch.cuda.stream(stream): |
| 914 | + quantizer(static_in) # compile the kernel before capture |
| 915 | + torch.cuda.current_stream().wait_stream(stream) |
| 916 | + graph = torch.cuda.CUDAGraph() |
| 917 | + with torch.cuda.graph(graph): |
| 918 | + static_out = quantizer(static_in) |
| 919 | + |
| 920 | + new_in = torch.randn_like(static_in) |
| 921 | + static_in.copy_(new_in) |
| 922 | + graph.replay() |
| 923 | + assert torch.equal(static_out, quantizer(new_in)) |
| 924 | + |
| 925 | + |
844 | 926 | def _quantize_and_summarize(self): |
845 | 927 | """Run on the worker via ``LLM.collective_rpc``. |
846 | 928 |
|
|
0 commit comments