Conversation
# Conflicts: # mstar/engine/resources/position/manager.py
| # | ||
| # TODO: derive this from (backend, dtype, ...) once there is more than one | ||
| # backend to ask, and drop the knob. | ||
| disable_sink_slot: bool = False |
There was a problem hiding this comment.
If this is ever True, padding rows are given NO_SLOT=-1. Decode and convolutional kernel are okay with this, but GDNPrefillWrapper::run calls torch.index_select(state, 0, slots) (resources/linear_attn/wrappers.py:310) and state.index_copy_(0, slots, ...) (:341) which both raise index out of range when given -1.
We should either get rid of this flag as suggested by the comment, or have the prefill wrapper mask out -1.
There was a problem hiding this comment.
I added functionality for the prefill wrapper to mask out -1
| self._cached_plan_output = self._current | ||
| return self._current | ||
|
|
||
| def clear_preplan(self): |
There was a problem hiding this comment.
If we fork in preplan, here we do not clear state, only set dst.has_state = False. This slot may then be reused in a subsequent step. GDNDecodeWrapper or the convolutional block only check has_state (they trust that a slot with has_state=False is all zeroes) when planning, then end up with bad values when they actually run.
There was a problem hiding this comment.
There is no model that forks on this resource, so this is not a problem now. We can either zero in the undo loop, or defer copying.
There was a problem hiding this comment.
Good catch, this is handled by not actually initiating the copy in preplan (the fix to the race condition in #258 (comment))
| "should have allocated it" | ||
| ) | ||
| for tensor in self._blocks.values(): | ||
| tensor[:, dst.index].copy_(tensor[:, src.index]) |
There was a problem hiding this comment.
Can this race?
Worker preplans N+1 after N commit_done (before N's kernels are finished). pre_plan_for_batch runs plan in dedicated stream; the synchronization between that stream and the default stream is on run_forward which has the default stream wait. The planning stream then does not have to wait for step N.
This copy can then be issued while step N replays; however, tensor[:, dst.index] is written in-place by default-stream kernels e.g. GDN index_copy_ and convolutional window shift. There is then a race between the default and planning stream.
This would actually also be possible on the KVManager, although Bagel does not run into this because any pre_forks are in prefill_text which is a PackedCudaGraphConfig which cannot get preplanned.
There was a problem hiding this comment.
Yes this can race; I think the solution is to exempt forks from the pre-plan (and always run them in the actual plan, even if a cached preplan is available). We have a similar policy in the sampler for staging the repetition penalty masks. In that case, the copy would always happen on the same stream. @stephen-dwq thoughts?
There was a problem hiding this comment.
Yes, this makes sense. I think it is the only way to make this work without races. @NSagan271
|
@NSagan271 I will adopt your pool in #257 and port KDA onto it. I think your implementation is better and more consistent architecturally. A few things that my implementation has that yours doesn;t:
As a side, just noticed that the PRs also collide outside the pool (worker.py, submodule_base.py, engine.py, communication/tensors.py, registry.py, docs). Your BatchedModelOutput overlaps my inline-results and coalescing changes, so one side (likely me) has to rebase regardless of the pool decision. |
stephen-dwq
left a comment
stephen-dwq
left a comment
There was a problem hiding this comment.
I think we may have to change the flash-infer floor to 0.6.8 because of the _GDN_DECODE_BF16_STATE_AVAILABLE. I don't think <0.6.8 has a 128x128 that meets this.
Resources, components and model LGTM.
| # they already are. | ||
| v = v.contiguous() | ||
| a = a.contiguous() | ||
| b = b.contiguous() |
There was a problem hiding this comment.
v, a, b are made contiguous but q and k are not. On the prefill path F.normalize happens to return contiguous tensors, but on decode plan.qk_l2norm_in_kernel is True so normalize is skipped and strided views from torch.split go straight into gated_delta_rule_decode_pretranspose, wrong inputs?
Not sure if https://github.com/flashinfer-ai/flashinfer/blob/v0.6.8/flashinfer/gdn_kernels/gdn_decode_bf16_state.py#L2546 would handle in this case so either for all or none..
|
|
||
| def _split_count(batch: int, vocab: int, device: torch.device) -> int: | ||
| """How many chunks to cut the vocab into, or 1 to keep the fused kernel.""" | ||
| del batch, device |
There was a problem hiding this comment.
_split_count deletes batch and device so every vocab over 16k takes the split path (every model vocab >16k). So the fused kernel is not reachable, Fix: use batch in the decision
| """ | ||
| dummy_rids = self.slot_for(lease).dummy_rids | ||
| self._dummy_rows.reset(dummy_rids[real_bs:lease.bucket.bs]) | ||
| self._dummy_rows.reset(dummy_rids[real_bs:lease.bucket.bs], free=True) |
There was a problem hiding this comment.
free=True makes every step re-acquire KV pages for padding rows. The piecewise runner still uses free=False. So when on a near full arena the allocation can fail, maybe we can free per resource?
|
|
||
| self._cu_buffer[: len(cu)].copy_( | ||
| torch.tensor( | ||
| cu, dtype=torch.int32, pin_memory=torch.cuda.is_available() |
There was a problem hiding this comment.
_cu_buffer is int64, the pinned source is int32. copy_ can't DMA across dtypes so it copies to a GPU temporary then converts, adding a kernel + an allocation per step.. We can match the dtypes (int32 would also match query_start_loc)
Claude says: "FlashInfer's SM90 prefill casts cu_seqlens to int64, and only the SM100 path wants int32. On H100 keep the device buffer int64 and make the pinned source int64. Switching to int32 would add a device cast per prefill call."
|
Thanks, the resources and model math lgtm (gates, layouts, RoPE, loaders and the ViT match HF and vLLM with real weights, and single-request greedy output matches HF). I ran it on H100s on marlowe and found two major issues plus a TP alignment problem that I guess the benchmarks did not see because they used --ignore-eos and power-of-two concurrencies. NOTE: There were quite a few issues so instead of leaving comments (which would have been fairly time-consuming), I tried fixing them all in PR #285 . @NSagan271 @vasilevklart @stephen-dwq please take a look, stress test, and approve/comment as needed. |
…s, TP alignment, sampler (#285) * step: is_padding_row on StepContext * recurrent pool: padding and capture rows take no slot * cuda graph runner: keep padding rows' pages resident again * BatchedModelOutput: carry the forward's row order * engine: stamp row_request_ids on collected outputs * worker: map stop-check rows through the forward order * gdn: 32-byte gate pad, zero the pad block * gdn tp: zero the rank's gate pad * sampler: split softmax survives -inf leading blocks * qwen3.5: image-only prompts, 400s for bad input, rgba, top_k * deps: flashinfer >= 0.6.14 * gdn wrappers: pass the pool's pad index to the conv kernels * gdn manager: null slot id for the conv kernels * rope: declare force_double_buffer * numa: only pin when nobody narrowed us * norm: honest RMSNormGated comment * qwen3.5: fix stale capture comment * benchmark: fix Qwen3.5 docstring * docs: qwen3.5 slot sizing * test: padding rows take no pool slots * test: row mapping through the forward order * test: gate pad alignment * test: pad accounting for the 32-byte alignment * test: split softmax with -inf regions * test: conv kernels honour the pad index * test: numa policy cases * test: qwen3.5 prompt edge cases * test: skip vision parity without qwen3_5 in transformers * test: skip the gdn graph test without CUDA * worker: call the row mapper through the class * test: fix the reference shape in the masked-rows test * kv: padding rows run against SINK_PAGE and hold no pages * cuda graph runner: padding rows hold nothing, so free per capture and per step * attn wrappers: room for one sink index per padding row * test: kv padding rows take no pages * padding rows for the GDN prefill kernel: newer flashinfer can't run GDN on cuda 12.8, and the flashinfer 0.6.14 prefill kernel fails when given length-0 rows --------- Co-authored-by: NSagan271 <nsagan@stanford.edu>
|
@merceod I did some testing and merged your branch in; can you approve and merge if it looks good? |
NSagan271
changed the title
…ut (merge fix: main's capture_into_graph with #258's collector)
This branch has not been deployed
4 participants
Qwen3.5 Dense + recurrent state / linear attention resources
Adds the Qwen3.5 dense family and the two engine resources it needed: a pool of fixed-size recurrent state, and linear attention planned over it. This might overlap with #257; @merceod thoughts on how to resolve the two recurrent pool implementations?
MoE variants can be a follow-up PR.
1. Two new resources
Qwen3.5 is a hybrid: most layers are gated DeltaNet (a fixed-size recurrent state), and only a few are full attention over a KV cache.
resources/recurrent/A slot is a fixed number of bytes per layer, held for as long as a request needs it. The pool allows arbitrary number recurrent state geometry.
Like in the KV cache, slot 0 is a sink slot; FlashInfer's bf16 GDN kernel redirects padding rows into slot 0, so that slot is by default not allocated to real requests.
resources/linear_attn/Plans and runs linear attention against a pool.
LinearAttnVariantisGDN(scalar decay per head — Qwen3.5, Qwen3-Next) orKDA(diagonal decay per K channel, e.g., Kimi Linear, GLM-5.3); only GDN is wired up.Also relevant are:
utils/causal_conv1d/(the short conv that precedes the delta rule) andmodel/components/linear_attn.pywith its TP variant.2. The model
mstar/model/qwen3_5/— config, weight loader, the hybrid LLM stack, a ViT tower, and the submodules that bind them to the engine. Five dense sizes registered (qwen3_5_{0.8,2,4,9,27}b); MoE is a follow-up. TP is wired (configs/qwen3_5_9b_tp2.yaml,qwen3_5_27b_tp4.yaml).For improved I2T TTFT and throughput, the model performs one prefill walk per prompt: a multimodal prompt runs a single
prefill_visioncarrying every text span and image.3. Performance
1×H100, closed loop,
--ignore-eos, output 256–4096 tokens, 5 trials/cell, against vLLM 0.29.0 (--max-model-len 8192 --gpu-memory-utilization 0.85):We lose I2T TTFT, which appears to be an API server issue (I'll investigate further, maybe after this PR):
recv → preprocess donefrom the--log-statsoutput is 51.8 ms at conc=1 and 107.8 ms at conc=4, and the entire engine path only takes 28.4 ms.9B, TP2
2×H100,
MSTAR_TP_ASYNC_SCHED=1, 3 trials/cell. Same work as the 4B rows(2304.4 / 2564.9 / 2452.9 output tokens per request):
Same 9B on one GPU (t2t, 512-token outputs), where we win as the 4B does:
So it seems like the issue is a TP problem and not a model problem; that'll be a separate investigation from this PR.
4. Engine-Level Performance Changes
forward_batchedto be a new dataclass,BatchedModelOutput(with backward compatibility maintained for other models). This dataclass allows the output offorward_batched, and the buffer that is D2H'd forcheck_stopto be batched across requests (instead of cloning or transferring tiny per-request buffers). This usually doesn't matter for performance, but Qwen 3.5-4B decode steps only take 4-5ms, so they are close to becoming host-bound.I (+Claude) also added some more timing instrumentation to
worker.py; it can be removed if it's too much clutter.Checklist
ruff check .passes