|
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 ( |
@@ -834,6 +836,73 @@ def test_update_kv_cfg_for_mla_skips_non_mla_and_warns_on_affine(): |
834 | 836 | } |
835 | 837 |
|
836 | 838 |
|
| 839 | +class _NativeMLAAttention(torch.nn.Module): |
| 840 | + def forward(self, query, kv_c, k_pe, *args, **kwargs): |
| 841 | + return query, kv_c, k_pe |
| 842 | + |
| 843 | + |
| 844 | +@pytest.mark.skipif(VllmMLAAttention is None, reason="this vLLM has no MLAAttention") |
| 845 | +def test_kv_nvfp4_mla_unit_quantizes_only_the_latent(monkeypatch): |
| 846 | + """The unit fake-quantizes the cached MLA latent; the empty NoPE ``k_pe`` passes through.""" |
| 847 | + monkeypatch.setattr( |
| 848 | + vllm_plugin, |
| 849 | + "create_parallel_state", |
| 850 | + lambda: vllm_plugin.ParallelState(data_parallel_group=None), |
| 851 | + ) |
| 852 | + |
| 853 | + class _TestQuantVLLMMLAAttention(vllm_plugin._QuantVLLMMLAAttention, _NativeMLAAttention): |
| 854 | + pass |
| 855 | + |
| 856 | + mla = _new_attention(_TestQuantVLLMMLAAttention) |
| 857 | + mla._setup() |
| 858 | + set_quantizer_by_cfg( |
| 859 | + mla, |
| 860 | + [{"quantizer_name": "*", "enable": False}, *load_config("configs/ptq/units/kv_nvfp4_mla")], |
| 861 | + ) |
| 862 | + |
| 863 | + assert mla.kv_c_bmm_quantizer.is_enabled |
| 864 | + assert mla.kv_c_bmm_quantizer.amax == 6.0 * 448.0 |
| 865 | + assert not mla.q_bmm_quantizer.is_enabled |
| 866 | + assert not mla.k_pe_bmm_quantizer.is_enabled |
| 867 | + |
| 868 | + mla.to("cuda") |
| 869 | + query = torch.randn(5, 4, 512, device="cuda", dtype=torch.bfloat16) |
| 870 | + kv_c = torch.randn(5, 512, device="cuda", dtype=torch.bfloat16) |
| 871 | + k_pe = torch.empty(5, 1, 0, device="cuda", dtype=torch.bfloat16) # qk_rope_head_dim=0 |
| 872 | + out_query, out_kv_c, out_k_pe = mla(query, kv_c, k_pe) |
| 873 | + assert out_query is query |
| 874 | + assert out_k_pe is k_pe |
| 875 | + assert not torch.equal(out_kv_c, kv_c) |
| 876 | + # NVFP4 with a fixed global scale is idempotent: the output is already on the grid. |
| 877 | + assert torch.equal(mla.kv_c_bmm_quantizer(out_kv_c), out_kv_c) |
| 878 | + |
| 879 | + |
| 880 | +def test_kv_nvfp4_mla_quantizer_replays_in_cuda_graph(): |
| 881 | + """The constant amax is created on the CPU; on the GPU (FakeQuantWorker moves every quantizer |
| 882 | + there before CUDA graph capture) the latent fake quant captures and replays like eager.""" |
| 883 | + layer = torch.nn.Module() |
| 884 | + layer.kv_c_bmm_quantizer = TensorQuantizer() |
| 885 | + set_quantizer_by_cfg(layer, load_config("configs/ptq/units/kv_nvfp4_mla")) |
| 886 | + quantizer = layer.kv_c_bmm_quantizer |
| 887 | + assert quantizer._amax.device.type == "cpu" |
| 888 | + quantizer.to("cuda") |
| 889 | + |
| 890 | + static_in = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16) |
| 891 | + stream = torch.cuda.Stream() |
| 892 | + stream.wait_stream(torch.cuda.current_stream()) |
| 893 | + with torch.cuda.stream(stream): |
| 894 | + quantizer(static_in) # compile the kernel before capture |
| 895 | + torch.cuda.current_stream().wait_stream(stream) |
| 896 | + graph = torch.cuda.CUDAGraph() |
| 897 | + with torch.cuda.graph(graph): |
| 898 | + static_out = quantizer(static_in) |
| 899 | + |
| 900 | + new_in = torch.randn_like(static_in) |
| 901 | + static_in.copy_(new_in) |
| 902 | + graph.replay() |
| 903 | + assert torch.equal(static_out, quantizer(new_in)) |
| 904 | + |
| 905 | + |
837 | 906 | def _quantize_and_summarize(self): |
838 | 907 | """Run on the worker via ``LLM.collective_rpc``. |
839 | 908 |
|
|
0 commit comments