Skip to content

docs(examples): Nemotron-3.5-Lightning-30B-A3B, one launch a decode step - #143

Merged
zhen8838 merged 1 commit into
tile-ai:mainfrom
zhen8838:docs/example-nemotron-3.5-lightning
Aug 31, 2026
Merged

docs(examples): Nemotron-3.5-Lightning-30B-A3B, one launch a decode step#143
zhen8838 merged 1 commit into
tile-ai:mainfrom
zhen8838:docs/example-nemotron-3.5-lightning

Conversation

@zhen8838

Copy link
Copy Markdown
Collaborator

Why

  • The four existing examples each keep one agent run whole, but in all of them
    the authored HIR is a set of per-op @funcs and the fusion, where there is
    any, lives only on the runtime side. So check compares a description of the
    program rather than the program that runs.
  • Nothing in examples/ yet shows a model whose decode step is a single
    cooperative launch, and nothing shows the low-level TileLang surface
    (T.tma_load, T.mbarrier_*, T.wgmma_gemm) being used end to end.

What

  • Adds examples/nemotron_3_5_lightning_30b_a3b-tilelang: a 52-layer
    Mamba2 / attention / MoE hybrid where the authored HIR is the mega program
    the whole decode step is one @func with 474 parameters, stage boundaries are
    mesh-level barriers rather than @func boundaries, and the twin runs it as a
    single launch with no CUDA graph.
  • Ships both generators next to what they emit (gen_model.py, gen_kernel.py);
    the rule stated in the README is to change the generator. TileLang rewrites
    every for in a @T.prim_func into a device loop and does not bind a nested
    def, and HIR's for is a runtime loop over one weight tensor, so trace-time
    repetition has to be written out.
  • attention.py keeps the placement ladder that chose the length dispatch, with
    the crossover read off analyze rather than tuned.
  • ISSUES.md plus repro/ and kbench/: 7 findings in TileFoundry and 8 in
    TileLang, each with a minimal repro. TF-3 is marked fixed by fix(diagnostics): preserve authored failure context #141.
  • Adds one line to the README's Latest News.

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

  • No library change. examples/ is pruned from the sdist and excluded from
    ruff, and nothing in CI runs it.
  • Verified at 74abc97 in a fresh venv from that wheel: check_all.py 59 of 59
    outputs, compare_transformers.py 64 of 64 tokens identical,
    tools/launches.py 1 against 3212, tools/ctxcmp.py agreeing at 9 of 11
    sampled 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

  • check covers only the short-context arm. Splitting a long reduction over
    a 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.md section 6). The long arm rests on the
    transformers comparison and on mega-against-op-by-op instead. The README says
    this where the number is. When the evaluator covers it,
    check_all.py --ctx-full 262016 is the test that says so.
  • The 262080 timing point is taken by dialling the cache counter rather than
    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.
  • Nothing re-runs any of this automatically, as with the other four examples.

@zhen8838
zhen8838 force-pushed the docs/example-nemotron-3.5-lightning branch from ba7fe56 to 6e3ab4a Compare August 31, 2026 03:45
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
zhen8838 force-pushed the docs/example-nemotron-3.5-lightning branch from 6e3ab4a to 7339b51 Compare August 31, 2026 03:54
@zhen8838
zhen8838 merged commit ce84368 into tile-ai:main Aug 31, 2026
1 check passed
@zhen8838
zhen8838 deleted the docs/example-nemotron-3.5-lightning branch August 31, 2026 03:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant