ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

qwen3_5 flex: fix bs=max_batch_size serialization regression

Browse source 
PageTable reserves batch_idx=0 as a no-op slot, so passing
``max_batch_size=N`` to ``PageTable.__init__`` yields only N-1
user-allocatable slots. When a user requests bs=N concurrent sequences,
the Nth sequence gets stranded in the ``waiting`` queue until the first
N-1 sequences finish, then runs serially. For bs=8 × 64 tokens this
split the decode into 63 iterations at B=7 followed by 63 at B=1 —
throughput collapsed from a predicted ~800 tok/s to 143 tok/s.

Fix: bump the internal PageTable capacity by 1 so the user's
max_batch_size maps to that many concurrent slots. All batch-indexed
buffers (input_pos_buffer, block_mask_logical, _linear_conv_states,
_linear_recurrent_states, _linear_caches) size accordingly. Slot 0
remains internally reserved; user-facing ``self.max_batch_size``
unchanged.

Throughput (Qwen3.5-4B, 128 new tokens, post-warmup):

| bs | pre-fix tok/s | post-fix tok/s | change |
|---:|---:|---:|---:|
|  1 | 191.2 | 191.5 | same |
|  2 | 338.1 | 339.4 | same |
|  4 | 548.6 | 542.0 | same |
|  8 | 143.5 | **804.7** | **5.6x** |

MoE Qwen3.6-35B-A3B bs=8 × 64: 38.2 → **305.2 tok/s** (8x).

Parity at bs=4 × 16 tokens: was 3/4 exact (P1 matched 13/16), now 4/4
exact — the serialization also exposed an extra bf16-drift path.
This commit is contained in:
danielhanchen 2026-04-23 23:44:46 +00:00
commit 2b8a88aa76
1 changed files with 14 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth/inference/flex_qwen3_5.py
View file

@ -466,6 +466,14 @@ class FlexQwen3_5Inference:
self.base_model = base_model
self.peft_model = peft_model
self.max_batch_size = max_batch_size
# PageTable reserves batch_idx=0 as a no-op slot, so the first
# ``max_batch_size`` user-allocatable slots are 1..max_batch_size.
# Bump the page table's capacity by 1 so the user actually gets
# ``max_batch_size`` concurrent sequences (otherwise bs=N requests
# serialise into bs=N-1 + bs=1, collapsing throughput — see
# commit on this file for the bs=8 regression diagnosis).
page_table_cap = max_batch_size + 1
self._page_table_cap = page_table_cap
self.max_seq_length = max_seq_length
self.page_size = page_size
self.max_new_tokens = max_new_tokens
@ -505,16 +513,16 @@ class FlexQwen3_5Inference:
self.page_table = PageTable(
n_pages=n_pages,
page_size=page_size,
max_batch_size=max_batch_size,
max_batch_size=page_table_cap,
device=self.device.type,
)
_patch_qwen3_5_full_attn_forwards(self.text_model, self.page_table)
self.input_pos_buffer = torch.zeros(
max_batch_size, dtype=torch.int32, device=self.device,
page_table_cap, dtype=torch.int32, device=self.device,
)
self.block_mask_logical = self.page_table.create_causal_blockmask(
B=max_batch_size, L=max_seq_length,
B=page_table_cap, L=max_seq_length,
)
# Pre-allocate per-slot LinearAttention caches. Each sequence's
@ -522,7 +530,7 @@ class FlexQwen3_5Inference:
# we reset it via .reset() at prefill so the conv/recurrent
# states start clean.
self._linear_caches = [
self._build_linear_cache() for _ in range(max_batch_size)
self._build_linear_cache() for _ in range(page_table_cap)
]
# Static state buffers for CUDA-graph-captured batched decode.
@ -541,11 +549,11 @@ class FlexQwen3_5Inference:
for layer_idx in self._linear_layer_indices:
la = self.text_model.layers[layer_idx].linear_attn
self._linear_conv_states[layer_idx] = torch.zeros(
max_batch_size, la.conv_dim, la.conv_kernel_size,
page_table_cap, la.conv_dim, la.conv_kernel_size,
dtype=la.conv1d.weight.dtype, device=self.device,
)
self._linear_recurrent_states[layer_idx] = torch.zeros(
max_batch_size, la.num_v_heads, la.head_k_dim, la.head_v_dim,
page_table_cap, la.num_v_heads, la.head_k_dim, la.head_v_dim,
dtype=torch.float32, device=self.device,
)
try: