diff --git a/deepinfer/h20_ops/vendor/flashmla_sparse.py b/deepinfer/h20_ops/vendor/flashmla_sparse.py index 6399965..d2aecb9 100644 --- a/deepinfer/h20_ops/vendor/flashmla_sparse.py +++ b/deepinfer/h20_ops/vendor/flashmla_sparse.py @@ -39,13 +39,19 @@ # # Why vendored instead of depending on `flag_gems`: # -# 1. `flag_gems/__init__.py` registers ATen implementations process-wide -# (`torch.library.Library("aten", "IMPL")` then `runtime.replace_customized_ops`). Importing -# the package to reach one kernel would silently swap operators under the whole engine, which -# changes numerics of code that has nothing to do with attention. This module imports only -# torch and triton, so it cannot do that. -# 2. The released 5.0.2 file has no `flag_gems` internal imports at all, so there is nothing to -# gain from the dependency. +# 1. Import cost and a native dependency. `flag_gems/__init__.py` star-imports every op module and +# pulls in Triton TLE, which loads `libflagcx` through ctypes -- an .so this deployment does not +# otherwise need. A single attention kernel should not drag that in. This module imports only +# torch and triton. +# 2. This file has no `flag_gems` internal imports at all, so there was nothing to gain from the +# package for this operator. +# +# Correction, recorded because an earlier version of this header said the opposite: importing +# `flag_gems` does *not* register ATen implementations. It creates the `aten_lib` handle and imports +# the op modules, but registration happens only when `enable()` or `only_enable(include=[...])` is +# called. Measured by hashing `torch._C._dispatch_dump_table` before and after `import flag_gems`: +# unchanged for `aten::rms_norm`, `aten::silu`, `aten::softmax`, `aten::topk`. So a blanket +# operator swap was never a risk of merely importing the package -- the risk lives in `enable()`. # # Why 5.0.2 rather than FlagGems master: master's copy is ~1200 lines with a TLE fast path gated on # triton internals, and its accumulation is fp16. This one is 329 lines, accumulates in fp32, and diff --git a/scripts/sparse_mla_port/README.md b/scripts/sparse_mla_port/README.md index 075c3e6..6a56b15 100644 --- a/scripts/sparse_mla_port/README.md +++ b/scripts/sparse_mla_port/README.md @@ -9,11 +9,20 @@ operator; it is now vendored, verified against the TileLang kernel, and wired in `flagos-ai/FlagGems`, file `src/flag_gems/fused/flashmla_sparse.py`, function `flash_mla_sparse_fwd`. Vendored as `deepinfer/h20_ops/vendor/flashmla_sparse.py` from the **PyPI -5.0.2 wheel**, which has no `flag_gems` internal imports at all -- only torch and triton. That -matters: importing `flag_gems` itself registers ATen implementations process-wide -(`torch.library.Library("aten", "IMPL")` then `runtime.replace_customized_ops`), so depending on the -package to reach one kernel would swap operators under the whole engine. The file header records the -source md5 and the two local deviations. +5.0.2 wheel**. Two reasons, plus a correction to what this file said previously: + +* Import cost: `flag_gems/__init__.py` star-imports every op module and pulls in Triton TLE, which + loads `libflagcx` through ctypes -- an .so this deployment does not otherwise need. One attention + kernel should not drag that in. +* The released file has no `flag_gems` internal imports at all (only torch and triton), so the + package bought nothing for this particular operator. + +Correction: importing `flag_gems` does **not** register ATen implementations, contrary to the earlier +wording here. It creates the `aten_lib` handle; registration happens only in `enable()` or +`only_enable(include=[...])`. Measured by hashing `torch._C._dispatch_dump_table` before and after +`import flag_gems`: unchanged for `aten::rms_norm`, `aten::silu`, `aten::softmax`, `aten::topk`. The +blast radius lives in `enable()`, not in the import. That distinction is what makes a broader +FlagGems integration tractable -- see `H20-FLAGGEMS-PLAN.md`. Master's copy (~1200 lines, TLE fast path, fp16 accumulate) also works and was verified first; 5.0.2 is used instead because it is a release, is 329 lines, accumulates in fp32, and already guards the