From b5bd7fc9797f044252e6aeff6e165e4e7cfde37b Mon Sep 17 00:00:00 2001 From: Chris Qian Date: Wed, 2 Sep 2026 14:06:55 +0800 Subject: [PATCH] qwen3_5_moe: run lm_head on sampled rows only, fixing 32GB first-prefill OOM The engine samples one row per request (batch_logits = logits[:batch.size]); the remaining rows of the forward window are overlap context. Projecting the whole window through the vocab GEMM allocated M x vocab bf16 -- a default 8192-token chunk at Qwen3.8's 248k vocab is 3.79 GiB, a guaranteed OOM on 32 GB cards at the first prefill of every session -- and spent FLOPs on rows nobody reads. Slice to batch.size rows before the lm_head GEMM. All three call sites already consume exactly this contract: the eager path slices logits[:batch.size], graph capture assigns into buffer.logits[:bs], and the prefill warmup discards the output. Long prefills additionally get a much cheaper lm_head pass. Report: unsloth/Qwen3.8-27B-NVFP4 on RTX 5090D. (cherry picked from commit 1d285471affd3cd8a74a0e46ea5ef97f01e8253a) --- python/freetoken/models/qwen3_5_moe/model.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..763eceb84 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -110,8 +110,15 @@ def __init__(self, config: ModelConfig): super().__init__() def forward(self) -> torch.Tensor: - output = self.model.forward(get_global_ctx().batch.input_ids) - return self.lm_head.forward(output) + ctx = get_global_ctx() + output = self.model.forward(ctx.batch.input_ids) + # Project only the rows the sampler reads. engine.forward_batch slices + # logits[:batch.size] (one row per request, the first rows of the forward + # window); the remaining rows are overlap context. Running the vocab GEMM on + # the whole window allocates M x vocab bf16 (a default 8192-token chunk at + # Qwen3.8's 248k vocab is 3.79 GiB -- a guaranteed OOM on 32 GB cards) + # and spends ~M x vocab x hidden FLOPs whose results nobody consumes. + return self.lm_head.forward(output[: ctx.batch.size]) __all__ = ["Qwen3_5MoEForCausalLM"]