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.

[gdevenyi]

Data point from trying the same one-line slice on the sibling qwen4_exp model (Qwen4ExpForCausalLM.forward has the identical self.lm_head.forward(self.model.forward(...)) shape, and the 248k vocab makes the same 3.79 GiB transient at an 8192-token chunk):

hidden = self.model.forward(batch.input_ids, batch)
return self.lm_head.forward(hidden[: batch.size])

On RadixArk/Qwen3.8-Flash-Next-NVFP4 (2x RTX 6000 Ada, main 86214a9, --moe-backend offload --num-tokens 262144 --max-running-requests 8 --cuda-graph-max-bs 8) the server boots, warms up and captures graphs, then both at TP=1 and at TP=2 the scheduler dies on the first real request:

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:111: operator(): block: [4,0,0], thread: [64,0,0]
Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
...
  File ".../freetoken/engine/engine.py", line 944, in forward_batch
    next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32)

(the device assert surfaces at the sampler's first sync; without the slice the same run serves fine). So the "sampled rows are the first batch.size rows of the window" contract that engine.forward_batch relies on for its own logits[: batch.size] slice does not appear to hold for every model's eager window, or something downstream still reads rows past batch.size; either way the slice needs to be gated on the family, or better, the engine should hand the model the row indices it will sample so the model can gather them instead of assuming a layout. I have not chased which consumer indexes past batch.size; posting in case it saves you the surprise when extending this to qwen4_exp.