[jhchouuu]

CI pins flydsl<0.3 because the cco FlyDSL examples import flydsl.expr.buffer_ops, which 0.3.0 removed. aiter now pins flydsl==0.4.0 (requirements.txt, setup.py, pyproject.toml), so mori was only tested against a FlyDSL its users no longer get. This upgrades CI to 0.4.0 and keeps the older versions mori still supports covered nightly.

Changes

• mori.cco.device.flydsl.buffer (new): create_buffer_resource_from_addr / buffer_load / buffer_store and the CM_* cache policies, moved from ops/gemm_ar/_compat.py, which re-exports them so the gemm_ar kernels and benchmark keep their imports. This is the shim that keeps cache_modifier working on 0.3.0+; EPv2's flydsl_compat.py accepts and ignores it there, so the examples could not use that one. On 0.2.x it is flydsl.expr.buffer_ops; only that module being missing selects the other path, and any other import error is raised. On 0.3.0+ it calls the generated rocdl ops and probes whether aux is a positional value (0.3.0) or a keyword attribute (0.3.1+); any other signature fails at import. It uses the generated ops rather than flydsl.expr.rocdl because from 0.3.1 that module wraps the load (only) and drops an aux that is neither an int nor an Attribute.
• mori.cco.device.flydsl.signals (new): the u32 signal-slot accessors (signal_ptr, signal_store_u32, signal_load_u32, local_load_u32, local_store_u32, i32_type), moved unchanged from ops/gemm_ar/_compat.py for the same reason: example 10 (example(cco): cco-sdma all-gather with start/wait (async mode) in pcp scenario #798) imported them from that private module. _compat.py re-exports them and keeps what is tuned for gemm_ar (the tile-counter atomics, the copy-engine release fence, wave_uniform_i64). Named signals, not signal, so it cannot shadow the stdlib module.
• cco examples 04/05/07/08 use buffer instead of flydsl.expr.buffer_ops, and example 10 imports from signals instead of mori.ops.gemm_ar._compat. Examples 04 and 05 now use the shared CM_* constants. Their private ones were wrong on every version, main included: 2 and 3 are nt and sc0 nt in LLVM's CPol numbering, which both paths pass through as aux, so the signal store and the spin load did not bypass L2.
• tests/python/cco/test_flydsl_buffer.py (new): compiles a buffer load and store for each CM_* value and checks the sc0/sc1/nt bits in the emitted ISA (gfx94x/gfx95x).
• flydsl_compat.py: the docstring says why EPv2 keeps its own shim for now. It also bridges vector and T, and honouring its cache_modifier on 0.3+ would change the scatter combine's codegen, which needs a perf check.
• docker/Dockerfile.dev: new FLYDSL_VERSION build arg, default 0.4.0, declared next to the flydsl RUN so changing it does not rebuild the ROCm layers.
• ci_cco.yml: also runs examples 07 and 10 (with MORI_ENABLE_SDMA=1), so every FlyDSL example runs on 0.4.0, and the cache-policy test.
• nightly.yml: new flydsl-compat job builds the image with 0.2.4 (the flydsl.expr.buffer_ops path) and 0.3.0 (the last release that takes aux positionally; 0.3.1+ is the keyword form CI's 0.4.0 covers) and runs test_sdma_api.py, the gemm_ar numerics (test_gemm_ar*.py, test_flydsl_ar.py), ops-v2 test_dispatch_combine_v2_intranode.py (FlyDSL backend), the cache-policy test and examples 04/05/07/08/10. The legs run one at a time, because each uses all 8 GPUs. They use the runner's RDMA settings from ci_cco.yml, and each removes its image afterwards.
• READMEs: install flydsl==0.4.0.

The library itself needed no change: outside the two shims nothing uses a 0.2-only API.

Testing

8x MI355X, CI image (ROCm 7.14), CI build flags, every FlyDSL-related ci_cco.yml step plus example 07:

	0.4.0	0.3.0	0.2.4
examples 01/02/04/05/07/08	pass	pass	pass
example 10, 2 and 8 ranks	pass	pass	pass
test_sdma_api.py	13/13	13/13	13/13
gemm_ar numerics	114 passed, 1 skipped	114 passed, 1 skipped	114 passed, 1 skipped
test_gemm_ar_bench.py	4/4	4/4	4/4
ROCm import order, Triton, JIT binding	pass	pass	pass
ops-v2 intranode, FlyDSL/HIP parity, routing snapshot, bench_ep	pass	pass	pass

On the final branch (main with #798 merged in, the signal accessors moved, the examples on the shared constants), I re-ran examples 04/05/07/08/10 and the cache-policy test on all three versions; example 10 was run as the CI command is written. I also re-ran the gemm_ar numerics on 0.4.0: 114 passed, 1 skipped. All three versions emit the same cache bits, for both load and store:

CM_*	0	1	2	16	17
bits	none	sc0	nt	sc1	sc0 sc1

With CM_SC1 set back to 4, the old wrong value, the test fails.

0.3.2 was also run (examples, SDMA API, gemm_ar, Triton, JIT binding: pass). The new example commands were run as written: without MORI_ENABLE_SDMA example 07 exits 1, with it both the ci_cco and nightly forms pass.

On main the four examples fail on 0.3.x/0.4.0 with cannot import name 'buffer_ops'; everything else above already passed there. Performance: gemm_ar is unchanged across versions; the EPv2 FlyDSL backend's combine is 3-6% faster on 0.3.4.1/0.4.0 than on 0.2.4 (bench_ep, M=512..8192), dispatch unchanged.

Note: test_fused_is_stable_across_repeats hung once on 0.3.2 in a full-suite run. It did not reproduce in 4 isolated runs (2 on main, 2 on this branch) or a second full-suite run, so it looks like an intermittent pre-existing issue rather than a version or branch effect -- worth knowing if the nightly job flakes.

🤖 Generated with Claude Code

[kawhil-amd]

Thanks for the shared buffer.py; having the cache-policy encoding in one place is a good change. The first comment is the one I would resolve before merging: the migrated examples appear to pass cache-modifier values from the old scheme into the new helper. The rest are smaller points on error handling, tests, coupling and CI. These come from reading the diff. I did not run them on hardware.

One point has no line in this diff because the file is unchanged: ops/dispatch_combine_v2/flydsl_compat.py is a second FlyDSL 0.3+ buffer shim. It ignores cache_modifier on 0.3+ and probes with ModuleNotFoundError, while buffer.py uses ImportError. The two shims will drift apart. Please reuse buffer.py there, or note why EPv2 must stay separate.

[kawhil-amd]

_CM_SC1 is still the old private constant (_CM_SC1 = 2 at line 68, _CM_SC0_SC1 = 3 at line 69), but this call now goes through buffer.buffer_load. On FlyDSL 0.3+, buffer.py passes aux straight through as the LLVM CPol operand. By the table in buffer.py, 2 is CM_NT there, not "bypass L2" (CM_SC1 = 16). The store at line 91 has the same issue: 3 is sc0|nt, not sc0|sc1 (CM_SC0_SC1 = 17).

If that is right, a spin on a slot that a peer writes can keep reading a stale L2 line and hang, and a signal store may not reach the peer. A 2-rank CI pass does not show that the policy is right. Please confirm which scheme buffer_ops uses on 0.2.x and which one the raw ops use on 0.3+, and use buffer.CM_SC1 and buffer.CM_SC0_SC1 here. If the 0.2.x buffer_ops uses a different numbering, buffer.py also needs a per-version mapping, because it passes CM_* to buffer_ops unchanged there.

[jhchouuu]

You're right, and it is older than this PR. 0.2.x buffer_ops passes cache_modifier to RawPtrBufferLoadOp/StoreOp as aux unchanged (aux_flags = _create_i32_constant(cache_modifier)), so both paths use LLVM's CPol numbering and no per-version mapping is needed. That also means _CM_SC1 = 2 and _CM_SC0_SC1 = 3 were nt and sc0 nt on every version, main included.

Example 05 now uses buffer.CM_SC0_SC1 for the signal store, buffer.CM_SC1 for the spin load and buffer.CM_CACHED for the output. The new tests/python/cco/test_flydsl_buffer.py compiles a load and a store for each CM_* value and reads the bits back from the ISA. On gfx950, 0.4.0, 0.3.0 and 0.2.4 all give the same result, for both load and store:

CM_*	0	1	2	16	17
bits	none	sc0	nt	sc1	sc0 sc1

[kawhil-amd]

Same issue as in example 05: _CM_UNCACHED = 3 is a private constant. The shared module defines CM_SC0_SC1 = 17. Please use the shared constant.

[jhchouuu]

Done: the private constant is gone, and the store uses buffer.CM_SC0_SC1.

[kawhil-amd]

flydsl_compat is not part of this PR, but it is the second shim for the same problem. See the review summary: it should reuse this module or say why it cannot.

[jhchouuu]

EPv2 keeps its own shim for now, and flydsl_compat's docstring now says why. It also bridges vector and T, and making it honour cache_modifier would change EPv2's codegen on 0.3+ (the scatter combine's stage-3 loads ask for nt). That needs a perf check first, so I'd rather do it as a follow-up than in the version bump. The import probes do agree now: buffer.py also catches ModuleNotFoundError, and only for flydsl.expr.buffer_ops.

[kawhil-amd]

This says gemm_ar's cross-GPU barrier and SDMA signals need the cache policy. gemm_ar/_compat.py and the note at line 67 below say those are atomics and use no aux bits. Please make the two statements agree.

[jhchouuu]

Fixed. The docstring now says gemm_ar's use is the kernels_sdma loads of data a peer's copy engine delivered (CM_SC1). It also says the barrier and signal slots are atomics in mori.cco.device.flydsl.signals and use no aux bits.

[kawhil-amd]

except ImportError also catches an ImportError raised from inside a working 0.2.x buffer_ops or flydsl.expr.typing import. The code then silently switches to the raw rocdl path. Consider catching only the buffer_ops import, so another missing module fails loudly.

[jhchouuu]

Done. It catches ModuleNotFoundError and re-raises unless e.name == "flydsl.expr.buffer_ops", so any other missing module fails loudly.

[kawhil-amd]

The version notes are inconsistent. The comment says aux moved in 0.3.1 and that sglang has 0.3.2, and aiter's venv has 0.3.0. Dockerfile.dev says aiter pins 0.4.0. The nightly matrix tests only 0.2.4 and 0.3.0, so the keyword path is covered only by 0.4.0. If the real boundary is 0.3.2, the nightly comment is wrong and the cutoff is not tested. Also, _AUX_IS_KWARG only recognises two spellings of aux. If the signature changes again, the probe may pick the positional branch and pass an i32 value where an attribute is expected. There is no unit test for buffer.py. A small test that checks which aux operand is emitted for each CM_* value would have caught the silent bug that the comment at line 65 describes.

[jhchouuu]

I checked the wheels' generated rocdl ops. aux is a positional i32 Value on 0.3.0 and a keyword-only IntegerAttr from 0.3.1 on (the same on 0.3.2 and 0.4.0). So the boundary is 0.3.1: nightly's 0.3.0 leg covers the positional path, and CI's 0.4.0 covers the keyword path. I removed the venv remarks from the comment.

Two more changes came out of this:

• The probe is strict now. It accepts only "keyword-only on both ops" or "positional on both", and otherwise raises ImportError with the FlyDSL version.
• It reads the generated ops in flydsl._mlir.dialects.rocdl, not flydsl.expr.rocdl. Running the strict probe showed that from 0.3.1, flydsl.expr.rocdl wraps raw_ptr_buffer_load but not the store. That wrapper also silently drops an aux that is neither an int nor an Attribute.

I added the unit test: tests/python/cco/test_flydsl_buffer.py. It runs in ci_cco.yml (0.4.0) and in both nightly flydsl-compat legs (0.2.4 and 0.3.0). It only runs on gfx94x/gfx95x, because sc0/sc1/nt are the CDNA3/4 spellings. As a check, I set CM_SC1 back to the old value 4, and the test fails with CM=4 load: [[]] != ['sc1'].

[kawhil-amd]

_compat now imports a private name (_dtype) from the cco device package. Importing mori.cco.device.flydsl.buffer also runs mori.cco.device.flydsl/__init__ first, which loads the bitcode locator, handles and _bindings. So any failure in the cco device bindings now breaks import mori.ops.gemm_ar, even for callers that only need buffer ops. A shared utility fits better in a neutral location, and _dtype should be public or stay private to buffer.py.

[jhchouuu]

_dtype is no longer imported; wave_uniform_i64 uses the public i32_type().

On the coupling: since d9ba2c4 the signal helpers live in mori.cco.device.flydsl.signals, so example 10 no longer imports gemm_ar's private module. import mori.ops.gemm_ar itself does not load _compat: the package imports only _shuffle and layout, and the kernels load lazily. The three modules that import _compat (kernels_lsa, kernels_sdma, kernels_fused) already import mori.cco.device.flydsl themselves. So this adds no new import-time dependency for any caller.

[kawhil-amd]

Nit: this comment is wrapped badly (llvm.inttoptr is identical in sits alone on a short line). __all__/the re-export list also omits CM_NT and CM_SC0, which buffer.py exports.

[jhchouuu]

Fixed: the comment, which now lives in signals.py, is rewrapped, and CM_SC0/CM_NT are in __all__.

[kawhil-amd]

Minor: ARG FLYDSL_VERSION is declared before the apt and ROCm layers but is used only by the last RUN. Some builders include declared build args in the cache key of later RUN steps. The nightly matrix builds with different FLYDSL_VERSION values, so each variant may rebuild the ROCm layers. Declaring the ARG just before the flydsl RUN avoids this.

[jhchouuu]

Done: WITH_FLYDSL and FLYDSL_VERSION are now declared just before the flydsl RUN, with a comment on why.