Repository navigation
Settle a step's shape once, and fix the three defects that exposed - #2088
Merged
Merged
Conversation
Every consumer used to re-derive its own width. `bs * max_seqlen_q` was computed in nine places, the flat replay bucket was searched twice with different operands, and the two searches came to disagree -- a step could report a captured graph at a width nothing had recorded. `ForwardMode.decide` is now the only thing that answers the question, and it answers it in the only two units the engine has: `running_bs` for everything per-sequence, `running_tokens` for everything per-row. `scheduled_*` stays this rank's own count, and is what a result slices back to. The batch is one rule on every kind of step: the DP reduction, rounded onto the capture ladder. It used to branch four ways on how the step dispatched, and two of those arms handed back a per-rank count -- so a graph key, a draft pass and an attention plan could each be holding a different width, and every consumer had to know which kind it had. Dispatch is now asked separately: `use_cudagraph` says whether the target replays and sizes nothing. The two units are agreed across DP to different extents, deliberately. `running_bs` always is. `running_tokens` is agreed exactly when `running_tokens_are_unified`, because off it MoE takes the variable-length gather, whose contract is that the rows it was handed are the rows this rank scheduled -- an over-wide per-sequence array costs a few sentinel rows, an over-wide height is a contract break. `_running_tokens` searches the ladder with `searchsorted` over an ascending int32 array built where the ladder is finalised. The list it used to scan is sorted descending for the length of `capture_cudagraph`, so a first-match scan read the largest size for that whole window and said nothing about it. Also here: - The draft pass reports the rows it RAN, not the rows that were real. It pads its own inputs and computes every one of them, so reporting `scheduled_bs * num_draft` against a `running_bs * num_draft` tensor was a contract break, not a conservative choice; `dp_gather_hidden_and_router` asserts the two agree and aborted 8/8 ranks on the first decode step whenever a batch had no recording to replay (`ATOM_DRAFT_CUDAGRAPH=0`). - `assert_shape_contract` covers the padded step, which is the only one that can be wrong. It used to return early whenever `running_bs != scheduled_bs`, exempting exactly that case; a mismatch then reached the model as lost accuracy rather than as a failure here. - The prefill draft pass reaches a warmed shape: prefill rounds onto the same ladder as a decode step, and its per-sequence tail is padded to match. - One name for this rank's sequence count. `real_bs` and `bs` were two spellings of `Context.scheduled_bs` across the drafters; 47 sites. - `ScheduledBatch` no longer carries `max_seqlen_q`, `dspark_dp_bs` or `dynamic_num_tokens_pad`. A scheduler object was holding numbers the DP group settled after it was built. V4-Flash-DSpark tp8 + DPA, k=7, fp8 KV, GSM8K 1319 x3: 0.9492 / 0.9515 / 0.9477. Unit suite 4712 passed.
… own A rank idling behind someone else's long chunked prefill faulted all eight ranks on V4-Flash-DSpark at 100k/c50: the target's embedding gather took an index outside the vocab. Three defects on the draft path, in the order found. 1. `DraftGraph.run` returned the recording's own tensors. A replay hands back the tensors the capture allocated. Those live in the pool every captured size shares, and the pool packs a later size's capture into memory the earlier ones released -- so replaying an earlier size rewrites addresses a later one's live output now holds. Holding a Python reference does not stop that, which is the whole difference from the eager allocator. DSpark's draft ids are the only value in the engine that outlives the step that made it: deferred output feeds them back as the next decode step's `input_ids` without a D2H. Read a step later they were another size's activations, float bit patterns arriving as vocab indices. Shown without the fault: stamp a sentinel into the bs=2 recording's output, replay the bs=1 graph, and bs=1's activations come back -- while the first-captured size is never dirtied. `run` now copies a replay's result out of the pool; the eager path allocates normally and is untouched. 2. A DP-sync dummy passed for a carried-over request. `dummy_execution` fabricates its sequence with a fixed id, so a second dummy matched the first in `get_token_locations` and read that "request's" anchor and drafts out of `prev_token_ids` / `draft_token_ids`. Measured: `src=[0]` on a dummy batch, whose seven draft slots then held 968440611 and similar. A fabricated batch carries nothing over, and the early return also skips the per-request loop -- most of the steps under DP attention are dummies. It still becomes `prev_batch` on purpose: reporting lags a step, so the dummy is what flushes the last real step's output to the host, and its own token lands on the key the deferred flag overwrites. 3. The draft pass reported the target's counts as its own. `_publish_draft_shape` now states both units honestly, declares the pass uniform across DP and not a prefill whatever the target just did, and fills the DP table locally rather than all_reducing it: `running_bs` came out of `decide`'s one collective and the width is config, so there is nothing to ask. `DPMetadata.make`'s assert cannot catch a wrong table -- it checks this rank's own entry, which a uniform fill satisfies -- so a probe compared it against the all_reduced one on every pass and found no disagreement across 8 ranks, over a 100k/c50 benchmark and a full GSM8K. Verified on V4-Flash-DSpark, tp8 + dp-attention, fp8 KV / fp4 indexer, --num-speculative-tokens 7: 100k/10/c50 drains clean, 106,043 tok/s (was an 8-rank abort) GSM8K 1319 0.9522 flexible / 0.9522 strict, fewshot 3 acceptance 48.6%, 4.40 tokens per forward unit tests 4719 passed, 0 failed The local DP table was suspected of costing 0.9pp and does not: n=3 each way, fully interleaved (0.9477/0.9515/0.9522 local, 0.9454/0.9522/0.9522 all_reduced). Single GSM8K runs on this model span 1.5pp.
`_publish_draft_shape` asserted two things on behalf of every drafter that
called it: that the pass was not a prefill, and that every DP rank was
running the same height. Both hold for the DSpark block pass, which is
`running_bs * num_draft` -- `decide`'s reduction times a config width. Only
one of the three call sites is that pass.
Eagle's first draft step runs the target's own token stream against the
target's still-prefill metadata. Clearing `is_prefill` there sends the MLA
layers, the indexer and PCP down their decode branches on a pass that is a
prefill, and declaring that height uniform states an equality no reduction
checked -- `input_ids.shape[0]` is this rank's `total_tokens_num`. Its
mid-steps are uniform only when a recording exists to widen to; without one
`running_bs` falls back to `scheduled_bs`. Kimi-K3's separate-draft path
publishes `scheduled_bs * T`, also rank-local, and nothing rules it out
under DP: `--enable-dp-attention` folds TP into the DP size.
`make` asserts only THIS rank's entry, which a uniform fill satisfies by
construction, so a wrong claim survives to the fixed-size gather in
`dp_gather_hidden_and_router` and lands as a mismatched collective.
So the seam states neither on anyone's behalf:
- `running_tokens_are_unified` is a required keyword. DSpark's block pass
passes True, Kimi-K3 False, Eagle `i > 0 and self.step is not None`.
- `is_prefill` goes back to being the caller's write. Two of the three
already had one; the block pass never did, and needs one now that it
reports `scheduled_bs * num_draft` -- `pad_for_all_gather` checks rows
against the scheduled count when the flag is True and would assert on
the first widened batch.
- `DPMetadata.make` takes `unified` and owns both ways to a table, so the
answer travels instead of the tensor it implies. Supplying a table and
claiming uniformity is now rejected rather than silently ordered.
- `warmup_draft_graphs` states the pair it bakes rather than inheriting
two defaults that happen to agree.
Eagle and Kimi-K3 are restored to their pre-9a3313953 behaviour; the DSpark
block pass is unchanged, only its writer moved.
Verified on V4-Flash-DSpark, tp8 + dp-attention, fp8 KV / fp4 indexer,
--num-speculative-tokens 7 -- the one configuration here that reaches the DP
branch. GSM8K 1319, fewshot 3, server restarted between runs:
0.9469 / 0.9469 / 0.9538 mean 0.9492, acceptance 47.6%, 4.37 tok/fwd
against 0.9505 (local table) and 0.9499 (all_reduced) on 9a3313953, whose
sample ranges this one overlaps in both directions. Single runs on this model
span 1.5pp. Eagle and Kimi-K3 have no recipe on this box and were not run.
`DPMetadata.make` needs no aiter, so tests/test_dp_metadata.py runs in CI --
the drafter-level tests cannot, since spec_decode/drafter.py imports aiter at
module load.
unit tests 4725 passed, 0 failed
Rounding every step onto the capture ladder gave prefill a `running_bs` above its own batch, and `prepare_prefill` pads `cu_seqlens_q` out to it so a following draft pass reaches a warmed shape. The LM head slices prefill rows by that array, so `compute_logits` returns `running_bs` rows -- while `prepare_sample` sizes `temperatures` / `top_ks` / `top_ps` at `batch.total_seqs_num`, because those describe requests. `postprocess` cut the SAMPLED IDS back to the scheduled batch, one step too late. Three prefill sequences on the default ladder run four rows, so any request setting `top_p < 1` or `top_k != -1` on a non-greedy batch reached `scaled_logits = logits / temperatures.unsqueeze(-1)` as `[4, V] / [3, 1]` and raised. The logprob path was already broken the same way and stayed broken after that cut: `log_probs.gather(-1, sampled_tokens.unsqueeze(-1))` gathered `[running_bs, V]` by `[scheduled_bs, 1]`. Both come from one uncut tensor, so the fix is to cut it: `logits[:bs]` before the sampler rather than `sampled_tokens[:bs]` after it. Greedy batches took `argmax` and temperature-only ones over-read the parameter buffer silently, which is why GSM8K and the benchmarks never saw it. Also here, from the same review pass: - `_stage` grows `pad_to`. Both V4 state-slot publishers had hand-inlined its body to get a padded width, dropping its buffer-capacity and dtype assertions along the way. - The packed DP wire loses `any_dummy`. Nothing has read it since it was added, its own comment describes an AND-reduce over an `.any()`, and the layout was being renumbered anyway -- a block drafter now appends one int32 instead of two. - `_running_tokens` drops a `batch` parameter it never read (and a `cls`), `_local_tbo_eligibility` an argument with one caller and one value, whose `np.asarray` now runs below the `enable_tbo` early return rather than on every step. `capture_sizes_np` stops sorting an already-sorted list. - `_dspark_decode_replay` is deleted: its only reader, `_piecewise_replay_shape`, went with the shape refactor and the per-step write survived it. - Dangling names: two comments cited `ScheduledBatch.max_seqlen_q`, which that refactor deleted; three cited `_refresh_dp_metadata`, renamed before this branch began -- one of them a `raising=False` monkeypatch that therefore neutralised nothing. `is_captured` credited `warmup` with an assert it does not contain, a state-slot comment cited `_populate_state_slot_mappings` as evidence for the opposite of what that function documents, and the call site of `assert_shape_contract` still said it skips cudagraph steps -- covering those is what the refactor changed it to do. `docs/environment_variables.md` still documented the EPLB capture opt-out that went away with `DraftGraph.pads`. - Comment volume, where the same fact was written more than once: the draft loop's eight-line publish carried 27 lines with one sentence in it three times; `assert_shape_contract` closed with eight lines about two assertions deliberately not written; the ladder's sortedness was explained three times, two of them disagreeing. Net 55 lines. unit tests 4729 passed, 0 failed
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
`tests/test_postprocess_width.py` aborted collection on the non-GPU runner,
which is not one failed test -- pytest exits 2 and the other 4700 never run.
`pytest.importorskip("aiter")` only skips on ModuleNotFoundError, and on that
runner `aiter` is importable by the time this file is reached:
`tests/test_pd_pp.py` does `sys.modules.setdefault("aiter", ModuleType("aiter"))`
at module scope and never restores it, so every module collected after it
(pd < po) sees the name resolve and dies reading a symbol off it --
`cannot import name 'destroy_dist_env' from 'aiter' (unknown location)`.
`tests/import_guard.py` already documents this shape as how the non-GPU job
actually fails.
Naming the module the test actually needs catches both shapes, and `exc_type`
is what makes the second one a skip:
mod = pytest.importorskip(
"atom.model_engine.model_runner", exc_type=ImportError, reason=...
)
`test_dspark.py` and `test_mtp_deferred_status_queue.py` carry the same
name-only guard and survive today only by sorting before `test_pd_pp.py`.
Alphabetical order is not a guard, so they move to the same idiom.
The probe that was supposed to catch this modelled `aiter` ABSENT, where a
name-only guard skips correctly -- so it passed the file that broke CI.
`/app/logs_claude/check_tests_import_without_aiter.py` now runs an ABSENT and
a HOLLOW pass; HOLLOW reproduces the runner's error verbatim, and the whole
tests/ tree is clean in both.
The stub leak in `test_pd_pp.py` is left alone: restoring `sys.modules` there
is a change to a PD test whose fixtures are not examined here, and no file
depends on the leak.
unit tests 4731 passed, 0 failed
…eeds it
EagleProposer declared a DraftGraph with no epilogue, so a mid-step replayed
the backbone and then ran `compute_draft_ids` eager -- an LM head, a TP
all_gather and an argmax, once per draft step, outside the recording.
Two things had to change, and the first is why this was never done:
`compute_argmax_token` reached `all_gather` without `use_custom`, whose
default is False, so it took RCCL. `graph_capture()` arms the custom
implementation and only that one, which is what `logits_in_graph =
world_size == 1` has meant since `0cad4aabe` ("NCCL all_gather which is not
graph-capturable on HIP when TP > 1"). The logits path two functions up
already passes `envs.ATOM_USE_CUSTOM_ALL_GATHER`; the argmax path did not.
With it, capture succeeds at tp8 -- that rationale was stale, not wrong.
The epilogue then returns `(hidden_states, ids)`, both: the next mid-step
reads the hidden states and they exist only inside the recording. It runs the
head over every row rather than `[:scheduled_bs]`, since a capture bakes the
length it was made at; the caller slices on the way out.
Measured on V4-Pro MTP3, tp8, same 100k/10/c50 shape both ways, rank_0
`propose_eagle` spans -- host kernel launches left outside the graph:
mid-step 1 18.29 -> 9.53 (the rest is prepare_mtp_decode at the i==1
metadata boundary, not the head)
mid-step 2 10.31 -> 1.51
GSM8K 1319 0.9515 flexible, was 0.9545
acceptance 64.47% / 66.11%, was 64.27% / 66.24%; 2.93 / 2.98 tok/fwd,
unchanged -- the probe that would see a token differ
Throughput is NOT measured. Both arms are single-shot and profiled, and step
0's span moved 19.1 -> 9.6 ms between the two runs while nothing touched it
(step 0 never enters a graph), so this shape cannot size a sub-percent
effect. A decode-dominated shape with alternating arms would.
`_owned` clones both tuple elements on replay. Neither outlives a replay here
-- `stage()` copies the hidden states before the next one, and the ids land in
`draft_token_ids` immediately -- so the hidden-state clone is waste, ~688 KB
per mid-step at bs=48. Left alone: teaching the seam which output to skip is a
knob on the seam that returned pool storage three commits ago, and the cost is
already inside the launch counts above.
No unit test: `EagleProposer.propose` has no harness, and asserting on the
one-line `_step_head` would not reach the property (the graph returns both,
the loop slices the ids). GPU evidence only.
unit tests 4731 passed, 0 failed
…ey describe
`CommonAttentionBuilder.prepare_prefill` pads `cu_seqlens_q`, `cu_seqlens_k`
and `context_lens` out to `running_bs` and uploads them at that width, so a
draft pass that follows the prefill sees clean tail rows. Every consumer that
still counted requests off one of them broke the moment a batch's sequence
count fell off the capture ladder -- which is most of them:
RuntimeError: The expanded size of the tensor (21) must match the existing
size (32) at non-singleton dimension 0
Six accuracy jobs failed that way. The heavy matrix had been SKIPPED on this
branch's earlier pushes, so this is the first run that exercised it; the green
"ATOM Test" on 7ef28c3 and 64f088d says nothing about accuracy. DeepSeek-V4
passed throughout because it has its own backend and never reaches the common
builder -- which is also why local V4 testing never saw it.
Each site cuts to the count already in scope. Nothing gains a parameter or a
field, and every cut restores the width the same line had before 8234bf4:
aiter_mla.py context_lens[:bs] under the kv_indptr cumsum
aiter_attention.py both per-seq args of make_sparse_prefill_metadata,
sliced like the block tables directly above them
attention_mla.py the two chunked-prefill varlen calls take as many
q-cums as the chunk built k-cums; the paged call
takes kv_last_page_lens' width
attention_mha.py unified_attention's two per-seq args, cut to the
block table it indexes with them
triton_mha.py the fake block table's row count, from the batch
rather than from the padded array's length
gdn_attn.py query_start_loc, PREFILL ONLY -- that function also
serves decode, whose padding predates this and whose
consumers are written for it
`minimax_m3/sparse_attn.py` and `index_topk.py` need no edit: their widths come
from the aiter_attention call above, and clamping them locally would hide a
real disagreement instead of fixing one.
Verified by reproducing first. GLM-5.2-MXFP4 MTP tp4 died on its first prefill
batch with the error above on four ranks and zero completions; after the fix:
GLM-5.2-MXFP4 MTP 0.9249 flexible / 0.9280 strict, acceptance 75.3%
(threshold 0.92, accept threshold 0.595), 0 tracebacks
MiniMax-M3-MXFP4 0.9393 flexible / 0.9401 strict (baseline 0.9363,
threshold 0.93), tp8 5-shot chat, 0 tracebacks
unit tests 4731 passed, 0 failed
Between them those two cover four of the sites: MiniMax reaches the
aiter_attention pair and the two minimax_m3 consumers downstream of it, GLM the
aiter_mla one. The other four models that failed -- Qwen3-Next-80B,
Qwen3.5-397B, DeepSeek-R1-0528-FP4 MTP, Kimi-K2.7 -- have no recipe on this
box and were NOT run; their sites were checked statically against 8234bf4^
instead. The gdn_attn edit landed after both e2e runs and neither model uses
that backend, so it carries static evidence only.
Known cost of fixing it this way: six sites now spell "the requests this rank
scheduled" six different ways, two of them by reading a length off another
array that happens to be unpadded. The invariant that keeps those two correct
is not stated anywhere.
Jasen2201
added a commit
that referenced
this pull request
Aug 31, 2026
Resolve against #2088, which renamed the decode step-shape parameters: prepare_decode now takes running_bs/running_tokens/max_seqlen_q instead of bs, so the replicated index-cache path adopts those names. The index_slot_mapping fill and its H2D copy are sized by running_tokens, the direct replacement for the old bs * max_seqlen_q, and the identity token_to_seq_idxs map for one-token-per-request decode uses running_bs.
PerryZhang01
pushed a commit
that referenced
this pull request
Aug 31, 2026
Two conflicts, both the same rename from #2088 ("Settle a step's shape once"): the padded per-sequence width in the MLA builder is now the caller-supplied `running_bs`, not a locally re-derived `bs`. - `prepare_prefill` takes main's new `(batch, running_bs)` signature and forwards it to `CommonAttentionBuilder`. The pooled-kpool helper `_sparse_selected_counts`, which this branch added immediately above it, is kept. - The sparse decode indptr keeps this branch's `_sparse_selected_counts` over main's `np.clip(..., index_topk)` -- the pooled indexer selects whole pools plus an unscored tail, so the count is not a clip -- and reads it off `context_lens[:running_bs]`, main's width. The two agree below the sparse threshold by construction, so the token-granular models main was written for are unaffected. Checked on the merge result: every `prepare_prefill`/`prepare_decode` override in the tree carries the new signatures; this branch's own additions to the other five co-touched files (glm5_next registration, KDA 0-based layer numbering, fp32 temporal state, the rope zero pad) all survive alongside main's `cu_seqlens_q` cuts. `ruff check` on the six files reports only the five findings main already has.
ganyi1996ppo
added a commit
that referenced
this pull request
Sep 1, 2026
…ompile The two triton launchers in triton_fused_sigmoid_mul_quant.py launched with a raw grid=(M, ...) inside the @support_torch_compile Kimi-K3 forward. With no compile guard, Dynamo traced into them and specialized the token dim M to a constant (16384 = max_num_batched_tokens). After main #2088 ("Settle a step's shape once") newly marks that dim dynamic, warmup compile raised ConstraintViolationError on hidden_states.size()[0], failing K3 accuracy CI. Wrap both launchers as torch_compile_guard custom ops with matching fakes, mirroring situ_and_mul_quant in kimi_k3/activations.py: - _fused_sigmoid_mul_fp8_per_token_quant (the path K3 actually uses) - _fused_sigmoid_mul_fp8_group_quant (per-1x128; extracted from the public fused_sigmoid_mul_fp8_quant, which now resolves the env transpose_scale to a concrete bool then dispatches) Guarded ops are opaque to Dynamo, so the raw triton launch never forces M to a constant. Non-quant path was already safe (pure torch). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
valarLip
pushed a commit
that referenced
this pull request
Sep 2, 2026
…and full attention (#2089) * [feat](k3): fuse activation+fp8 per-token quant for dense/shared MLP and full attention Two ptpc_fp8 fusions for Kimi-K3, each collapsing an activation kernel and a separate per-token FP8 quant into a single pass so the bf16 activation never round-trips through HBM. Both gate on the consuming Linear's runtime scheme (_effective_layer_quant == per_Token/fp8) and fall back to the unfused path otherwise. 1. Dense MLP + shared-expert down_proj: SiTUv2-and-mul + per-token fp8 quant via the aiter `situv2_and_mul_quant` kernel. New compile-safe wrapper `situ_and_mul_quant` (torch_compile_guard, returns (fp8, scale)); SituAndMul grows a `fused_quant` flag and KimiMLP forwards the scale into down_proj. Routed experts are excluded by the online-quant globs; only dense + shared (which resolve to per_Token/fp8) take the fused path. 2. KimiFullAttention o_proj input: sigmoid(g_proj) * attn_out + per-token fp8 quant. Adds a `per_token=True` mode to the existing triton_fused_sigmoid_mul_quant kernel (previously per-1x128 only, used by qwen3_next), returning a single [M,1] scale; KimiFullAttention.forward feeds (fp8, scale) straight to o_proj. Validation: - Kernel UTs: scale err ~1e-9, fp8 dequant rel err ~3%, zero/empty inputs clean. - Device time (torch.profiler per-kernel, reconciled against a kineto trace): ~1.7x (decode) to ~2.6x (prefill) for #1; ~1.5x-2x for #2. - GSM8K: 0.9674 strict / 0.9666 flexible (tp8, ptpc_fp8) -- no regression. Depends on aiter `situv2_and_mul_quant` (branch situv2-and-mul-quant); no-ops gracefully where the scheme doesn't match if that kernel is absent at build. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [refactor](k3): unify sigmoid-mul quant path behind one QuantType-selected helper Address review feedback (valarLip): collapse the two-branch KimiFullAttention.forward into a single `o_proj(x, x_scale=scale)` call site. - fused_sigmoid_mul_maybe_quant(attn, gate, quant, quant_type): returns (fp8, scale) when quant=True and (bf16, None) otherwise, so o_proj always takes the same x_scale= path (None -> o_proj self-quantizes). - quant_type (QuantType.per_Token or QuantType.per_1x128) names the scheme explicitly instead of a per_token bool; per-1x128 delegates to the existing block-quant path, unsupported types raise. K3's o_proj fusion gate now also engages for per-1x128 block fp8, not just per-token ptpc_fp8. - Import hoisted to the module top-level (no circular dep). No behavior change for K3's per-token path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [refactor](k3): unify SiTUv2 quant/non-quant path behind situ_and_mul_maybe_quant Mirror the fusion #2 unification onto fusion #1: SituAndMul.forward now always returns (activation, scale) — (fp8, scale) when it fuses SiTUv2 + per-token FP8 quant, else (bf16, None) — so KimiMLP takes a single down_proj(x, x_scale=scale) call site with no isinstance branch. Unlike fused_sigmoid_mul_maybe_quant there is no scheme selector: the aiter situv2_and_mul_quant kernel implements only the per-token scheme. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [fix](k3): guard fused sigmoid-mul-quant triton launchers for torch.compile The two triton launchers in triton_fused_sigmoid_mul_quant.py launched with a raw grid=(M, ...) inside the @support_torch_compile Kimi-K3 forward. With no compile guard, Dynamo traced into them and specialized the token dim M to a constant (16384 = max_num_batched_tokens). After main #2088 ("Settle a step's shape once") newly marks that dim dynamic, warmup compile raised ConstraintViolationError on hidden_states.size()[0], failing K3 accuracy CI. Wrap both launchers as torch_compile_guard custom ops with matching fakes, mirroring situ_and_mul_quant in kimi_k3/activations.py: - _fused_sigmoid_mul_fp8_per_token_quant (the path K3 actually uses) - _fused_sigmoid_mul_fp8_group_quant (per-1x128; extracted from the public fused_sigmoid_mul_fp8_quant, which now resolves the env transpose_scale to a concrete bool then dispatches) Guarded ops are opaque to Dynamo, so the raw triton launch never forces M to a constant. Non-quant path was already safe (pure torch). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Revert "[fix](k3): guard fused sigmoid-mul-quant triton launchers for torch.compile" This reverts commit fccfcc8. The guards did not fix the K3 accuracy CI failure: CI reproduced the same ConstraintViolationError with them in place. The triton launchers were never what specialized the token dim; the real cause is fixed in the next commit. Unguarded triton launchers inside a @support_torch_compile forward are the norm here (rmsnorm_nw in deepseek_v4, triton_mrope in qwen3_next). Co-Authored-By: Claude <noreply@anthropic.com> * [fix](k3): pass o_proj quant scheme into forward as QuantType.value KimiFullAttention stored the raw QuantType and passed it into fused_sigmoid_mul_maybe_quant, which compared it against QuantType.per_Token inside the traced forward. QuantType is a pybind11 type, so Dynamo cannot compare two of them: Failed to compare UserDefinedObjectVariable(QuantType) with UserDefinedObjectVariable(QuantType), because ... is not a Python constant That graph break split KimiModel.forward, so AttentionResidual.forward compiled as its own frame, where block_residual re-entered as a fresh input with a concrete token dim. Its torch.cat then pinned the dim to max_num_batched_tokens, raising ConstraintViolationError on hidden_states.size()[0] at warmup and failing K3 accuracy CI. Store o_type.value and compare .value on both sides, matching qknorm_quant_type_value in the same __init__ (and deepseek_v2, llama, activation.py). Verified with my_script/run_atom.sh: no ConstraintViolationError, and the graph-break log is clean -- compile now traces past the attention layer. Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: ganyi <ygan@amd.com>
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Four commits on one thread: the engine now settles a step's shape in one
place, and the three defects that centralisation exposed are fixed.
refactor(engine): settle a step's shape once, in the two units it hasbs * max_seqlen_qwas computed in nine places and the flat replay bucket searched twice with different operands.ForwardMode.decideis now the only answer, inrunning_bs(per-sequence) andrunning_tokens(per-row);scheduled_*stays this rank's own. The per-step DP scalars ride one packed all_gather.fix(spec-decode): stop the draft pass handing out storage it does not ownDraftGraph.runreturned the recording's own tensors. Draft graphs share one memory pool, which packs a later capture into memory an earlier one released, so replaying one size rewrote addresses another size's live output held.prepare_draft_idsholdsdraft_token_idsacross a step, so it read another size's float activations as vocab indices and the target's embedding gather faulted 8/8.fix(spec-decode): let the caller say whether its draft pass is uniform_publish_draft_shapeasserted "not a prefill" and "every rank runs this height" for every flavor. Both hold for the DSpark block pass; neither holds for Eagle's first draft step, which runs the target's own token stream against still-prefill metadata. Both are now the caller's to state.fix(engine): cut a prefill's logits to the requests before sampling themrunning_bsabove its batch, so the LM head emittedrunning_bsrows againstscheduled_bs-wide sampler parameters. Found by review of the three commits above.The one to look at
The last commit is a live crash, not a latent one. A three-sequence prefill
rounds to four rows on the default ladder, so any request setting
top_p < 1or
top_k != -1on a non-greedy batch divided[4, V]by[3, 1]. Greedybatches take
argmaxand temperature-only ones over-read the parameter buffersilently -- which is why GSM8K and the throughput benchmarks never saw it, and
why the fix arrives with a unit test rather than a measurement.
Verified
V4-Flash-DSpark, tp8 + dp-attention, fp8 KV / fp4 indexer,--num-speculative-tokens 7-- the one configuration here that reaches the DPbranch these commits change.
The GSM8K arm overlaps the pre-change measurement in both directions (0.9505
and 0.9499 on two earlier arms of the same config); single runs on this model
span 1.5pp, so these are indistinguishable rather than equal.
Not covered
box and were not run. Both were restored to their pre-branch behaviour by
the third commit, so the risk is a restoration rather than a new claim.
tests/test_postprocess_width.pyand the drafter-level tests needaiterat module load, so CI skips them.
tests/test_dp_metadata.py,test_forward_mode.pyandtest_dp_sync_layout.pydo run there.nothing here measures that; the numbers above are a single arm.
Deliberately left
running_tokens_are_unifiedno longer carries theor not enable_dp_attentiondisjunctmainhad. That is a decision, argued intests/test_forward_mode.py: one flag, one question. The cost is that plain DPwithout dp-attention drops decode off CUDA graphs whenever a peer prefills.
Both behaviours are self-consistent --
main's claim was self-fulfillingbecause it also padded every rank to the group max -- and no configuration
here exercises that combination, so it is called out rather than changed.
Test plan
bash .github/scripts/run_unit_tests.sh