coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 08:46:21 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/scripts
History
loopyt f25e5d8b6c feat(h20): run sparse MLA on FlagGems' Triton kernel, so TileLang is no longer required 
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`).
2026-10-25 14:02:58 +08:00
..
simulate_no_tilelang feat(h20): make the pure-torch kernel references usable without TileLang 2026-10-22 16:08:46 +08:00
sparse_mla_port feat(h20): run sparse MLA on FlagGems' Triton kernel, so TileLang is no longer required 2026-10-25 14:02:58 +08:00
tile_kernels_patch fix(h20): let the reference backend sync, and stop the reference raising under the guard 2026-10-24 19:10:49 +08:00
ab_cases.py test(h20): catch a factory that returns nothing before eight ranks do 2026-10-22 14:16:41 +08:00
build_vendor.sh initial commit 2026-10-20 12:09:42 +08:00
build.sh initial commit 2026-10-20 12:09:42 +08:00
compare_layers.py fix(sm90): expand expert scales on the contraction axis only 2026-10-21 16:50:31 +08:00
copy_site_packages.py feat(h20): make the FlagTree image setup reproducible instead of a shell history 2026-10-22 13:20:03 +08:00
enable_torch_kernel_refs.sh feat(h20): make the pure-torch kernel references usable without TileLang 2026-10-22 16:08:46 +08:00
flagtree_triton_probe.py docs(h20): record what a FlagTree/FlagOS image can and cannot do 2026-10-22 10:53:47 +08:00
format_csrc.sh initial commit 2026-10-20 12:09:42 +08:00
gen_stub.sh initial commit 2026-10-20 12:09:42 +08:00
h20_ask.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_attn_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_e2e_report.sh docs(h20): point at the setup script, and let the report script carry its own defaults 2026-10-22 13:20:08 +08:00
h20_engram_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_expert_weight_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_index_decode_check.py test(sm90): rule out the KV cache and the decode indexer 2026-10-21 20:02:19 +08:00
h20_kv_roundtrip.py test(sm90): rule out the KV cache and the decode indexer 2026-10-21 20:02:19 +08:00
h20_ops_ab.py test(h20): catch a factory that returns nothing before eight ranks do 2026-10-22 14:16:41 +08:00
h20_philox_check.py feat(triton): reproduce cuRAND's Philox stream, so the sampler can be ported 2026-10-22 08:37:08 +08:00
h20_quant_check.py test(sm90): decode FP4 by table lookup in the quant check 2026-10-21 17:15:16 +08:00
h20_real_requant_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_recipe_probe.py debug: capture final logits, and keep the profile forward out of the dumps 2026-10-21 17:44:21 +08:00
h20_rng_dump.py feat(triton): reproduce cuRAND's Philox stream, so the sampler can be ported 2026-10-22 08:37:08 +08:00
h20_smoke_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
h20_state_check.py test(h20): check what the model actually says, in a file rather than on the command line 2026-10-22 10:53:01 +08:00
h20_topk_check.py style(sm90): satisfy ruff and add the missing type annotations 2026-10-21 17:09:08 +08:00
ref_layer_dump.py debug: capture final logits, and keep the profile forward out of the dumps 2026-10-21 17:44:21 +08:00
ref_next_token.py debug: per-step decode dumps, and prove the flip is decode-only 2026-10-21 19:07:49 +08:00
setup_flagtree_image.sh feat(h20): make the FlagTree image setup reproducible instead of a shell history 2026-10-22 13:20:03 +08:00
verify_tile_kernels_patch.py feat(h20): make the pure-torch kernel references usable without TileLang 2026-10-22 16:08:46 +08:00