qwen3_5_moe: run lm_head on sampled rows only (fixes 32 GB first-prefill OOM)

[chrisqianz]

Symptom

Serving any qwen3_5_moe dense model (vocab 248,320 — e.g. Qwen3.6-27B-NVFP4, unsloth/Qwen3.8-27B-NVFP4) on a 32 GB card dies on the first request of every session:

File ".../freetoken/models/qwen3_5_moe/model.py", line 124, in forward
    return self.lm_head.forward(output)
File ".../freetoken/kernel/triton/fp8_pertensor_linear.py", line 167, in _gemm
    out = torch.empty((M, N), dtype=compute, device=a.device)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.79 GiB.
GPU 0 has a total capacity of 31.36 GiB of which 2.71 GiB is free.

The worker exits and the supervisor stops the API (Backend worker is gone and cannot be restarted — the #20 symptom family). A shorter first prompt "works": the fatal allocation scales with the prefill window, so it fires exactly when the first request's chunk fills.

Root cause

The eager model forward runs the vocab GEMM over the entire forward window, but the engine only ever consumes one row per request:

• engine.py:933 — batch_logits = logits[: batch.size] (the rest is overlap context, discarded);
• graph.py:175 — CUDA-graph capture already assigns self.buffer.logits[:bs] = model.forward(), i.e. the graph path is built around a bs-row output;
• the prefill warmup (engine.py:985) discards the output entirely.

So for a default 8192-token chunk the eager path transiently allocates 8194 × 248,320 × bf16 = 3.79 GiB — exactly the number in the traceback (8194 = max_extend_tokens 8192 + sampled rows) — and spends ~M × vocab × hidden FLOPs on rows nobody reads. On a 32 GB card holding ~21 GB of weights + the KV pool, the first prefill of every session deterministically OOMs.

Fix

Slice to the sampled rows before the lm_head GEMM:

return self.lm_head.forward(output[: ctx.batch.size])

qwen3_5_moe's forward window already carries the sampled rows first (that ordering is exactly what logits[:batch.size] assumes), so no gather is needed. This is the same reduction deepseek_v4 already ships — F.linear(h[0, last_indices], self.head) # [B, vocab] — just expressed as a slice for this family's layout.

Transient buffer: 3.79 GiB → batch.size × vocab × 2B (~0.5 MB at concurrency 1); long-prefill lm_head cost collapses from a full-window GEMM to a single-row GEMM.

Scope / prior art

Present since the initial release (3af9d90); surfaced while serving unsloth/Qwen3.8-27B-NVFP4 on an RTX 5090D (32 GB) in the #208 discussion, but #208 never touched forward() — this is a pre-existing engine bug that also threatens Qwen3.6-27B-NVFP4 on the same hardware. Distinct from the other M-proportional prefill OOM reports: #171 (dsv4 pool-derived chunk budget), #172 (dsv4 sliding-window re-prefill), #110 (MoE expert workspace). Several sibling families (llama, qwen3, glm4_moe, … ) still return self.lm_head.forward(output) over the full window and could adopt the same one-liner; this PR keeps to the family verified end-to-end.

Verification

• RTX 5090D 32 GB, unsloth/Qwen3.8-27B-NVFP4 (fp8 lm_head): a 9,500-token prefill that previously died at exactly this allocation returns HTTP 200 in 2.7 s (prompt_tokens=9500), server stays healthy; chat conversations run normally.
• Decode path unaffected: decode batches are already batch.size rows (slice is a no-op), CUDA graphs capture with the same shapes as before.
• tests/models/test_qwen3_5_moe_config.py + test_qwen3_5_moe_weight.py: 29 passed.