Save ~5-8 GiB GPU VRAM: move dummy weight to CPU
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).
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user