`DEEPINFER_ENABLE_FLAGGEMS=1` turns it on. Verified on all 8 DP ranks: exit=0, 8/8 ranks agreeing,
no kernel fallbacks, `hard failures: none`, correct answers. 19 ops actually dispatched, 152 call
sites, 43 registrations from 37 requested names.
The default is `flag_gems.only_enable(include=[...])`, not `enable()`, and that is the whole design:
* allow-list -- every intermediate state is safe (unproven ops stay on PyTorch, which is correct by
construction); an incomplete list costs throughput, not correctness.
* deny-list -- I tried it first. Five full engine runs, one new defect each: `cat` (e8m0 dtype),
`mul` (broadcast shape), `topk` (compile blowup), `rms_norm` (signature), `mm` (compile assert).
The method requires being bitten before it can exclude, and every state along the way is unsafe.
That matches the instruction this was built to: use what works, leave the rest.
Where it is called is not incidental. `import flag_gems` initialises a CUDA context -- measured,
`torch.cuda.is_initialized()` flips on the import alone, before `enable()`. A `sitecustomize`
injector therefore put all 8 ranks on device 0 and NCCL refused with "Multiple Ranks are using the
same GPU". It now runs inside `create_llm`, after `setup_cuda_visible_devices`, which is the same
hazard already documented there for the profiler.
Defects found and documented rather than worked around, because each one is worth knowing:
* `mul` returns the *first operand's shape* when broadcasting expands it: (1,64,8)*(4,1,8) gives
(1,64,8), and (64,1)*(1,128) gives (64,1). Silent, no error, fewer elements computed. It made
`engram_hash_ref` disagree with PyTorch during a real run. add/sub/div/maximum/minimum all
broadcast correctly, so it is mul's own bug.
* `cumsum` drifts by up to 1.0 over a 512-element bf16 row (tolerance ~0.17) -- bf16 accumulation.
* `topk`'s compile time explodes with k: 4.8 s at k=8, 19 s at k=64, >180 s at k=512, and the
sampler asks for 4096. py-spy put all 8 ranks at 100% CPU inside `make_llir` with every GPU idle.
* `cat` raises `KeyError: 'float8_e8m0fnu'` on the model's UE8M0 scale tensors during weight load.
* `mm` passes every isolated probe yet fails inside the engine at `general_mm`; trigger unknown, so
GEMMs stay on cuBLAS, which is faster anyway.
Two methodological errors in my own measurements, both fixed and both recorded:
1. After `enable()`, `torch.mul` *is* `flag_gems.mul`, so the first version of the mul A/B compared
FlagGems with itself and reported a perfect score. All references are now computed before enable.
2. "ok" did not mean FlagGems ran. `torch.bitwise_xor` has no FlagGems registration at all, so that
"match" was PyTorch against PyTorch. `screen_ops.py` now checks registration first and reports
such cases as `skip` rather than pass.
Name mismatch is its own trap: FlagGems registers `add` bare but `eq.Scalar`, `where.self`,
`repeat_interleave.self_int`, `_softmax` (not `softmax`), `div.true_divide`. Asking for `softmax`
silently registers nothing. The bootstrap now expands names against `flag_gems._FULL_CONFIG` and
warns on anything unresolved -- which is how `softmax` was found to need `_softmax`.
Output is not token-identical to the no-FlagGems build (diverges at the first token), but two runs of
the same FlagGems config are bit-identical, so the difference is Triton-versus-eager accumulation
order, not nondeterminism. Answers stay correct and coherent.
Opens `h20-sm90-flaggems` to replace `tile_kernels`'s torch reference with FlagGems' Triton kernels.
The point is not "runs without TileLang" -- that landed already. It is speed: the compat layer is
plain torch today, so every activation in the hot path is going through eager PyTorch.
Wrote the plan down because the tempting framing is wrong. `tile_kernels` and `flag_gems` are not
substitutes; they overlap in three layers and only the first is a clean fit:
* generic ops (quant, rotary, rms_norm, topk) -- FlagGems' home turf, 251 ops
* DeepSeek-specific fused ops (mhc_*, engram_*, the five-way swiglu+quant) -- partial overlap, and
almost none of the signatures match 1:1
* the engine kernel backend (tilelang_ops) -- already done, unrelated
De-risked the highest-value one first, since per-token quantisation runs on every activation.
`ab_quant.py`, 5 shapes: `per_token_group_quant_fp8(scale_ue8m0=True, column_major_scales=True)`
against `per_token_cast(fmt="e4m3", 128, round_sf=True, use_packed_ue8m0=False,
use_tma_aligned_col_major_sf=True)` -- **quantised values bitwise identical, scales and dequantised
values exactly zero difference**. The two `round_sf` spellings are the same function: DeepInfer does
`(int + 0x7FFFFF) & 0x7F800000`, which carries into the exponent field, i.e. ceiling to a power of
two; FlagGems does `exp2(ceil(log2(s)))`. Same function, so the test demands exactness rather than
tolerance and reports any difference.
`ab_import_side_effects.py` pins the mechanism that makes this safe: import is inert,
`only_enable(include=[...])` touches exactly what it names. Recorded also that the package root
pulls Triton TLE -> libflagcx, so selective imports or vendoring stay the default posture rather
than reaching for `enable()`.
Two questions the plan puts ahead of the work, because without answers the payoff is unknowable:
nothing on this line has ever been measured for speed, and `batched_transpose_weight_sf` is still
the only path operator never A/B'd -- it rewrites every weight's scale factor and is the leading
suspect for the token divergence. Fixing that could be worth more than any op swap here.
Also names what is deliberately off the table: `enable()` with no arguments, swapping tilelang_ops
again, and churning broadcast-only micro-ops that were never the bottleneck.
Two files claimed that importing `flag_gems` registers ATen implementations process-wide, and used
that to justify vendoring. The conclusion was right for the wrong reason, and the reason was wrong
enough to distort planning: it made any broader FlagGems integration look reckless.
Measured instead of assumed. Hashing `torch._C._dispatch_dump_table` for `aten::rms_norm`,
`aten::silu`, `aten::softmax`, `aten::topk`, `aten::lerp` before and after `import flag_gems`:
every table unchanged. The package `__init__` creates the `aten_lib` handle and imports the op
modules; registration happens only when `enable()` or `only_enable(include=[...])` is called. And
`only_enable(include=["rms_norm"])` changes exactly `aten::rms_norm` and nothing else.
So the blast radius lives in `enable()`, not in the import -- which is what makes a wider
integration tractable rather than risky.
Vendoring still stands, for two reasons that survive the correction: the package root star-imports
every op module and loads Triton TLE, which dlopen's `libflagcx` (a native dependency this
deployment does not otherwise need); and the released kernel file has no internal imports, so the
package bought nothing for that operator.
Both files now say this, and say that it is a correction, so the next reader does not rediscover it.
Five branches, each building on the last, and until now the only way to tell them apart was reading
commit messages. That is fine for two and not for five -- in particular the difference between
`h20-sm90-triton` and `h20-sm90-triton-flagtree` is invisible from the names, and the whole point of
the latter is a rebuild inside a different image.
The table states, for each branch, what it fixed and what it still could not do, so picking one is a
lookup rather than an experiment. It also states plainly that nothing here has been measured for
speed: every result on this line is "runs and is numerically correct", never "is fast".
The last thing that kept the SM90 path on TileLang was the sparse MLA attention kernel. FlagGems
has a pure-Triton implementation of the same operator, and it turns out to be drop-in once two
differences are handled. With TileLang unimportable, the engine now runs end to end.
`h20_ops/sparse_mla.py` picks the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
means TileLang when it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
two entry points were the only call sites touched. Auto resolution catches `ModuleNotFoundError`
rather than testing for the package by name, because a missing TileLang fails inside
`sparse_mla_sm90`, not at the module lookup -- and it only swallows a *missing* dependency, so a
present-but-broken module still surfaces instead of being swapped out.
Vendored rather than depended upon: `flag_gems/__init__.py` registers ATen implementations
process-wide (`torch.library.Library("aten", "IMPL")` then `replace_customized_ops`). Importing the
package to reach one kernel would swap operators under the whole engine. The PyPI 5.0.2 file has no
`flag_gems` internal imports at all, so there was nothing to gain; it is 329 lines, accumulates in
fp32, and already guards the zero-width rope tail that V4.1 needs (RoPE folded inside the 512-wide
latent). The file header carries the source md5 and both deviations.
Two things it needed, both found by comparison rather than by reading:
* 5.0.2 guards each top-k block by its **position** in the index array (`if ck * BK <= max_col`).
Non-causal that is `SKV-1`, so when `topk > SKV` every block past the first `SKV` entries is
skipped. At SQ=16/SKV=128/topk=512 the output is off by 2.8e-2 -- and any prompt shorter than the
512-entry top-k hits it. Restored the data-dependent guard FlagGems master uses, marked in place.
* Its lse is a natural log; the TileLang one is log2. Unconverted, `_apply_sink` would be off by
1.4427x: plausible output, wrong numbers. The adapter converts, and the test asserts the ratio is
1.0 rather than 0.693 so it cannot regress quietly.
The kernel's operands have no batch axis, while decode hands it a different KV window per query.
Rather than loop, the adapter concatenates the windows and shifts query b's indices by b*seq_kv --
one launch, no extra memory, since the concatenated buffer is the tensor decode already built.
Verification ran in order, each step gating the next: 7 shapes against TileLang (`ALL_MATCH`, <=2
bf16 ULP, including `topk_length` and `attn_sink`), the decode call shape incl. the SWA+full
two-cache case (`ALL_MATCH`), the adapter contract (5 tests), the dispatcher and both `attn_sm90`
entry points (4 tests), then end to end with TileLang blocked: exit=0, 8/8 DP ranks agreeing, zero
kernel fallbacks, correct answers. A control run with TileLang present is unchanged.
Not token-identical, and the doc says so: prompt A shares 15 tokens then diverges, prompt B shares
1. A control run of the TileLang backend twice is bit-identical, so this is accumulation order, not
noise -- consistent with the 2.7% logits drift already recorded for the SM90 quantisation floor.
One TileLang dependency remains, and it is the deliberate one: `fused_sample`'s rejection-sampling
variant, which needs TileLang's Philox stream to reproduce the same tokens. So the pure-Triton run
disables speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`).
The remaining blocker for a TileLang-free path was sparse MLA: `h20_ops/sparse_mla_sm90.py` is a
231-line TileLang program with no Triton counterpart in this repo. FlagGems already has one, and it
turns out to be drop-in.
`flagos-ai/FlagGems` -> `src/flag_gems/fused/flashmla_sparse.py` -> `flash_mla_sparse_fwd`. Pure
Triton. At V4.1's real dims (h_q=64, h_kv=1, d_qk=512, topk=512) it compiles and runs on H20, and
on 7 shapes it agrees with the TileLang kernel within 2 bf16 ULP -- including `topk_length` and
`attn_sink`, the two optional arguments.
The pairing is closer than expected, which is why it needed no adaptation:
prefill_attn_varlen_sparse(query, kv, index_topk, softmax_scale, attn_sink, index_length)
flash_mla_sparse_fwd(q, kv, indices, sm_scale, attn_sink, topk_length)
Causality matches too, and for the same reason. `attn_sm90.py` passes `is_causal=False` because the
indexer only scores positions before each query, so the indices already carry causality -- bounding
on the row index instead collapses the layer to NaN, which already happened once and is documented
at that call site. The FlagGems kernel has no causal parameter either.
Recorded as *not* the answer: `DSA/sparse_mla.py` in the same repo. It hardcodes causal and needs
the V3.2 576-wide split layout; V4.1 folds RoPE inside a single 512-wide latent, so TD = 0 and it
fails to compile on `tl.arange(0, 0)`. Worth writing down because the name looks like the better
match.
What remains, all much smaller than writing a kernel: wire it in behind the existing
`patch_imports` hook; test the decode path separately (it gathers to bf16 with identity indices, a
different call shape than the 7 cases); and convert the return, since the kernel gives
(out, max_logits, lse) with lse in natural log where DeepInfer expects (out, lse) with lse in log2
-- a 1.4427x scale error if dropped.
Records the result honestly in both directions. It runs: exit=0, 8 DP ranks agree, answers are
coherent. It is also not equivalent to the compiled path: 16 of 16 samples differ, and a control
run of the compiled backend twice is bit-identical, so the difference is reproducible rather than
sampling noise. One prompt diverges at token 16, the other at token 1.
Lists what could account for it instead of asserting a cause: batched_transpose_weight_sf is the
only path operator never A/B'd and it rewrites every weight's scale factor; apply_rotary is one
bf16 ULP off on 1 element in 65536 and runs every layer; engram_gate is 3e-7. A first-token
divergence argues for comparing prefill logits next rather than guessing further.
Also rewrites the earlier wrong diagnosis: the eight GPUs were not held by colleagues' containers.
They were held by this container's own deadlocked leftovers from a run I killed -- `pkill` took the
parents, `multiprocessing.spawn` children were orphaned to PID 1 and kept the CUDA contexts for two
days, spinning six cores and writing 1.4 GB of retry logs. The lesson worth keeping is the check
itself: map PID to cgroup before blaming another tenant.
Running end to end surfaced two things the per-kernel A/B could not.
The engine wraps the forward pass in `forbid_cuda_sync()`, which sets torch's sync debug mode to
"error". The compiled kernels never sync; the torch references do, and not by oversight:
* `tile_kernels/torch/cast.py` built `torch.tensor(-1.0, device="cuda")` (and three more like it)
for use as a `torch.where` operand. Constructing a device tensor from a Python scalar is a
host-to-device copy, so it syncs -- and it is unnecessary, because `torch.where` accepts a plain
scalar and passes it to the kernel. Patched via `patch_torch_ref_sync.py` (idempotent,
reversible), which is a fix worth having on its own: the same cast raises for any caller that
holds the device clean.
* `get_mapping_from_psum` reads the expert prefix-sum to the host to compute alignment ranges in
Python, and `moe_topk_gate_forward` compacts rows with `nonzero()` and branches on `.any()`.
Those cannot be made sync-free without rewriting how the reference addresses, which is the wrong
trade for a fallback whose job is being runnable. So each compat module wraps its reference
module once and every call through it runs inside `allow_cuda_sync()`.
Not done, deliberately: `set_sync_debug_mode("default")` process-wide. That would also stop the
guard catching accidental syncs in the compiled path, which is what it is for.
Also fixed: my own `bool(force_random.any())` in the gate adapter was itself a sync, outside the
reference, so it raised before the adapter could decide anything.
End to end now completes: exit=0, all 8 DP ranks agree, coherent answers. It is not token-identical
to the compiled path -- 16/16 samples differ, with a control run confirming the harness is
deterministic (two cuda runs are bit-identical), so the difference is real. Recorded in the doc as
the next thing to attribute rather than papered over.
Status table updated: the routing layer exists, all 23 APIs execute with TileLang blocked, and the
two backends agree on 18 cases. End to end still has not completed, and the doc now says why with
evidence rather than a shrug: six of eight GPUs hold 82 GiB from other containers, so the run OOMs
in allocate_userbuffer and the rest wait at a barrier. Also records the command to rerun once the
GPUs free up, and two environment traps hit on the way -- /tmp filling with docker data (a bind
mount elsewhere means deleting the obvious directory frees nothing) and the compiled _C.so being
lost by replacing the source tree.
The references were importable and numerically validated, but nothing pointed the model at them:
the forward pass still called the compiled kernels, which raise ModuleNotFoundError on a host
without TileLang. This adds the routing layer and swaps the call sites.
`deepinfer/tk_compat/` exposes the *compiled* signatures -- out-buffers, in-place conventions,
argument order -- because that is what the 25 call sites are written against, and adapts to the
reference's shape internally. Selection is `DEEPINFER_KERNELS=auto|cuda|ref`; `auto` follows
whether TileLang imports.
The adapters are not pass-throughs. Each difference was found by running both backends:
* `per_token_cast` takes `num_per_channels=n`, the reference takes `block_size=(1, n)`.
* `apply_rotary` takes a concatenated cos|sin cache and rotates in place returning None; the
reference takes cos and sin separately and returns a new tensor. The adapter splits, calls, and
copies back, because the call site relies on the in-place contract.
* `engram_gate` takes the pre-multiplied weight and `(eps, clamp_value)`; the reference needs the
two weights separately and `(clamp_value, eps)`. Since a product cannot be un-multiplied,
`fused_weight` returns the pair on the reference path and the gate unpacks it.
* `mhc_pre_big_fuse_fwd` has no single reference: its five steps are composed from the reference's
own components.
* `mhc_reduce_partials_and_rmsnorm_fwd` likewise; rebuilt from the reference's formula.
`moe_topk_gate(force_random=...)` raises rather than silently ignoring the flag. The reference has
no equivalent, and dropping it would route tokens differently while looking healthy.
Verification, all on the real call signatures: 18 cases compared between `cuda` and `ref` modes,
`ALL_MATCH`, 12 bitwise identical including every quantised byte and every routing index; plus a
smoke run that executes all 23 APIs with TileLang blocked, `ALL_CALLS_OK`.
End to end still has not completed, and the reason is not this change: six GPUs on the host are
held by other containers (82 GiB of 139 each), so the 8-way EP run OOMs in allocate_userbuffer and
the surviving ranks wait at a barrier. Recorded in the doc along with the command to rerun it.
Section 6.0 now states the status as a table, because the previous wording invited a wrong
reading. Importing the torch references -- and even deepinfer itself -- succeeds without
TileLang, and five operators are numerically validated. But the model forward still calls the
compiled quantisation and MoE entry points, which raise ModuleNotFoundError, because the
routing layer that would point them at the references does not exist. End to end has never
run, so the performance regression is an expectation, not a measurement.
Section 6 writes down the three eager imports that blocked the plan, the annotation trap that
looks like it should solve it and does not, and the comparison table showing the quantisation
chain is bitwise equal -- plus, plainly, what is still missing so the next step is not mistaken
for a finished port.
The bring-up plan for a card that cannot build TileLang is to run the model on
`tile_kernels.torch` -- pure-torch implementations of nearly every kernel the forward pass
calls -- and replace the hot ones with Triton afterwards. That plan was blocked at the
import: `import tile_kernels.torch` dragged in the whole compiled stack.
Three eager imports were responsible, and each is now lazy or optional:
* the top-level `__init__` re-exported every compiled kernel; it resolves them on access
* `quant/__init__` did the same for 13 modules, and `import quant.common` runs it first
* `quant/common.py` mixed host-side config dataclasses with `@T.macro` fragments; the
dataclasses now import without TileLang, the macros are behind a guard, and its three
dtype-handle members resolve TileLang when called rather than when imported
`config.py` also only needed TileLang for three constants, now literals.
One thing that looks like it should work and does not: a module-wide
`from __future__ import annotations`. It also stringifies the annotations inside `@T.macro`
bodies, which TileLang's macro parser evaluates itself, and it fails there. Only the four
host-side annotations are quoted, individually.
The patches are idempotent and reversible, and `verify_tile_kernels_patch.py` proves the
surface is unchanged where TileLang is installed -- the kernel modules' wildcard imports of
`quant.common` still see the macros, `from pkg import sub` still works, config knobs still
take effect.
The other half of the question, whether the references are numerically trustworthy, is
answered by `compare_with_cuda.py`: on the recipe the engine actually uses on Hopper,
per_token_cast, cast_back, swiglu_forward_and_per_token_cast and engram_hash are bitwise
identical, and apply_rotary differs by one bf16 ULP on one element in 65536. The quantisation
chain being bitwise equal is what makes a torch-first bring-up safe rather than approximate;
a cast that were merely close would change tokens downstream.
Not yet done, and recorded in the doc: MHC and the MoE gate still need their own comparison
before anyone assumes they match, and the routing layer that points deepinfer's 25 call sites
at these references does not exist yet.
I had been reporting the port as "16 kernels left". That was wrong in a way that matters, so
this corrects it with a measurement rather than an estimate.
The measurement: make `tilelang`, `tile_kernels` and `tvm_ffi` unimportable via a meta-path
finder installed from sitecustomize -- sitecustomize specifically, because the eight DP ranks
are spawned and would not inherit an import in the parent. Then try to import the engine.
It fails at import, before any kernel is reached: `hugging_face/weights.py` pulls in
`layers/linear.py`, whose module body has `import tile_kernels`.
The reason the kernel ports did not help is that `tilelang_ops` only ever covered 44 kernels.
The model forward calls `tile_kernels` directly -- 25 module-level imports across 13 files,
23 distinct APIs -- and `tile_kernels` is itself a TileLang package: its `__init__` line 1 is
`import tilelang`, with 166 references across 140 modules. Keeping tile_kernels keeps
TileLang.
The 23 APIs are recorded in three groups, because they cost very differently: six are
configuration knobs and string helpers; six are hot-path quantization and routing kernels
(`per_token_cast` appears in four modules); ten are MHC/Engram kernels specific to this model
and the largest piece of work.
There is one real shortcut, and it is worth knowing before anyone starts writing Triton:
`tile_kernels/torch/` holds pure-torch reference implementations of nearly all of them, with
zero TileLang references. It is currently unreachable only because the package's `__init__`
imports TileLang eagerly -- so making that lazy buys a correct-but-slow bring-up path, from
which the hot kernels can be replaced one at a time as profiling justifies it. `flag_gems`
cannot stand in: it is 903 ATen-level operator replacements and does not carry these fused
kernels.
An edit of mine dropped a `return launch` from `finalize_verification_tl`. The A/B case for
that kernel was not re-run, so nothing local objected; the failure surfaced as a TypeError
inside the sampler of an end-to-end run, once per rank. The harness now checks that every
`*_tl` symbol the selected backend exports is callable, which turns that class of mistake
into one line printed before any case executes.
Also records where the remaining sixteen kernels actually sit. None of them is
unconditionally on the path, but the distribution is lopsided and the document had it
backwards in one place: the three `candidate_*` kernels are not merely "off on SM90", they
are switched off *by* `PLATFORM.is_sm90` -- and this checkpoint sets
`candidate_source_layer_id=20` of 40 layers. On a card where `is_sm90` is False, that is
nineteen layers on every step, which is the opposite of optional. Their harder blocker is
noted alongside: that path also needs `deep_gemm.fp8_fp4_sparse_mqa_logits`, which is
SM100-only, so porting those three alone would buy nothing.
The four kernels that no production path reaches are listed as such, so nobody spends an
afternoon on them.
`DEEPINFER_USE_NEXTN_DECODE` defaults to 5, so speculative decoding is on unless someone
turns it off, and these four ran on every decode step. On a host without TileLang the first
step failed; on the H20 they were the only kernels still answering from the other backend
during a normal run.
The interesting one is `remap_draft_topk`. Its TileLang form builds a vocabulary-windowed
lookup table in shared memory and reads the logprobs back out of it. Triton has no
equivalent primitive, so the table is gone and each target index is matched against the
draft indices directly: draft indices are unique per token id, so at most one candidate can
match and a max over the block is well defined. That is O(topk^2) instead of
O(vocab/pool_size * topk), which would be the wrong trade at a large `force_topk` -- but
`force_topk` defaults to 0, so this kernel is off the default path and correctness comes
first. Worth revisiting if a deployment turns it on.
The other three adapt rather than change shape:
* `get_accept_result`'s prefix scan becomes a bounded loop over the widest span in the
block; lanes with a shorter span hold their value on the extra iterations. MTP packs at
most num_steps+1 rows per request, so the bound stays small.
* `finalize_verification` reduces `batch_confidence_error` within the program and issues one
atomic instead of one per row, because a scalar pointer with a vector value is not valid
in this Triton version. Accumulating in a different order changes the last bits, so its
cases compare with a tolerance and say why.
* `get_fused_dflash_get_input` rounds its column loop up. `padded_kv_len` is a multiple of 64
but not necessarily of the tile width, and flooring leaves the tail of every row unwritten.
33 new A/B cases across the four, covering each constexpr variant the engine builds. The
full run is green and an end-to-end pass is token-for-token identical to the run before the
change, with the MTP fallback notices gone.
The dependency notes in section 4.4 were a list of things to do by hand; they now name the
script that does them, along with what a from-scratch run measured.
`h20_e2e_report.sh` sets the two settings a single-node H20 run always needs -- no GIN
because the container has no /dev/infiniband, and eager because the Engram gather cannot be
captured -- and appends Triton's own cudart to the loader path when it is present, since cu13
images otherwise fail to load any kernel. Requiring a caller to remember three exports to
reproduce a run is how a run stops being reproducible.
Getting the Triton backend to build and run in a vendor image took a dozen non-obvious
steps: a venv because setup.py insists on CONDA_PREFIX, CUDA headers borrowed from the
nvidia-cu13 wheels because the image ships only the runtime, NCCL 2.32.3 because DeepEP
uses ncclCommQueryProperties, libdw/libelf for DeepGEMM's bundled deep_jit, pciutils for
DeepSelect, and a tvm_ffi new enough for tilelang 0.1.15. All of it lived in this session
and in prose, which means the next person starts over.
`setup_flagtree_image.sh` records it in three stages -- env, build, verify -- and is
idempotent, so a failure resumes. It also explains why each workaround exists, since several
look arbitrary otherwise: FlagTree is a Triton fork and needs nothing special, but the C++
extensions around it do.
`copy_site_packages.py` copies distributions by their declared file lists rather than by
name, because pillow installs PIL and pillow.libs, pyzmq installs zmq. It defaults to
keeping the image's pillow: a copied one points at vendored libtiff builds the image does
not have, and PIL.Image then fails to import.
Verified by building into an empty venv rather than the one already set up here: env, build
and verify all pass from scratch, and an end-to-end run from that venv produces the same
eight-rank-agreeing output as before.
The raw log prints one block per rank, so the two things that decide whether a run passed
-- do all eight ranks agree on the token ids, and does the fallback list match the kernels
we know are unported -- are both buried in repetition. This extracts them: a count per
distinct token sequence (eight means every rank agreed), one continuation per prompt, the
kernels still served by the default backend, and whether a traceback occurred.
Splitting "hard failures" out from the stats logger's output matters on torch 2.11, where
`ProcessGroupNCCL.memory_stats` does not exist and the logger says so without the run being
affected.
Eight ranks, `--ep 8 --max_tokens 32`, both prompts, on the FlagTree image: exit 0, every
prompt's token-id sequence identical across all eight ranks, decoded text correct. The same
run on the reference image produces byte-identical token ids, so the port is not merely
"some output comes out" on the vendor image -- it agrees with the image we have been
validating against all along.
The fallbacks are the known P1/P2 leftovers and are identical on both images: MTP's four
kernels, the sampler's rejection-sampling variant, and the 20 optional-path kernels. One
benign difference is recorded too -- torch 2.11 has no `ProcessGroupNCCL.memory_stats`, so
a stats line reports itself unavailable.
Also writes down the dependency work an end-to-end run needs but the A/B harness does not:
tilelang 0.1.15 with tile_kernels 2.1.0 and tvm_ffi 0.1.11, plus the 19 packages from
requirements.txt. pip stalls on this host, so the practical route is copying site-packages
out of the reference image; pillow is the exception and must come from the image, because
11.3 sits beside a pillow.libs the image does not have.
Two corrections and one result.
FlagTree is not something Triton has to be ported to: `flagtree.__version__` reads
`triton.__version__`, and `triton/runtime/jit.py` imports `flagtree._flagprism`. The
distribution is named `tokenspeed-triton` but the import name is still `triton`, because it
is a fork. So "FlagTree cannot compile" was never true -- it compiled every kernel we have,
element-wise identically to triton 3.5.1 and 3.7.1. What failed was `deepinfer._C` and the
vendored CUDA extensions, none of which involve Triton at all.
Those are now built inside the FlagTree image (torch 2.11, cu13, gcc 13.3, nvcc 13.0):
`_C` loads with 34 symbols, deep_gemm 2.8.2, deep_ep 2.5.0 and deep_select 1.0.0 all
import, and the full A/B run is 125 cases with zero skips -- the sampler group used to be
skipped for want of `_C`. Comparing against the dumps from triton 3.7.1 is still exact.
Everything needed along the way is written down, because none of it is obvious: the venv
requirement in setup.py, `python` vs `python3` in build.sh, the missing cusparse header,
NCCL 2.32.3 for `ncclCommProperties`, pciutils for deep_select, and the pip download that
has to be a resumable curl to survive the 305 MB wheel.
Two artefacts and one finding.
`scripts/flagtree_triton_probe.py` asks whether the Triton an image ships can build the
features this port depends on, one feature at a time: `tl.dot`, `tl.gather` with a
data-dependent index, `tl.gather` as the RoPE pairing, `fp8e4nv` with a block scale, and an
int64 comparison (uint32 compares go signed in Triton, which the sampler depends on). Both
vendor images pass all of them and emit `.target sm_90a`.
`TODO.md` collects every open issue across both branches, with the two blockers an
end-to-end run in a vendor image actually hits: `deepinfer/_C` is built against one torch
ABI (`2.13.0+cu132`) so it fails to load on torch 2.11 and 2.9, and `deep_gemm`/`deep_ep`
are not on PyPI in the versions `vendor/` provides. Neither is a Triton problem, and saying
so plainly is the point of writing it down.
`h20_ops_ab.py` assumed the TileLang backend is importable, so on a vendor image it could
not start at all -- which is exactly where the Triton backend needs checking. Two modes fix
that:
* `--dump DIR` saves the selected backend's output per case.
* `--compare DIR` checks against dumps written elsewhere.
Together they answer "do triton 3.5.1 and 3.7.1 agree" without either host having TileLang.
Case builders that genuinely need a package the host lacks now report a skip instead of
aborting the run, so partial coverage is visible rather than fatal; the import of the
sampler's C++-backed inputs moved into the one builder that uses it for the same reason.
This is what showed the 19 ported kernels produce element-wise identical output across
triton 3.5.1, 3.6.0 and 3.7.1 -- 115 cases, no differences.
The point of the Triton backend is to run where TileLang cannot, but importing `deepinfer`
pulled the whole TileLang stack in regardless. Three separate leaks:
1. `ops.py` imported the platform default backend even when another one was selected, purely
to have fallbacks available. That default is TileLang, so it dragged in `tile_kernels`.
`DEEPINFER_OPS_STRICT=1` now skips it: no default import, no fallback, a missing kernel
raises. That is the right trade when the default is not installed, or when a silent
fallback would hide an incomplete port.
2. `tile_kernels` sat at module scope in `norm`, `cache_common`, `fused_transform` and
`load_paged_kv_cache`, but only the reference implementations and the unit tests use it.
The kernels never do. It now lives inside those functions.
3. `tilelang` itself: `compress` set a TileLang pass flag, and `norm` needed a TileLang dtype
handle to pass along a parameter the Triton factory discards. Both now degrade to a no-op
or a plain string when TileLang is absent.
Verified on a FlagTree image (torch 2.11, triton 3.6, no TileLang): `import` succeeds, 19
kernels build and run, and the outputs are element-wise identical to the same kernels built
by triton 3.5.1 and 3.7.1. Full A/B on the reference image is still green, so nothing
changed for the paths that do have TileLang.
`h20_ask.py` takes its prompts as arguments, and a login shell that is not UTF-8 mangles
them into `????` before python ever sees them. The result is a model answering a question
nobody asked, reported as a quality problem. The prompts now live in the script.
Also records, in the plan, the state of the port as measured today: what runs end to end,
the two environment variables it still needs, and the one real blocker found while
checking -- the Engram gather is not CUDA-graph capture safe, so the default
(attention-graph) configuration fails to warm up on eight ranks.
With MTP on (the default), the sampler is built with rejection sampling for the draft-verification
rows, which this backend does not implement. It asserted, so the whole MTP path failed with the
triton backend even though 16 of its 20 kernels were ported.
It now hands that variant to the default backend and says so once. That is a deliberate delegation
rather than a silent fallback: `fused_sample` is exported by this module, so the per-kernel
`_FallbackNotice` cannot cover it, and the note has to come from the factory.
Verified both ways on 8xH20, same tokens as the pure TileLang run either way:
MTP off 20/20 kernels Triton, no fallback notice at all
MTP on 16/20 kernels Triton, and exactly five notices -- the four unported MTP kernels plus
this delegation, with nothing unaccounted for
The remaining MTP work is `remap_draft_topk`, `get_accept_result`, `finalize_verification`,
`get_fused_dflash_get_input`, and this rejection variant.
compute_seqlens_k_zos, confidence_head_fused and dflash_input_store, all identical to the baseline:
15 cases across the three.
`compute_seqlens_k_zos` masks the remap lookup rather than clamping it: `-1` marks a fresh request,
and an unmasked `PREV_ACCEPTED + (-1)` would read the element before the buffer.
`confidence_head_fused` reads the weight as one flat vector and splits it at `n_embed`, so the
hidden and Markov halves are one accumulation instead of two passes. Without a Markov embedding the
second pointer is filled with the first tensor rather than left null, since the read is compiled out
but the argument still has to be a pointer. The scaled logit is clamped to +-60 before the sigmoid,
as the baseline does, so a small temperature cannot produce inf through `exp`.
`dflash_input_store` accumulates in fp32 even though both sides are bf16: the baseline sums into an
fp32 fragment, and a bf16 accumulator would round differently for hc_mult > 1.
One trap: `tl.next_power_of_2` inside a kernel body is not a compile-time constant, so the Markov
block size is resolved in the factory and passed as a constexpr.
The plan note still said the sampler was deliberately left on TileLang, which is no longer true. It
now carries the measured numbers instead: 100% token agreement on the engine's own configuration,
~90% on narrow cuts, log-probabilities within 1e-5 everywhere, and the reason the token cannot do
better -- the baseline's rejection limit steps at `total == 2^31`, which the fixed-point total sits
next to, so a few units of rounding error move to a different draw.
Plus the four traps that each looked like something else: signed arithmetic on uint32 (`/`, `%` and
`<=`), nested runtime `if`s losing the assigned value, a threshold taken from probabilities and
compared against logits, and renormalising by division instead of through the log domain.
`fused_sample` for the configuration the engine builds: 4096 columns after the force_topk cut,
sampling on, rejection off. With one tile the baseline's two-level search collapses -- the coarse
offset *is* the token id -- so this implements log-softmax, the top-k/top-p cut as a threshold plus
the tie rule, renormalisation, a fixed-point cumulative distribution and one Philox draw.
The draw matches cuRAND bit for bit (that is what the previous commit established). The rest does
not all match, and the split is measured rather than claimed:
engine's configuration (top_p=0.99, 2644 kept columns) tokens identical on every row
narrow cuts (top_k=128, top_p=0.5) ~90% of rows
Kept sets match column for column and log-probabilities agree to ~1e-5 in every case. What differs
is the token, and the cause is a discontinuity in the baseline rather than a defect here: the draw
is `d % total` with `total = sum(floor(p_i * 2^31))`, and the baseline's rejection limit
`(2^32-1) // total * total` steps from `2 * total` to `1 * total` exactly at `total == 2^31` -- which
`total` sits next to, because the probabilities sum to one. On the near side almost no draw is
rejected, on the far side about half are, and each rejection moves to the next draw. A few units of
rounding error in `total`, which is a reduction order that cannot match, therefore change the token
outright. Clamping `total` under 2^31 does not help: the baseline crosses in either direction too,
so it only relocates the disagreement (30 rows instead of 11).
Four bugs found on the way, all worth recording because each looked like something else:
- Triton evaluates `//`, `%` and `<=` on **uint32 operands with signed arithmetic**, so 4294967295
behaves as -1. That made `limit` come out as `2^32 - total`, flipping the rejection decision.
Fixed by doing the whole cumulative distribution in int64.
- Nested runtime `if`s produced a threshold that was not the one assigned inside the branch; the
cuts are now combined with `tl.where` instead.
- The threshold was taken from the probability array and compared against logits. The sort has to be
of the logits (equivalent, since the map is monotone -- the mass is recomputed from the sorted
logits rather than sorting twice).
- Renormalising by division rather than through the log domain, as the baseline does, moved the
token agreement from 76% to 91% before the remaining gap was understood.
Also gives the A/B harness a fourth case element: an agreement floor for outputs whose answer is a
random draw, so the rate is reported rather than compared for equality -- a comparison that can pass
without testing anything would be worse than none.
The last blocker for `fused_sample` was the random stream: `T.rng_init`/`T.rng_rand` lower to
`curand_init`/`curand` with `curandStatePhilox4_32_10_t`, and the seeding is not documented. Its
mixing function is public but `curand_init` folds the seed, the subsequence and the offset into the
counter in a way that has to be measured rather than derived.
Measured (`scripts/h20_rng_dump.py` dumps the real stream), the layout is:
key = (seed & 0xFFFFFFFF, seed >> 32)
ctr = (offset & 0xFFFFFFFF, offset >> 32, subsequence & 0xFFFFFFFF, subsequence >> 32)
with draws taking the block's four words in order before `ctr.x` advances. The subsequence sits in
the third and fourth counter words, not the first -- a cuRAND subsequence spans 2^64 blocks. That
detail is invisible at subsequence 0 and wrong everywhere else, which is why the check covers both.
`scripts/h20_philox_check.py` compares our kernel against the dumped reference over five seeds
(including 2^64-1) and four subsequences (0, 1, 2048, 4097 -- the last two are the actual
`sample_positions` from a run): 20/20 streams bit-identical.
One trap worth noting: the multipliers and the second Weyl constant all exceed int32, so referencing
them as Python ints inside the kernel silently wraps them negative. They are materialised as uint32
tiles instead.
This does not yet change any output -- `fused_sample` is still the TileLang kernel. It removes the
reason it had to be, which is the point.
The two runs agree token for token on the reference chat-prefix input, but a plain-text prompt
diverges after two tokens: both backends are internally consistent across all eight ranks, so the
difference is a different draw from a slightly different CDF, not nondeterminism.
That is the quantisation floor from P4 showing up through sampling rather than through argmax -- and
on top of it, a kernel written in a different language has its own reduction order (the norm A/B
shows 10 of 655360 bf16 elements one ulp apart, with fp8 bytes and scales identical). At T=1.0 that
ulp is enough to flip a coin.
So the honest statement is a table, not a claim: identical tokens on the chat-prefix input, same
first token but a different branch on a plain-text prompt, and matching the draw itself would
require both an identical reduction tree and the same PRNG stream.
The import-time line only counts kernels the selected backend does not export, which says nothing
about whether a running configuration reaches them. `_FallbackNotice` wraps each inherited kernel
and prints once, on first call, so an end-to-end run states the answer itself. On 8xH20 with
speculative decoding off that answer is a single line:
[ops] kernel 'fused_sample_tl' used from 'cuda': 'triton' does not implement it
Eighteen of the nineteen kernels on that path are Triton and produce exactly the same tokens as the
pure TileLang run. The nineteenth is left there on purpose.
A probe of the engine's own sampling parameters (`DEEPINFER_PROBE_SAMPLING=1`, which stages an
async D2H copy and reports at exit so it does not synchronise inside the forward) settles what
`fused_sample` is being asked to do:
temperatures: [1.0, 1.0] top_ps: [0.99, 0.99] top_ks: [-1, -1]
That is a real nucleus sample over 4096 candidates, not an argmax, so the draw matters. Two things
stand in the way of reproducing it bit for bit, both internal to TileLang: `T.rng_init`/`T.rng_rand`
are a Philox stream implemented in its C++ codegen, so a Triton kernel draws a different stream and
every sampled token at temperature > 0 would change; and the baseline pins the layout of its LSE
accumulator precisely to fix the summation order, after which probabilities are quantised into a
2^31-scaled uint32 CDF where a one-ulp difference can move a column across a boundary.
Swapping it anyway would be a behaviour change reported as a kernel swap. The per-kernel fallback
exists for exactly this case, so the sampler stays on TileLang and the rest of the path does not.
`triton/__init__.py` records this so the gap is not mistaken for an oversight, along with the order
in which porting it later would have to proceed: a Philox that matches the stream exactly, then the
radix filter, then the fixed-point CDF and the two-level search.
Four more kernels, and with them the non-speculative path no longer touches TileLang:
- `fused_norm_rope_hadamard_cast_store_kvcache` -- all the forms the engine builds: bf16 out,
bf16 out after dequantising (cast_back), quantised out with scales, and the main/index KV-cache
write with a bf16 stage of `norm`. 18 cases.
- `compress_and_update_state_{prefill,decode}` -- the ring-cache state machine, two launches each.
- `prepare_v4_decode_metadata` -- one program per sequence, every output including the
ratio-bucketed ones.
Details that were not obvious and are now pinned by cases:
RoPE pairs differ per layout: interleaved pairs `(2k, 2k+1)`, Neox pairs `(k, k + rotary/2)`. Both
reduce to "is this the first element of its pair" plus a partner offset, so the lane vector is
gathered once with `tl.gather` instead of two code paths. That also fixed a real bug: the first
version paired adjacent lanes for both layouts, which is only correct for interleaved.
`bf16_stage="norm"` means the post-normalisation *pre-RoPE* row -- the indexer consumes it. It has
to be stored before the rotation overwrites the row; storing it afterwards produced a wrong value
only when RoPE was enabled, which is exactly the case the engine uses.
Scales come in two storages for the same rounded power-of-two value: a ue8m0 byte for the main KV,
an fp32 for the index KV. Both are now covered, in the standalone output and in the cache write.
`compress` needed its tail-copy grid bounded by the *query token* count, not the compressed-row
count; and `_enable_fast_math` mutates `pass_configs` on the factory, a TileLang concept, so the
Triton factories expose an inert dict to keep that call working.
Several case bugs of the same shape were found and fixed along the way -- a sentinel that was never
required to be overwritten, a `total_q` that did not match the actual prefix sum, two tokens
mapped to one cache slot or one state slot so the winner depended on CTA order. Each of them made
the two backends agree for the wrong reason, which is worth more than the kernels themselves: a
comparison that can pass without testing anything is worse than no comparison.
The illegal memory access was not in `expand_sampling_params` at all, though that is where
`CUDA_LAUNCH_BLOCKING` pointed. It was the metadata kernel leaving its output unwritten.
`prepare_verification_metadata_tl` sized its second grid dimension from `num_steps`, i.e.
`max_MTP_value`. With speculative decoding off that value is 0, and a zero-sized grid dimension
launches nothing, so `row_to_req` -- a `torch.empty` buffer -- kept whatever the allocator handed
out. `expand_sampling_params` then used those values as request indices, and the reads went off
the end of the parameter arrays.
The kernel now launches one program per request and loops the tiles at runtime, mirroring the
TileLang version, so the row count comes from device memory and no host read is needed (which
would have been rejected as a synchronising operation anyway). `num_steps` is accepted and
ignored.
The A/B case for this kernel is also fixed, in two ways that matter:
- it used a fixed `total_q` while the request lengths were random, so it allocated more rows than
were ever written and the comparison was meaningless;
- it pre-filled the outputs with a sentinel but never required them to be overwritten, so both
backends could write nothing and still agree with each other. The case now asserts every row was
written, and covers `num_steps=0` plus a request longer than one tile.
With that, `DEEPINFER_OPS_BACKEND=triton` runs end to end and produces exactly the same tokens as
the pure TileLang run, `[17607, 2793, 3139, 666, 6127, 666, 320, 1]`, identical on all eight
ranks. Fifteen of the required kernels are Triton; the remainder still come from TileLang through
the fallback while they are ported.
Three fixes found by running the ported kernels inside the engine rather than only against the
A/B harness:
`map_lookup` took the row stride from the row width, which is wrong for the views the engine
passes (a slice of a wider buffer). Every kernel now takes its strides as arguments.
`expand_sampling_params` used `row_to_req` directly as an offset; padding rows carry a negative
request index, so that read out of bounds. It is part of the mask now, and `store_mtp_info` treats
a non-positive slot as "no write" for the same reason.
`prepare_verification_metadata` computed its grid from `seqlen_q.max().item()`, a device-to-host
read that the engine rejects inside the forward ("called a synchronizing CUDA operation"). The
grid now comes from the factory's `num_steps`.
Integration is not finished, and `H20-TRITON-PLAN.md` says so plainly: with
`DEEPINFER_OPS_BACKEND=triton` the engine runs the 15 ported kernels into the sampling phase and
then faults on an illegal access inside `expand_sampling_params_tl`, even though that kernel is
element-for-element identical to the baseline in its own A/B cases, and even though the negative
index, the reserved slot and uint64 seeds have all been ruled out. The note records the two
bisection steps worth trying next, the env-var tooling already in place, and the fact that one of
my own edits silently did nothing because the search string carried a trailing space.
Also records the two new Triton traps: an `other=` literal must match the destination dtype, and a
grid must never be derived from a device read.
store_kv_cache and load_paged_kv_cache, both byte-for-byte identical to the TileLang baseline on
the layouts SM90 uses: main/SWA KV (e4m3 values, ue8m0 scale bytes, 528 B per token) and index KV
(128 fp8 + one fp32 scale, 132 B per token). The store also covers fp32 scales, because the index
layout uses them. FP4 caches and the packed-ue8m0 index layout are refused instead of approximated.
Store cases compare the whole cache buffer, including the `-1` invalid-slot path and a block size
(134) that is not a power of two. Load cases cover dequantising to BF16, the index passthrough,
and holes expressed as `src_block_id == -1`.
Two traps worth recording:
The load kernel's `other=` default has to be a float literal. With an fp8 source pointer, `other=0`
fails to type-check with a message about casting an int32 tile to fp8e4nv, which points at the
wrong line.
The load cases originally filled the cache with random bytes. Random bytes decode as fp8 NaN and as
NaN scale factors, and NaN compares unequal to itself, so the harness reported mismatches that said
nothing about the kernel. The cases now build the cache with `store_kv_cache` and compare index
outputs as raw bytes, where a NaN bit pattern is just a byte pattern.
15 kernels and 62 cases agree.
norm, map_lookup, expand_sampling_params, prepare_verification_metadata, copy_block_tables.
Plus the measured parameterisation that made the port tractable: `DEEPINFER_TRACE_KERNEL_ARGS=1`
records each factory's construction arguments, which showed that e.g. `norm_tl` is only ever built
as one combination on this path (d=5120, bf16 source, fp8 out, per-128 fp32 scales, col-major,
round_sf, no residual) rather than the dozen the signature allows.
Three findings worth keeping:
`norm` has to round the fp32 normalised value back to bf16 *before* computing the per-block amax
and the scale. That is what makes its fp8 output bit-identical to quantising the bf16 result
separately, which is how the engine defines the reference. Without it the fp8 bytes scatter by an
ulp wherever a value lands on a rounding boundary -- which looked like a scale bug at first
(224 vs 240) and was not.
With `round_sf=True`, which is the production setting, fp8 output and every scale factor come out
byte-identical. With `round_sf=False` 240 of 327680 bytes differ by one ulp, because `amax/448`
depends on the order of a 5120-element reduction and two compilers will not agree on that order.
For the same reason the bf16 output differs on 10 of 655360 elements by one ulp. Both are
documented, and the A/B harness now picks its tolerance from the dtype instead of one global pair.
`map_lookup` resolves duplicate keys by taking the last match, not the maximum: the TileLang kernel
overwrites on every hit while its own Python reference uses `max`, and the two agree only when keys
are unique. Top-k indices always are, so the Triton version matches the kernel -- and the case
generator now produces unique keys, with a comment saying why.
13 kernels and 53 cases now agree with the baseline.
gather, gather_input_ids_zos, gather_engram_tokens_zos, expand_swa_index, kv_index_cat,
convert_kv_index, load_kv_index, store_mtp_info. 45 cases, every one bitwise identical to the
TileLang baseline through the shared interface function.
`scripts/h20_ops_ab.py` swaps `ops.<name>_tl` between the two backends and runs the public call
once per side, so the comparison exercises the interface layer's shape handling too, not just the
kernel. `scripts/ab_cases.py` holds the cases: real engine widths and dtypes, negative indices
for gather, -1 padding for the index kernels, page-boundary lengths, the slot-0 no-write path,
and the engram image-sentinel rule.
Five Triton semantics traps worth writing down, all documented in the plan:
- `tl.where` evaluates both arms, so `tl.where(i == 0, 0, LOAD(i - 1))` really reads `LOAD[-1]`;
that was an illegal memory access, fixed by clamping the index instead.
- A scalar pointer cannot carry a block-shaped mask; it has to be broadcast with
`+ tl.zeros((BLOCK,), tl.int32)` first.
- Two arms of one kernel cannot use different tile shapes, which is why `kv_index_cat` became two
launches rather than one fused kernel with a column-block branch.
- Triton's `data_ptr()` already includes a view's storage offset, so passing `out[:, w:w+k]` and
adding `w` again double-counts -- that showed up as 0x01010101 in the second half.
- Test-case capacity matters: `store_mtp_info` addresses `slot * nextn + step`, so the pool needs
`(max_slot + 1) * nextn` rows; too few rows let both backends write out of bounds and disagree.
Groundwork for replacing TileLang with Triton on the H20 path, plus the first ported kernel.
The scope is measured rather than guessed. `DEEPINFER_TRACE_KERNELS=<path>` proxies every `*_tl`
factory in the dispatcher and records its first use; one full text-generation run on 8xH20 says
exactly 24 kernels execute. `H20-TRITON-PLAN.md` lists them, in dependency order.
`DEEPINFER_OPS_BACKEND` selects the backend package (unset keeps the platform default, so nothing
changes by omission) and `DEEPINFER_OPS_FALLBACK` names individual kernels that must stay on the
default. Kernels the selected backend does not export at all are taken from the default with a
one-line notice, which lets a partial port run end to end the whole way through instead of only
at the end.
That coexistence is the reason for a parallel backend rather than an in-place rewrite: the
TileLang kernel is a validated baseline, so each Triton kernel can be diffed against it on real
shapes through the same public interface function, which is a much stronger check than comparing
against a fresh Python reference. `scripts/h20_ops_ab.py` does that swap and reports per case.
First kernel: `gather_tl`, all 9 cases bitwise identical (zeros filled for out-of-range indices,
int8/bfloat16 payloads, strided rows, int32/int64 indices, empty and single-row inputs).
Every comparison in this note feeds the reference's chat-template token ids, which is the
worst case for a quantization perturbation -- it is a continuation task with a single
high-probability token, so the margin at the contested position is only 1.64. Asking the same
model in plain text gets the first token right with no drift at all:
prompt: 问题：中国的首都是哪里？\n答案：
output: 北京。\n问题：中国的首都是哪里 tokens [6127, 876, 2056, ...]
Identical on all eight ranks. Without this note the document reads as though ordinary Q&A
carries the ** artefact, which would understate the port.
The note still ended at "next: P4" with an open checklist. Records the outcome instead: the three
decode candidates (KV cache, decode indexer, MTP) are all ruled out with measurements, and the
remaining flip is quantified -- noise std 1.63 against a reference margin of 1.638 at the one
position that differs, out of eight.
Also renumbers the sections, which had a stray unnumbered conclusion between 四 and 五.
Brings in the P4 investigation (per-step logits, KV round-trip and decode-indexer checks, the
recipe probe) along with the docs that record what each check ruled out.
The substantive bug is the debug_dump move: five platform-neutral files imported it through
h20_ops, which registers the SM90 import patches with override=True at module scope, so loading
any of them on SM100 replaced the real flash_mla and deep_select entry points.
The dump hooks live in five platform-neutral files -- `sampler.py`, `mlp.py`, `deepseek.py`,
`decoder_layer.py`, `model_runner.py` -- and all of them imported
`deepinfer.h20_ops.debug_dump`. That is wrong on every platform except SM90: `h20_ops/__init__.py`
calls `extend_import_patches(override=True, ...)` at module scope, so importing it from shared code
installs the SM90 import shims globally. On an SM100 machine that replaces
`flash_mla.decode_mla_mqa_sparse`, `flash_mla.prefill_attn_varlen_sparse` and
`deep_select.interface.topk` with the Hopper implementations, and on NPU it would do the same
before the Ascend shims get a chance to matter.
`debug_dump` itself has nothing platform-specific in it -- `os`, `pathlib`, `numpy`, `torch` --
so it moves to `deepinfer/debug_dump.py` and the five imports follow. Verified that importing it
no longer pulls `h20_ops` into `sys.modules`, and re-ran the end-to-end check on H20: same output
token for token, `[17607, 2793, 3139, 666, 6127, 666, 320, 1]`.
This bug is already in the open PR, so it needs to land before that PR is worth reviewing further.
P4.2 completes the decode comparison, and it closes the investigation rather than opening
another thread.
The one prerequisite was `DEEPINFER_USE_NEXTN_DECODE=0`. Its default is 5, so MTP speculative
decoding is on by default, which means several draft/verify rows per step and `logits[0]` no
longer addresses the position you think it does -- that is why an earlier attempt misread the
pos-15 vector as pos-14 and looked like a missing step. With MTP off, one row per step, clean
alignment. Disabling it is also the right comparison, since the reference `generate.py` has no
MTP at all.
Per position, 11-token prompt, 8 decode steps:
pos ours ref ours margin ref margin d margin
11 17607 17607 10.695 9.861 +0.835 (prefill)
12 2793 2793 8.448 10.328 -1.880
13 3139 3139 15.276 11.946 +3.330
14 666 6127 1.480 1.638 -0.157 <-- only flip
15 6127 6127 15.950 15.109 +0.842
16 666 666 20.118 17.874 +2.244
17 320 320 9.829 10.402 -0.573
18 1 1 12.262 11.091 +1.171
d margin: mean +0.711, std 1.628, max 3.330. Seven of eight positions are right. The single
flip sits where the reference's own margin is 1.638, against 9.9-17.9 everywhere else -- the
noise (std 1.63) and that margin are the same size, so it is a coin toss. The reference takes
`6127` (北京) first, we take `666` (`**`) first, and both go on to produce the same text:
`**北京**` is valid markdown and means the same thing.
So the remaining difference is per-128 quantization, which SM90 gets no choice about, landing on
a position where the model itself is undecided. Not something to fix: the only lever is the
lossless e8m0 per-32 recipe from P3.6, and DeepGEMM's `gran_k == 128` assertion forecloses it
(see `H20-DRIFT-ANALYSIS.md`).
Also rules out two more suspects. MTP is not the cause: with `USE_NEXTN_DECODE=0` the output is
byte-identical, `[17607, 2793, 3139, 666, 6127, 666, 320, 1]`. And no logits post-processing is
involved: at pos 14 our top-1 is 666 with 6127 second by 1.480, so 6127 cannot have been masked
-- it would not be the runner-up if it had been.
Records one methodological correction: the original criteria assumed prefill and decode should
agree, and treated a large rel_err between them as a decode defect. The reference's own
prefill-to-decode transition moves the margin by about 2x, because the KV cache goes through a
quantized round trip. The meaningful comparison is d margin against that position's margin, not
rel_err on its own. The old criteria stay in the document as process history, marked revised.
Two closed-loop checks, both clean, which finally retires the two standing suspects.
P4.3 KV cache. The repo's own store/load tests pass for the fp32-scale index layout SM90 uses
(6 each). Running every layout gives 48 passed / 12 failed, and every failure is a ue8m0 index
layout -- `FP4_INDEX_KV_FORMAT` or `MXFP8_INDEX_KV_FORMAT`, neither of which SM90 selects. Those
failures are in the test's reference, not the kernel: it reinterprets ue8m0 bytes as
`PLATFORM.fp8_sf_type`, which is float32 here, so `index_copy_` hits a dtype mismatch.
`scripts/h20_kv_roundtrip.py` then closes the loop without any reference at all -- write known
BF16 through the paged cache, read it back, compare. Both SM90 layouts, block 32 and 64, 1/233/
1029 tokens:
V41_FORMAT (main/swa KV) rel 0.0266-0.0295 cos 0.9996
FP8_FLOAT32_SF_INDEX_KV_FORMAT (132B) rel 0.0266-0.0327 cos 0.9996
That is the fp8 floor, the same order as the 2.7e-2 on weights, so the cache is not where the
decode error comes from and the 132-byte layout question from P3.9 is closed.
One trap: the first version of the script compared raw fp8 values against BF16 and reported
rel=126. Index layouts are deliberately *not* dequantized by `load_paged_kv_cache` -- the
consumer is a kernel that wants bytes plus scales. Multiplying by the scale gives 0.027. The
127x was the test, not the implementation.
P4.4 the decode indexer. `_decode_p_sum_sm90` reimplements DeepGEMM's paged MQA logits in torch
because Hopper's kernel requires 64-token pages; nothing else exercised it. Compared against a
float64 reference that shares no code -- explicit per-token offsets instead of strided views --
it agrees to rel~1e-7 / cos 1.0 across page sizes 32 and 64, boundary-crossing lengths,
multi-page contexts and 32 heads. The `q_sf[..., :1]` dequantization is also sound: head_dim is
128, equal to the granularity, so one scale per (row, head) is correct.
What this changes. `logits_prompt` from the max_tokens=1 run puts the prefill error at cos
0.9846 with an absolute error of 3.30 on the winning token, against a reference margin of 3.112:
the margin is thin because the error is the same size as the margin, not by coincidence. The
reference's own prefill-to-decode transition also costs about 2x margin (3.112 -> 1.638) purely
from the quantized KV round trip. Our 0.256 halved would still be positive, yet it flips, so
roughly 0.5-1.0 of logit-level error in decode remains unexplained -- now narrowed to attention
over the cached main KV and the SWA-window / compressed-position mapping.
Also corrects an earlier overstatement: "the difference is introduced by decode" does not mean
decode has a discrete bug, and P4.3/P4.4 show its two main suspects do not.
One more dump trap worth recording: with max_tokens > 1, `logits_prompt` is contaminated too.
mysteps6 (max_tokens=7) gives cos 0.68 and a 16.5 top-1 error for the same input where mylogits3
(max_tokens=1) gives 0.9846 and 3.30. Prefill logits are only trustworthy from a
max_tokens=1 run.
Adds `record_step` so every decode step is kept instead of just the first, a `--force_ids` mode
to `ref_next_token.py` that mirrors `generate.py`'s incremental loop, and four fixes to the dump
path without which the data was unusable:
- filenames carry `_rank<R>_pid<P>`: with data parallelism every engine is its own
single-process group and reports RANK=0, so eight processes wrote the same file and the last
one won;
- decode reads row 0, not the last row: the batch is padded and the padding rows hold logits for
placeholder tokens, so the last row was a constant `top1=69, margin=0.0528`;
- only the driver rank records: a follower only joins the embedding all-gather, so its logits are
placeholders that got recorded dozens of times and buried the real steps;
- `DEEPINFER_DUMP_TOKENS` gating plus `profiling()` now cover both profile paths.
The conclusion itself comes from a black-box four-way comparison rather than the dumps, because
per-step dumps still do not line up with the emitted token sequence (`max_tokens=7` produced 75
sampler calls). At the same 14-token context:
reference prefill 6127 (北京) correct
reference decode 6127 (北京) correct
our prefill 6127 (北京) correct
our decode 666 (**) wrong
So prefill is right and the remaining flip is introduced by decode alone. Note the reference's
own prefill and decode margins differ (3.112 vs 1.638), so the two paths are not expected to be
identical -- KV cache goes through a quantized round trip -- but they should not disagree on the
direction.
Remaining suspects are therefore narrowed to decode-only state: the KV cache round trip (P4.3)
and the indexer scored in torch versus the TileLang prefill kernel (P4.4), which have never
been diffed against each other.
Comparing final logits was impossible before: `ref_layer_dump.py` read `logits[0, -1]`, but the
head already collapsed the sequence axis (`x = x[:, -1]`), so logits are `[b, vocab]` and that
index silently stored a 0-dimensional scalar. Every earlier dump had an unusable
`final_row_logits`.
With that fixed, the comparison reframes the remaining drift. On the same 14 tokens:
prefill top-1 = 6127 (北京) margin(6127 - 666) = +0.256
reference top-1 = 6127 (北京) margin = +3.112
Our prefill picks the right token; the error only compresses the margin, by ~12x. But the
end-to-end run, which reaches that context via decode steps, picks 666 (`**`) instead. Same
context, different path, different answer -- so the remaining drift is in **decode**, not in the
per-128 quantization. The quantization floor is still 2.7e-2 and still irreducible
(`H20-DRIFT-ANALYSIS.md` has the DeepGEMM assertion that forces `gran_k == 128` on SM90), but it
is not what flips the token.
Two traps had to be handled to get numbers that mean anything:
- `execute_model_profile_run` runs a full forward + sampler on synthetic tokens. The layer dumps
were filtered by their prompt-length gate; logits are not, because their row count is the
number of *selected* tokens (1). The first attempt captured that step and produced cos = 0.34.
`debug_dump.profiling()` now suspends capture for the whole profile step.
- With `max_tokens=1` both a prompt and a decode step run, and only the prompt one is comparable
with a prefill-only reference. They are dumped as `logits_prompt` / `logits_decode`.
Also re-runs the layer comparison now that the expert-scale fix has landed: layer 0 went from
1.27e-1 to 2.88e-2 and layer 39 sits at 2.00e-1, so the error curve flattened by an order of
magnitude.
