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:
2026-05-19 07:29:38 +00:00
parent cebc586014
commit 79a41d9197
2 changed files with 9 additions and 2 deletions

View File

@@ -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,
)

View File

@@ -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,
)