docs(examples): Nemotron-3.5-Lightning-30B-A3B, one launch a decode step - #143
Merged
zhen8838 merged 1 commit intoAug 31, 2026
Merged
Conversation
zhen8838
force-pushed
the
docs/example-nemotron-3.5-lightning
branch
from
August 31, 2026 03:45
ba7fe56 to
6e3ab4a
Compare
A 52-layer Mamba2 / attention / MoE hybrid taken end to end on tilelang, where
the authored HIR **is** the mega program: the whole decode step is one `@func`
with 474 parameters, its stage boundaries are mesh-level barriers rather than
`@func` boundaries, and it runs as a single cooperative launch.
1 launch a step against 3212 on the op-by-op path
64 of 64 greedy tokens identical to transformers
287.4 tok/s at ctx 32 SGLang 294.8, same card, same day
231.6 tok/s at ctx 262080 278.2
2.37x-2.44x off the measured bandwidth floor, flat across nine lengths
It does not beat SGLang -- 97.5% of it short, 83.2% at 262080 -- and the README
leads with that. What the example is for is the shape: because both sides of the
twin are the same program, `check` compares the thing that actually runs.
Two generators ship alongside what they emit, and the rule is change the
generator. TileLang rewrites every `for` in a `@T.prim_func` into a device loop
and does not bind a nested `def`; HIR's `for` is a runtime loop over one weight
tensor. Repetition that has to happen at trace time -- a prefetch prologue, a
fixed count of accumulator registers, one body per stage kind, 52 layers sharing
neither weights nor shape -- has to be written out.
## Verified
At `74abc97`, in a fresh venv from that wheel:
check_all.py 59 of 59 outputs pass, bounds derived per output
compare_transformers.py 64 of 64 tokens identical
tools/launches.py 1 launch a step against 3212
tools/ctxcmp.py 9 of 11 lengths agree; the 2 that differ are where
the reference's own top-2 gap is one bf16 ulp
Six of the nine table points were re-measured and agree to 0.1%. `reports/` keeps
that evidence.
**`check` only covers the short-context arm, and the README says so where the
number is.** Splitting a long reduction over a worker axis is what long-context
attention needs, and the evaluator models one mesh participant, so at
`ctx_full > 0` it refuses. The long arm rests on the other two rows instead.
When that is fixed upstream, `check_all.py --ctx-full 262016` is the test that
says so.
## Also here
`ISSUES.md` -- 7 things in TileFoundry and 8 in TileLang, each with a minimal
repro under `repro/` or `kbench/`. TF-3 (a diagnostic naming neither file nor
line) is marked fixed by tile-ai#141. `SGLANG_BASELINE.md` carries the baseline's method
and the four ways measuring it went wrong first.
zhen8838
force-pushed
the
docs/example-nemotron-3.5-lightning
branch
from
August 31, 2026 03:54
6e3ab4a to
7339b51
Compare
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.
Why
the authored HIR is a set of per-op
@funcs and the fusion, where there isany, lives only on the runtime side. So
checkcompares a description of theprogram rather than the program that runs.
examples/yet shows a model whose decode step is a singlecooperative launch, and nothing shows the low-level TileLang surface
(
T.tma_load,T.mbarrier_*,T.wgmma_gemm) being used end to end.What
examples/nemotron_3_5_lightning_30b_a3b-tilelang: a 52-layerMamba2 / attention / MoE hybrid where the authored HIR is the mega program —
the whole decode step is one
@funcwith 474 parameters, stage boundaries aremesh-level barriers rather than
@funcboundaries, and the twin runs it as asingle launch with no CUDA graph.
gen_model.py,gen_kernel.py);the rule stated in the README is to change the generator. TileLang rewrites
every
forin a@T.prim_funcinto a device loop and does not bind a nesteddef, and HIR'sforis a runtime loop over one weight tensor, so trace-timerepetition has to be written out.
attention.pykeeps the placement ladder that chose the length dispatch, withthe crossover read off
analyzerather than tuned.ISSUES.mdplusrepro/andkbench/: 7 findings in TileFoundry and 8 inTileLang, each with a minimal repro. TF-3 is marked fixed by fix(diagnostics): preserve authored failure context #141.
Measured on one H200 at 1500 MHz: 1 launch a step against 3212 op by op;
287.4 tok/s at context 32 and 231.6 at 262080, against SGLang's 294.8 and 278.2
measured the same day on the same card; 2.37x–2.44x off the measured bandwidth
floor, flat across nine context lengths. It does not beat SGLang, and the
example's README leads with that.
Contract
examples/ispruned from the sdist and excluded fromruff, and nothing in CI runs it.
74abc97in a fresh venv from that wheel:check_all.py59 of 59outputs,
compare_transformers.py64 of 64 tokens identical,tools/launches.py1 against 3212,tools/ctxcmp.pyagreeing at 9 of 11sampled lengths (the 2 that differ are where the reference's own top-2 gap is
one bf16 ulp). Six of the nine table points were re-measured and agree to 0.1%.
reports/keeps that evidence.Risk
checkcovers only the short-context arm. Splitting a long reduction overa worker axis is the placement long-context attention needs, and the evaluator
models one mesh participant, so it refuses at
ctx_full > 0(
docs/spec/evaluator.mdsection 6). The long arm rests on thetransformerscomparison and on mega-against-op-by-op instead. The README saysthis where the number is. When the evaluator covers it,
check_all.py --ctx-full 262016is the test that says so.walking a real decode, because the walk did not finish; the two methods agree
to 0.4% at all eight lengths where both can run, and the README explains why.