From 79a41d9197f9106745a38cc891b619ddec1c0e14 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 07:29:38 +0000 Subject: [PATCH] Save ~5-8 GiB GPU VRAM: move dummy weight to CPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CuTeDSL kernel never reads layer.weight — it uses the runner's pre-processed fp4/sf/gs tensors. The dummy BF16 weight exists only for vLLM model introspection. Moving it to CPU saves massive VRAM: - q_b_proj alone: 65536*1536*2 = 192 MiB on GPU → ~0 MiB - All layers combined: ~5-8 GiB saved This should fix the KV cache OOM (needed 10.28 GiB, had 9.36 GiB). --- vllm/cutedsl_quant_method.py | 5 ++++- vllm/kernels/linear/nvfp4/cutedsl.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/vllm/cutedsl_quant_method.py b/vllm/cutedsl_quant_method.py index d1c43cd7..85780781 100644 --- a/vllm/cutedsl_quant_method.py +++ b/vllm/cutedsl_quant_method.py @@ -112,9 +112,12 @@ class CuTeDSLNvfp4Method(LinearMethodBase): torch.cuda.empty_cache() # Replace weight with dummy BF16 (needed by vLLM module introspection) + # Replace weight with a minimal CPU dummy (vLLM introspection expects + # layer.weight to exist, but our kernel never uses it). Keeping it on + # GPU would waste ~5-8 GiB across all layers. layer.weight = torch.nn.Parameter( torch.zeros(out_features, in_features, dtype=torch.bfloat16, - device=device), + device="cpu"), requires_grad=False, ) diff --git a/vllm/kernels/linear/nvfp4/cutedsl.py b/vllm/kernels/linear/nvfp4/cutedsl.py index 7e3718ee..fb7b7519 100644 --- a/vllm/kernels/linear/nvfp4/cutedsl.py +++ b/vllm/kernels/linear/nvfp4/cutedsl.py @@ -99,9 +99,13 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): layer._cutedsl_runner_id = register_runner(runner) layer._cutedsl_out_features = out_features + # Replace weight with a minimal CPU dummy (vLLM introspection expects + # layer.weight to exist, but our kernel never uses it). Keeping it on + # GPU would waste ~5-8 GiB across all layers (e.g. q_b_proj alone + # is 192 MiB of zeros). layer.weight = torch.nn.Parameter( torch.zeros(out_features, in_features, dtype=torch.bfloat16, - device=device), + device="cpu"), requires_grad=False, )