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.

[gdevenyi]

This is not a qwen3_5_moe bug. The same call site exists unfixed in three more families, and the engine slice you are keying on is shared by all of them:

python/freetoken/engine/engine.py:1000:  batch_logits = logits[: batch.size]

python/freetoken/models/qwen3_5_moe/model.py:113   return self.lm_head.forward(output)     <- this PR
python/freetoken/models/qwen4_exp/model.py:297     return self.lm_head.forward(...)
python/freetoken/models/glm5_next/model.py:181     return self.lm_head.forward(output)
python/freetoken/models/llama/model.py:86          logits = self.lm_head.forward(output)

Every one projects the whole forward window and hands engine.forward_batch a tensor it immediately throws all but batch.size rows of.

Concretely, on a deployment that is not 32 GB. We serve Qwen3.8-Flash-Next (qwen4_exp, the same 248,320 vocab) on 2 x RTX 6000 Ada at TP=2, --max-extend-tokens 8192. A full-width chunk allocates

8192 x 248320 x 2 B = 4.07 GiB

of bf16 logits to read at most 8 rows of, and runs a vocab GEMM 1024x larger than the sampler needs — on the TTFT path, every chunk. We had not hit the OOM your report opens with because the allocator has expandable_segments on and there is headroom, but the prefill warmup at length 4096 alone is a 2.03 GiB transient against 2.4 GiB free after CUDA-graph capture. That is not comfortable, and on a 262k context it is roughly 32 chunks of GEMM nobody reads.

So the 32 GB OOM is the symptom that made this visible, not the boundary of the problem. Even where it fits, it is paid in TTFT.

I have applied your one-liner to qwen4_exp on our branch and it is CPU-clean; I will post the measured TTFT delta at 1k and at a long prompt once it has a GPU window, which should give you a throughput argument to go with the memory one.

Two suggestions for this PR. First, it is worth widening to all four call sites — it is the same line, the comment you wrote applies verbatim, and leaving three behind means the next person rediscovers it on a different card. Second, a test would pin it cheaply. Slicing before a row-wise projection is exactly equal to slicing after, so the property is assertable without a GPU:

full_then_slice   = head.forward(hidden)[:size]
slice_then_project = head.forward(hidden[:size])
torch.testing.assert_close(slice_then_project, full_then_slice)

plus a source guard on the call site so it cannot silently regress in a rebase. I have both in tests/models/qwen4_exp/test_lmhead_rows.py on our branch and they are yours if useful.

One thing worth stating explicitly in the comment, because it is the part a reviewer has to trust: the rows kept are the first batch.size rows of the window, not the last token of each sequence. That reads wrong at a glance. It is right because engine.forward_batch already slices exactly that way — but a reader who does not know that will flag it, so it is worth the half sentence.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

Correction to my previous comment, and it matters before you widen this PR.

I said I had applied your one-liner to qwen4_exp and it was CPU-clean, and that I would post the TTFT delta. It is not clean on hardware. At TP=2 the server dies during _warmup_prefill:

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:111: operator():
  Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
torch.AcceleratorError: CUDA error: device-side assert triggered
  File "freetoken/engine/engine.py", line 1161, in _warmup_prefill
    dummy_row.fill_(dummy_slot)      <- where the async fault surfaces, not where it is

The only delta from a branch that starts and serves is your one line. I have reverted it on our deploy branch.

What I can tell you about the shape of it:

• CUDA-graph capture passes. So batch.size does equal the captured bs on the decode path, and the premise holds there.
• It is the prefill path. The warmup batch is one request with length tokens (engine.py builds reqs=[warm_req] with input_ids of length), so batch.size is 1 while the window is up to 4096 rows. That is the widest gap between batch.size and the window, which is exactly where the slice bites hardest.
• The dependency is not visible by reading. Every indexing kernel in that path — the PLE n-gram table gather, the KV store, hybrid_swa_pool.translate_loc_from_full_to_swa — runs before lm_head, so none of them can be reading the tensor I truncated. Something less direct is going on, and naming it needs CUDA_LAUNCH_BLOCKING=1 on the box.

I do not yet know whether this is specific to qwen4_exp (GDN state slots, hyper-connections, the PLE layer, or the vocab-parallel FP8 head we run) or whether qwen3_5_moe has a latent version of the same thing that your single-card TP=1 testing would not reach. That is the question worth answering before this lands in four families. If your validation was TP=1, it would be worth one TP=2 run before widening.

Also worth noting for your own PR: the test I wrote alongside the change passes. It asserts that slicing before a row-wise projection equals slicing after, and guards the call site. Both are true. Neither is the property that breaks. So "row-wise projection commutes with slicing" is not a sufficient argument for this change, and I was wrong to treat it as one — the engine's prefill contract, not the linear algebra, is what needs pinning.

The memory and TTFT argument for the change stands and is worth chasing: 4.07 GiB of bf16 logits per 8192-token chunk at a 248k vocab, to read at most batch.size rows. I will report back if I get a kernel name.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

I have the kernel now, and the answer inverts my earlier comments. Please do not widen this PR — for any family whose head is ParallelLMHead, the slice is an out-of-bounds bug, and the waste it targets does not exist.

Run under CUDA_LAUNCH_BLOCKING=1, the assert resolves to one line:

File "freetoken/engine/engine.py", line 1152, in _warmup_prefill
File "freetoken/models/qwen4_exp/model.py", line 303, in forward      <- the sliced call
File "freetoken/layers/embedding.py", line 142, in forward            <- x = x[indices]
torch.AcceleratorError: CUDA error: device-side assert triggered

layers/embedding.py:142 is inside ParallelLMHead.forward:

def forward(self, x: torch.Tensor) -> torch.Tensor:
    bs = batch.size
    if batch.is_prefill:
        indices = batch.attn_metadata.get_last_indices(bs)
        x = x[indices].contiguous()        # line 142
    logits = self._logits(x)

The head already does the reduction, and it does it before the GEMM. get_last_indices(bs) is cu_seqlens_q[1:1+bs] - 1 — the last token of each sequence. So _logits never sees the full window: there is no M x vocab allocation to reclaim and no oversized vocab GEMM. My 4.07 GiB figure was wrong, and I withdraw it.

Two consequences.

1. Slicing first breaks it. output[: batch.size] hands the head a bs-row tensor, which the head then indexes with last-token positions — for a 4096-token warmup, index ~4095 into 1 row. Hence the assert. This is why my "row-wise projection commutes with slicing" argument was not just insufficient but beside the point: the head is not a bare projection.

2. This is not a qwen4_exp quirk. On current main, all four call sites I listed use the same head:

qwen3_5_moe/model.py:101   self.lm_head = ParallelLMHead(
qwen4_exp/model.py:145     Fp8ParallelLMHead(...)   # subclasses ParallelLMHead
glm5_next/model.py:161     self.lm_head = ParallelLMHead(
llama/model.py:74          self.lm_head = ParallelLMHead(

and git log -S get_last_indices -- python/freetoken/layers/embedding.py dates that reduction to 3af9d90, the initial open-source release. It predates this PR, so it was in place when you opened it.

Which leaves your OOM genuinely unexplained, and worth a second look. Your traceback allocates 3.79 GiB at fp8_pertensor_linear.py:167 — that is M=8192 at a 248,320 vocab, i.e. the full window reaching _logits. If ParallelLMHead.forward had run its prefill branch, M would have been batch.size. So on your build the reduction did not happen. Candidates worth checking before changing the call site: whether the dense qwen3_5_moe path reaches ParallelLMHead.forward at all or calls _logits directly; whether batch.is_prefill is true on that first request; and whether your attention backend implements get_last_indices. The fix probably belongs wherever that branch is being skipped, not in the four model files.

Sorry for the noise across three comments — the first two were wrong in opposite directions, and this one has the traceback behind it. Reverted on our side; our branch is back to a state that starts and serves.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[chrisqianz]

Closing as superseded — @gdevenyi's final comment is correct, and the real root cause was on our side.

On main this call site was never a waste: ParallelLMHead.forward reduces to the last-token rows via get_last_indices(bs) before the GEMM on prefill (present since 3af9d90, implemented across the attention backends). The 3.79 GiB transient in my traceback came from the sibling PR #208, whose first commit routed the unsloth export's FP8 lm_head to a plain replicated linear (Fp8PerTensorLinear) — a head class without the prefill reduction. The head swap removed the reduction; the full-window GEMM followed. I attributed to upstream what my own branch introduced, and the earlier "widen to four families" suggestion in this thread is withdrawn — slicing at those call sites is the out-of-bounds bug gdevenyi traced, not a fix.

The remaining question — quantized lm_heads keeping the reduction — is solved structurally by #418 + #438: heads are always ParallelLMHead, with the quant kind resolved per prefix through QuantConfig/quant_method, so the reduce survives for every storage the checkpoint uses. The #438 e2e matrix passes unsloth/Qwen3.8-27B-NVFP4 (the checkpoint behind my original OOM) and the fp8-lm_head variants.

Thanks for the CUDA_LAUNCH_BLOCKING trace that turned a wrong turn into the right answer.

[chrisqianz]

Superseded by #418/#438 — see comment above.