feat(flydsl): DSV4 FP4 MoE routing + blockscale kernel port (#37 W4.5) - #61
Merged
Merged
Conversation
…OCm#2702) Three independent bugs preventing release builds: 1. Runner labels don't exist: aiter-mi300-1gpu, aiter-mi325-1gpu were never registered. Use aiter-1gpu-runner (MI325 1-GPU, confirmed in aiter-test.yaml). Fix linux-aiter-mi355-1 typo to linux-aiter-mi35x-1. 2. Docker username typo: rocmshard -> rocmshared (missing 'e'), causing Docker login failure on build-only-aiter runner. 3. setuptools_scm 10.x breaks build: moved core to vcs_versioning package, causing ModuleNotFoundError. Pin to <10 until pyproject.toml is updated. Also protect tag-based builds from cancel-in-progress.
* [OPUS] Optimize opus.hpp compile time: 70% reduction (4829ms → 1465ms) Major optimizations to reduce template instantiation overhead in opus.hpp: - Replace recursive static_ford with flat static_for + compile-time index decomposition - Use runtime flat_to_coords in layout_to_offsets to avoid N unique coord_to_linear instantiations per layout - Replace static_for loops in load/store/MMA methods with runtime for-loops where compile-time indices aren't required - Use __builtin_convertvector for large vector casts instead of 64-element pack expansion - Use __builtin_shufflevector for vector slice instead of element-by-element make_vector - Use runtime loop for contiguous set_slice instead of fold expression - Eliminate std::common_type overhead in cast_impl by specifying D type directly - Flat unfold_x_stride_at bypasses concat_tuple with direct per-element stride computation - pickup_shape_impl uses filtered index seq instead of conditional concat_tuple - Direct flatten_tuple_impl for 1-level nested tuples without explode_tuple - Cache MMA tile sizes (mma_a/b/c_len, tile_a/b/c_len) as class-level constexpr - Cache layout issue-space computations in layout_load_traits and layout_imm_offsets - Add 5-arg concat_tuple overload and reduce_tuple_mul fold-expression fast path All changes preserve codegen quality: VGPRs=251, Spills=0, Occupancy=2. Verified with GQA flash attention kernel benchmark (~930 TFlops at N=1024). * [OPUS] Optimize opus.hpp compile time: 70% reduction (4829ms → 1465ms) Major optimizations to reduce template instantiation overhead in opus.hpp: - Replace recursive static_ford with flat static_for + compile-time index decomposition - Use runtime flat_to_coords in layout_to_offsets to avoid N unique coord_to_linear instantiations per layout - Replace static_for loops in load/store/MMA methods with runtime for-loops where compile-time indices aren't required - Use __builtin_convertvector for large vector casts instead of 64-element pack expansion - Use __builtin_shufflevector for vector slice instead of element-by-element make_vector - Use runtime loop for contiguous set_slice instead of fold expression - Eliminate std::common_type overhead in cast_impl by specifying D type directly - Flat unfold_x_stride_at bypasses concat_tuple with direct per-element stride computation - pickup_shape_impl uses filtered index seq instead of conditional concat_tuple - Direct flatten_tuple_impl for 1-level nested tuples without explode_tuple - Cache MMA tile sizes (mma_a/b/c_len, tile_a/b/c_len) as class-level constexpr - Cache layout issue-space computations in layout_load_traits and layout_imm_offsets - Add 5-arg concat_tuple overload and reduce_tuple_mul fold-expression fast path - Fix flatten_tuple SFINAE for non-tuple types (e.g. seq) All changes preserve codegen quality: VGPRs=251, Spills=0, Occupancy=2. Verified with GQA flash attention kernel benchmark (~930 TFlops at N=1024, ~1263 TFlops at N=16384). All OPUS device kernel tests pass (MFMA, cast, load/store, tr_load, predicated ops). * [OPUS] Update op_tests/opus README with post-optimization compile times Update per-file and total device test compile times to reflect opus.hpp compile-time optimizations (MFMA tests 1.3-2.1x faster, total parallel build 930ms → 625ms). Fix source count (18 .cu files). * [OPUS] Condense opus.hpp: remove redundant comments and merge short function bodies * [OPUS] Fix _if methods to pass multi-index to predicate (not flat index) The load_if/store_if/tr_load_if/async_load_if predicates receive multi-index (ids...) from static_ford iteration, which callers like opus_fmm use for boundary checks via layout_cached::operator()(ids...). Revert _if methods to use static_ford for predicate dispatch while keeping layout_to_offsets for data access. * [OPUS] Add 2D multi-index predicate test for load_if/store_if Add test_predicated_copy_2d that uses a 2D layout with a multi-index predicate (i_row, i_col) for boundary checking. This catches bugs where _if methods pass flat index instead of multi-index to predicates — the exact issue found with opus_fmm. * [OPUS] Update op_tests README: add predicated_copy_2d test, fix folder listing Add predicated_copy_2d to test summary table and total count (3→4 load_store_if tests). Add missing wmma_f16/f32/f8.cu to folder structure. * [OPUS] Add SKILL.md: compile-time best practices guide Document techniques for reducing HIP/C++ kernel compile time with opus.hpp, covering: header minimization, template instantiation reduction, LLVM builtins, intermediate type avoidance, parallel compilation, and measurement with -ftime-trace. Based on 70% compile-time reduction achieved on GQA flash attention kernel (4.8s → 1.5s) and 61x improvement in warp_sort_bitonic Python binding builds (21s → 346ms). * [OPUS] Move SKILL.md to .claude/skills/ for Claude Code skill discovery Move compile-time best practices guide from csrc/include/opus/SKILL.md to .claude/skills/opus-kernel-best-practice/SKILL.md following Claude Code's skill convention. Add YAML frontmatter for auto-discovery. Invokable via /opus-kernel-best-practice in Claude Code sessions. Update csrc/include/opus/README.md with link to the skill. * [OPUS] Fix Black formatting in test_opus_device.py * [OPUS] Merge device-side defs into hip_host_minimal.h, update SKILL.md Merge __launch_bounds__, __shared__/__device__/__global__ fallbacks, and __all() warp vote from opus_attn/hip_minimal.h into aiter's hip_host_minimal.h. The header now serves both host and device passes. Update SKILL.md: add Section 0 "Always Separate Device and Host Code" as the most important technique (hipcc 2-pass compilation). Show proper include paths (-I<aiter_root>/csrc/include) and the canonical __HIP_DEVICE_COMPILE__ guard pattern with opus/opus.hpp + hip_host_minimal.h. * [OPUS] Rename hip_host_minimal.h → hip_minimal.h The header now covers both host and device passes with proper guards: - Device pass (#if __HIP_DEVICE_COMPILE__): __all(), warp vote - Host pass (#if !__HIP_DEVICE_COMPILE__): dim3, hipMalloc, hipLaunchKernelGGL - Both passes: __launch_bounds__, __shared__/__device__/__global__/__host__ Old hip_host_minimal.h kept as a compatibility shim (#include "hip_minimal.h"). All .cu test files, README.md, and SKILL.md references updated. * [OPUS] Add hipMemcpy, hipEvent* to hip_minimal.h * [OPUS] Move hip_minimal to opus/hip_minimal.hpp Canonical include is now: #include "opus/hip_minimal.hpp" Old locations (hip_minimal.h, hip_host_minimal.h) kept as shims. All .cu test files, SKILL.md, and READMEs updated. * [OPUS] Guard hip_minimal.hpp types to avoid conflict with hip_runtime_api.h Wrap hipError_t, hipStream_t, dim3 in #if !defined(HIP_INCLUDE_HIP_HIP_RUNTIME_API_H) so hip_minimal.hpp can coexist with <hip/hip_runtime.h> when both are included. Guard __ockl_wfall_i32/__all with HIP header guards to prevent redefinition. * [OPUS] Add device intrinsic wrappers to opus.hpp, simplify hip_minimal.hpp opus.hpp now provides device intrinsic wrappers so kernels only need #include <opus/opus.hpp> without <hip/hip_runtime.h>: opus::thread_id_x/y/z(), block_id_x/y/z(), block_size_x/y/z(), grid_size_x/y/z(), sync_threads(), warp_all() hip_minimal.hpp is now host-only: dim3, hipMalloc, hipMemcpy, hipEventCreate/Record/Elapsed, hipLaunchKernelGGL, etc. Guarded to coexist with <hip/hip_runtime.h>. * [OPUS] Update SKILL.md and README with device intrinsic wrapper docs Document the opus:: device intrinsic wrappers (thread_id_x, block_id_x, sync_threads, warp_all, etc.) and the recommended pattern: opus.hpp for device code (self-contained), hip_minimal.hpp for host code only. Add mapping table from HIP runtime → opus:: → LLVM builtin. * [OPUS] Remove hip_minimal.h and hip_host_minimal.h compatibility shims All source files now include <opus/hip_minimal.hpp> directly. The old shim files at csrc/include/ are no longer needed.
* fix fused_dynamic_mxfp4_quant_moe_sort dispatch * Optimize fused_dynamic_mxfp4_quant_moe_sort_hip in small M * update
* Update runner-config.yml * CI: surface runner-config mapping in AMD CI job monitor Load GPU architecture and count from runner-config.yml so the runner fleet summary shows the configured inventory for each label. Trigger the monitor workflow when runner mappings change and install PyYAML for the runner report job.
Co-authored-by: gyohuangxin <42127654+gyohuangxin@users.noreply.github.com> Co-authored-by: Xin Huang <Xin.Huang@amd.com>
Add a dedicated GitHub Actions workflow for op_tests/opus so OPUS validation runs independently on MI35X and MI325 runners without being mixed into the main Aiter test shards.
Co-authored-by: solin <bingzhou@amd.com>
Split main-branch concurrency by event type so scheduled runs do not block push-triggered validation when a long queued job keeps the nightly workflow open.
* CI: Enable Deepseek ATOM tests on MI35X * CI: Use /models cache for MI35X ATOM DeepSeek test Route the MI35X DeepSeek job to the runner-local /models cache so it avoids downloading into /run, and make the output artifact name unique now that two DeepSeek variants run in the same workflow. * CI: Mount /models into MI35X ATOM test container Pass the runner's shared /models cache into atom_aiter_test so MI35X DeepSeek jobs can use the mounted model path.
…2717) * replace ck with opus * fix compile issue * fix waterfall and use buffer inst for lse. * Replace ck with opus for mla metadata. * Add dim=512 fp32 case for wave32
* docs: add ISA-level kernel optimization guide using LLVM tools Step-by-step guide covering the full LLVM-based workflow for inspecting, modifying, and recompiling AITER GPU kernel ISA: disassemble, extract reassemblable .s, round-trip recompile with binary-identical .text verification, and profile with rocprofv3. Includes Python extraction script handling branch label word-offset addressing, llvm-objcopy section swap for preserving kernel metadata, and rocprofv3 kernel-trace + ATT profiling instructions. * docs: add ISA optimization code examples and Dockerfile Runnable companion to the ISA kernel optimization guide: - extract_asm.py: standalone ASM extraction with CLI interface - analyze_kernel.py: instruction mix analysis and rocprofv3 profile parser - roundtrip.sh: end-to-end disassemble/extract/recompile/verify script - Dockerfile: ROCm 7.2.1 dev environment with all tools pre-installed including ATT trace decoder built from source * style: fix black formatting and ruff lint in ISA optimization examples - Rename loop var 'l' to 'ln' to fix E741 (ambiguous variable name) - Remove extraneous f-prefix on strings without placeholders (F541) - Apply black auto-formatting * style: fix black formatting for CI compatibility - Add blank line between module docstring and imports (E302) - Collapse multiline f-string call arguments --------- Co-authored-by: Peng Sun <pensun@Pengs-MacBook-Pro.local>
* Add run_config/compare support to GemmTuner (bf16) - Add config_env_name, _clear_op_caches(), and run_config() to GemmTuner so --run_config and --compare flags work for bf16 GEMM tuning - Update bubbly-exploring-turtle.md plan doc to reflect the full implementation including --compare, config_env_name, cache clearing, and post-tune config switching architecture Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add --run_config and --compare benchmark support to all tuners Add infrastructure in base_tuner.py for production operator benchmarking: - --run_config: benchmark only, no tuning - --compare: pre-tune benchmark, tune, post-tune benchmark with comparison table - Config env switching and cache clearing for post-tune benchmarks Implement run_config() and _clear_op_caches() in all CK-based tuners: gemm_a8w8, gemm_a8w8_bpreshuffle, gemm_a8w8_blockscale, gemm_a8w8_blockscale_bpreshuffle, gemm_a4w4_blockscale, gemm_moe_2stages, batched_gemm_a8w8, batched_gemm_bf16 * Revert unintended composable_kernel submodule change * Fix review comments and remove intermediate plan docs - Save/restore AITER_REBUILD original value instead of hardcoding 0 - Use defensive strip() for mixed-type object columns in _read_csv - Remove docs/bubbly-exploring-turtle.md and docs/run_config_benchmark.md (consolidated into csrc/.claude/add_run_config_to_tuner.md) * update ref rtol,atol * Fix tuner cache invalidation, run_config preshuffle, and compare workflow - Fix _clear_op_caches for all tuners: properly clear lru_cache and internal dict/attribute caches (a4w4, a8w8 variants, fmoe) so post-tune benchmarks use freshly tuned configs instead of stale ones. - Fix fmoe run_config: preshuffle weights before calling fused_moe to match production layout (tuner always tunes with bpreshuffle=True), preventing preshuffle_on/off module mismatch and 99%+ error. - Add defensive warning in fused_moe get_2stage_cfgs when tuned config is found but is_shuffled=False. - Fix run_config to read shapes from tuned CSV and set config env var. - Fix --compare workflow: run post-tune benchmark before tune_summary to avoid summary errors blocking verification. - Fix base_tuner _set_config_env_for_run_config return value. - Use print instead of logger.info for benchmark tables. * fix format * fix format * update readme * fix lint error * fix lint * update csv only when perf improves * format * fix lint * revert format for some files * clarify compare and gated update flow Make --compare keep a candidate csv and require --update_improved before writing back tuned results so the CLI stays explicit and easier to extend. Made-with: Cursor * fix flydsl GemmTuner review issues Trigger FlyDSL package validation before importing tuning kernels and align the FlyDSL bias cast order with the runtime path to avoid unintended dtype promotion. Made-with: Cursor * update * revert claude md * update shape_grouped * fix format * fix bug * fix lint error --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
* Move swiglu to a util file + add optional residual flag * refactor reduce and make it compatible with >65k tokens * Update aiter/ops/triton/_triton_kernels/moe/reduce.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: Lukasz Burzawa <lukasz.burzawa@amd.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…ROCm#2744) When `use_asm_v3` is `false`, `fmha_fwd_v3()` correctly returns `-1` to fall back to the CK path, but it also emits a misleading "unsupported condition in fwd_v3!!!" warning. This is not an unsupported condition — the caller intentionally opted out of v3. Separate the `use_asm_v3` check into an early return without a warning, so the `AITER_LOG_WARNING` only fires for genuinely unsupported parameter combinations (wrong head dims, unsupported dtypes, bias, dropout, wrong arch). Made-with: Cursor
* feat(aot): add MoE FlyDSL AOT pre-compilation module Co-Authored-By: Claude Opus 4 <noreply@anthropic.com> Signed-off-by: zhiding512 <zhimding@amd.com> * refactor(aot): remove --stage flag from MoE AOT module Co-Authored-By: Claude Opus 4 <noreply@anthropic.com> Signed-off-by: zhiding512 <zhimding@amd.com> * reformat * feat(aot): integrate FlyDSL MoE AOT precompilation into setup.py Move moe.py into aiter/aot/flydsl/, support multiple CSV configs, simplify compile_one_config to use COMPILE_ONLY=1 env var, and add MoE AOT pre-compilation step to the package build in setup.py. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com> Signed-off-by: zhimding <zhiming.ding@amd.com> * feat(aot): FlyDSL MoE AOT with COMPILE_ONLY dummy-tensor precompilation Rework FlyDSL MoE AOT to use COMPILE_ONLY=1 with dummy tensors instead of run_kernel, removing all HIP op dependencies (moe_sorting, shuffle_weight, etc.) from the precompilation path. Changes: - Replace _run_kernel with _precompile_to_cache using torch.zeros dummy tensors and COMPILE_ONLY=1 for pkl cache generation - Add sys.modules bridging in setup.py so aiter.jit.core reuses the same module instance loaded via sys.path - Auto-detect bundled flydsl_cache in aiter/__init__.py and set FLYDSL_RUNTIME_CACHE_DIR - Add KeyError to aiter/__init__.py exception handler for robustness - Support multiple CSV configs (dsv3 + kimik2) - Remove run_kernel parameter and test_bad_tile logic Signed-off-by: zhimding <zhiming.ding@amd.com> * update flydsl * update flydsl * adapt hgemm * fix black * add flydsl gemm aot precompile support --------- Signed-off-by: zhiding512 <zhimding@amd.com> Signed-off-by: zhimding <zhiming.ding@amd.com> Co-authored-by: Claude Opus 4 <noreply@anthropic.com>
* gather support qk_nope_head_dim != v_head_dim * fix 192 pad
…OCm#2733) * feat: add/retune BF16 GEMM configs with FlyDSL backend for 6 models Tuned on MI355X (gfx950) with all backends competing (ASM, hipBLASLt, Triton, FlyDSL). New tuned configs for Llama 70B, Llama 405B, Qwen 32B. Re-tuned existing configs for GPT-OSS, DSV3, Kimi-K2 to include FlyDSL. Backend wins across 708 total shapes: - hipBLASLt: 472 (66.7%) - ASM: 131 (18.5%) - FlyDSL: 70 (9.9%) - Triton: 7 (1.0%) - Mixed/other: 28 (4.0%) * feat: retune BF16 GEMM without hipBLASLt, add GLM-5 and 3 new models Re-tuned all BF16 GEMM configs on MI355X (gfx950) with --libtype asm,triton,flydsl (no hipBLASLt). Added GLM-5 (88 shapes from CI log) and new configs for Llama 70B, Llama 405B, Qwen 32B. Backend wins across 796 total shapes (7 models): - ASM: 437 (54.9%) - FlyDSL: 224 (28.1%) - Triton: 135 (17.0%) Per-model breakdown: - GPT-OSS (57): asm=54, triton=3 (bias=True, no FlyDSL support) - DSV3 (58): flydsl=22, triton=18, asm=18 - Kimi-K2 (125): asm=77, flydsl=46, triton=2 - GLM-5 (88): asm=42, flydsl=30, triton=16 - Llama 70B (156): asm=84, flydsl=49, triton=23 - Llama 405B (156): asm=89, flydsl=43, triton=24 - Qwen 32B (156): asm=73, triton=49, flydsl=34 Tuning time without hipBLASLt: 4h total (long pole: 405B @ 4h) vs with hipBLASLt: 10h+ total (long pole: 405B @ 8h+)
* Hoist inspect.signature/typing.get_type_hints out of per-call ctypes dispatch These two introspection calls were recomputed on every invocation of the ctypes caller closure (~79µs + ~91µs per call). Since the decorated function's signature and type hints are immutable, compute them once at decoration time and capture via closure. Made-with: Cursor * update ruff format * update black format --------- Co-authored-by: amd-ruitang3 <rui.tang2@amd.com>
…ise (ROCm#2732) * Handle FlyDSL LDS limits and candidate failures Use shared-memory-per-block queries to keep FlyDSL LDS checks architecture-aware, and surface candidate failures as concise runtime warnings so tuning can continue without noisy tracebacks. Made-with: Cursor * Keep tuner topk local per shape Avoid mutating the shared topk value while post-processing one shape so later shape groups keep the intended candidate limit. Made-with: Cursor * fix lint * Cache FlyDSL shared memory queries Avoid repeated device property lookups while validating FlyDSL kernel configs by caching the default device selection and shared-memory-per-block queries. Made-with: Cursor * fix lint * Update aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * parse kernel name to select flydsl kernel * fix black format error * refine --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: solin <bingzhou@amd.com>
Co-authored-by: Sergey Solo <ssolovye@amd.com>
…m#2221) * Make AiterAsmKernel load hsaco on each GPU it is used on * Replace unsafe uses of std::unordered_map with SynchronizedCache
…OCm#2723) * Fix Triton MoE GEMM shared memory exhaustion - Reduce num_stages in kernel configs - Lowered LDS usage to avoid shared memory OOR - Fix triton.runtime.errors.OutOfResources errors in MoE GEMM kernels * Fix: set num_stages=1 on gfx950 using get_arch() conditionally for gfx950 to ensure no bottlenecks for gfx942 * Add determinism for fused mul add test * Format fused mul add test with black
The type annotation bool was incorrect for moe_sorting_dispatch_policy, which accepts int values. The @torch_compile_guard decorator uses these annotations to generate PyTorch custom op schemas; with bool, PyTorch schema enforcement casts any value to bool, so dispatch_policy=2 becomes bool(2)=True (1), silently losing the intended policy. Using int allows callers to set dispatch_policy=2 correctly. Fixes: ROCm#2576 Signed-off-by: Tres Popp <tres.popp@amd.com> Co-authored-by: Tres Popp <tres.popp@amd.com>
* fix moe splitk aot and jit * split moe aot to serveral libs base on tuned_moe configs * update copyrights * fix typo * test shuffle as default and fix moe split jit
* add optimized prefill gdn kernels for qwen3_5 * refine code style * add ssm_state vk_layout kernels for vllm support * add default turnoff for triton autotune
…ROCm#2262) * Introduce asm fmoe kernels that do not require bf16->fp8 quantization * Update quantization division precision to be closer to IEEE correctness * Add transpose_scale flag * Update kernels with fixed s_waitcnt * Add 16x128 kernel merged with quantization * Remove legacy code * Remove one more redundant line * Revert changes to sub_X_cnt calulations for fmoe asm * Update the tuner to support x_bf16 kernel * Update pandas interface * Fix formatting * Generate config with the new tuner * Update base tuner to support merging csv files without strict column matching * Fix formatting * Fix core.py to handle columns merge * Fix formatting --------- Co-authored-by: Sergey Solo <ssolovye@amd.com>
ROCm#2893) * test: expand test_batch_prefill_large_kvcache for >4GB KV cache overflow Rewrite test_batch_prefill_large_kvcache to validate the per-tile SRD rebase fix for >4GB KV caches across all page sizes, dtypes, and attention configurations: - Add page_size=1 and 16 (page_size < kN0, exercises rebase path) - Add GQA (16, 8) in addition to MHA (8, 8) - Add causal masking with CK-compatible attn_mask for SDPA reference - Use full KV cache (4.5GB) with pages spanning the overflow boundary - Use torch SDPA as reference (memory-efficient backend, no score matrix materialization) - Add scatter_pages parameter (False only; True for future global_load_lds flat addressing) - Add GPU memory check to skip configs that exceed HBM capacity Test matrix: 24 cases (3 page_sizes × 2 dtypes × 2 causal × 2 GQA × 1 scatter) * test: add GPU sync after CK kernel in large_kvcache test Add torch.cuda.synchronize() after CK kernel launch in test_batch_prefill_large_kvcache to ensure all async GPU work completes before memory is freed between tests. Without this sync, repeated allocate/free cycles of large KV cache buffers (~20GB) with mixed dtype (bf16→fp8) can trigger GPU page faults when the HIP memory allocator reuses virtual addresses that are still referenced by pending async GPU work. The fault manifests as VM_L2_PROTECTION_FAULT at address 0x0 (NULL), causing GPU reset and kernel soft lockup. * feat(fmha): runtime dispatch for >4GB KV cache in batch prefill Add use_64bit_load to batch prefill traits and runtime overflow detection. When page_block_size < 128 and max_page_byte_offset > INT32_MAX, dispatch to the flat 64-bit load kernel variant for correctness. Also add vectorized KV layout coverage to test_batch_prefill_large_kvcache. * fix: remove unused k_vector_size variable in large_kvcache test * fix(mha): improve batch_prefill TORCH_CHECK error message for >4GB KV cache Include page_size, num_pages, and dtype in the error message when kernel dispatch fails. Add hint about CDNA3+ GPU requirement when KV cache exceeds 4GB with page_size < 128. * test: update scatter_pages comment in large_kvcache test The comment incorrectly stated scatter_pages=True was "expected to FAIL". This is no longer true — the flat 64-bit load path handles scattered pages correctly. Update to describe the test's purpose instead. * fix(mha): widen batch_prefill 64-bit threshold to total KV bytes The previous check used (num_total_pages - 1) * batch_stride * element_size which measures the last-page base offset, missing within-page offsets and producing an off-by-one at exactly INT32_MAX (the largest representable SRD voffset). Switch to total KV cache footprint (num_total_pages * batch_stride * element_size > INT32_MAX) so within-page reads on the last page are covered, and drop the redundant num_total_pages > 1 guard since single-page configs trivially fit in 32 bits. Also unify wording: 4GB → 2GB (INT32_MAX byte offset for SRD voffset), matching CK's TwoGB convention. The actual hardware bound has always been 2GB; the prior comments were imprecise. Found during batch prefill template dispatch review. * docs(mha): unify >2GB wording in batch_prefill error and test The 4GB number in the TORCH_CHECK error message and the test comment was imprecise — the actual SRD voffset bound is 2GB (INT32_MAX). Update both to match the threshold check and CK's TwoGB convention. Found during batch prefill template dispatch review. * refactor(mha): drop wrapper-side use_64bit_load; let CK dispatcher decide The wrapper hardcoded kN0_min = 128 to compute the >2GB KV cache predicate, which leaked CK tile config into aiter and would silently break if a new arm with bn0 != 128 were added. The CK auto-generated dispatcher now decides per-arm using its own compile-time bn0 and per-dtype kElementBytes, so the wrapper just forwards args. Remove the `use_64bit_load` runtime field from `mha_batch_prefill_traits`, the parameter from `get_mha_batch_prefill_traits()`, and the entire predicate computation block from the dispatcher call site. Bumps CK submodule to pull in the matching codegen change. * chore(mha): bump CK + update wrapper wording for kUseGlobalLoad rename Bumps 3rdparty/composable_kernel to dd8d293ea (refactor(fmha): batch prefill review polish — assert helper + setter guards) which builds on the prior 99a3ca9af kUseGlobalLoad rename. Wrapper-side updates to match: * csrc/cpp_itfs/mha_fwd_batch_prefill.cu: rename "64-bit-load" wording in the per-arm dispatcher comment to "kUseGlobalLoad" so the wrapper comment matches the CK-side identifier. Also drops the trailing `false /* skip_min_seqlen_q */` argument from the get_mha_batch_prefill_traits call to match the upstream CK API signature change. * csrc/py_itfs_ck/mha_batch_prefill_kernels.cu: change the >2GB error message from "page_size < 128" to "page_size < kN0" so the diagnostic tracks the tile-size constant rather than a magic number. * op_tests/test_batch_prefill.py (test_batch_prefill_large_kvcache): three documentation enhancements with no behavior change — - explain why qo_len caps at 128 (causal) / 1024 (non-causal): the causal cap is a math-backend cliff for the SDPA reference, not a kernel limit; - explain that the +256 padding on kv_page_indices is a batch_prefill ABI requirement (kernel may speculatively read up to bn0=256 entries past the last valid page index); - expand the torch.cuda.synchronize comment to call out the misattribution failure mode and GPU-reset cascade risk. * test(fmha): parametrize test_batch_prefill_large_kvcache over batch_size {1, 4} Adds multi-batch coverage to the >2GB KV cache regression test. The previous single-batch coverage left the kernel's per-sequence SRD rebase path unexercised: with cu_seqlens_q=[0, qo_len] and kv_indptr= [0, num_blocks], the kernel never walks the indptr to reposition K/V SRDs across batch boundaries. After the kUseGlobalLoad rename and the new positive static_assert(kUseGlobalLoad_) calls in update_physical_pages and set_page_stride_elements, we want a regression that catches any boundary-crossing SRD bug -- the failure mode no single-batch test can detect (one batch correct, others wrong). batch_size=4 partitions the >2GB page pool across 4 sequences (last sequence absorbs the remainder), exercising 3 cross-batch SRD transitions. The SDPA reference is computed per-batch and concatenated; per-iteration free + empty_cache keeps peak memory at one batch's worth. Verified on: - gfx950 (smci355-gfx950, MI355X): 160 passed, 32 skipped - gfx942 (smc300x-clt, MI308X): 160 passed, 32 skipped Skips are the existing vectorized + page_size=1 incompatibility (3D tensor layout), now 16 per batch_size value. --------- Co-authored-by: Xin Huang <Xin.Huang@amd.com>
--------- Co-authored-by: zhuyuhua-v <yuhzhu@amd.com>
…M#37) Host-side ABI validator for DSV4 sparse_attn call sites. This task covers §5 #1 (shape & rank) of the spec; subsequent tasks fill in dtype, device/contiguity, topk domain, slot domain, positions, cu_seqlens_q monotonicity & ownership. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…TOM#37) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…y513/ATOM#37) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…unway513/ATOM#37) Completes the AITER validator. Full host-side ABI checker per spec §5. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…acy) Skeleton for porting ROCm/FlyDSL kernels/moe_blockscale_2stage.py (Apache-2.0) into aiter for DSV4 per_1x128 FP8/FP8 g1u1 MoE. Task 1 of the W4.5 accuracy plan (sunway513/ATOM#58 merged): - Apache-2.0 SPDX header + dual copyright (FlyDSL + AMD) - Kernel-name pattern: flydsl_moe{stage}_afp8_wfp8_bf16_blockscale_t{M}x{N}x{K}[_{mode}][_w{WPE}][_bnt{B}][_xcd{X}] - 96 stage1 + 192 stage2 kernels registered (tile_m={32,64,128} x tile_n={128,256} x tile_k={128,256} x waves_per_eu={1..4} x b_nt={0,2}; stage2 also has mode={atomic, atomic_persist}) - DSV4 12-token prefill candidate (t32x128x256_w2) verified registered - compile_moe_blockscale_gemm1/2 + flydsl_moe_blockscale_stage1/2 are NotImplementedError stubs (Task 2 ports them) Verified preflight (Task 0): - FlyDSL upstream commit 8bee73f has moe_blockscale_2stage.py (136KB) - Installed flydsl 0.1.4.2 satisfies aiter's _MIN_FLYDSL_VERSION=0.1.3 - All FlyDSL APIs used by upstream import OK - aiter's mfma_preshuffle_pipeline.py + mfma_epilogues.py contain ALL symbols upstream moe_blockscale_2stage.py imports (despite being more evolved than upstream — keep aiter version, no backport needed) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…tream (#37 W4.5) Replaces the NotImplementedError stubs in aiter/ops/flydsl/moe_blockscale_kernels.py with the verbatim port of ROCm/FlyDSL kernels/moe_blockscale_2stage.py: - compile_moe_blockscale_gemm1 (stage1 fused gate+up SiLU GEMM) - compile_moe_blockscale_gemm2 (stage2 down-projection GEMM) - compile_moe_blockscale_gemm2_ex (atomic vs reduce mode dispatcher) - compile_moe_reduction (topk reduction kernel) - MoeGemm2Mode + _MoeGemm2ReduceWrapper - _if_then / _if_else SCF helpers The kernels' top-level imports are rewritten to point at aiter's copies of mfma_preshuffle_pipeline and mfma_epilogues under aiter/ops/flydsl/kernels/. Adds Python wrappers flydsl_moe_blockscale_stage1 / _stage2 that mirror the API of the FP4 versions in aiter/ops/flydsl/moe_kernels.py so the fused_moe dispatcher can route DSV4 (per_1x128 FP8/FP8) through FlyDSL instead of CK MoE (which has the ABI mismatch tracked in #37). The reduce-mode wrapper is intentionally guarded: upstream compile_moe_blockscale_gemm2_ex(REDUCE) currently passes an unsupported in_dtype= kwarg to compile_moe_blockscale_gemm2; until that is reconciled in FlyDSL, the wrapper raises NotImplementedError for mode='reduce' and only the atomic path is exercised. Smoke test: docker exec atom_dsv4_feat /opt/venv/bin/python -c " import sys; sys.path.insert(0, '/workspace/aiter-lingpeng') from aiter.ops.flydsl.moe_blockscale_kernels import ( flydsl_moe_blockscale_stage1, flydsl_moe_blockscale_stage2, compile_moe_blockscale_gemm1, compile_moe_blockscale_gemm2, compile_moe_blockscale_gemm2_ex, compile_moe_reduction, MoeGemm2Mode, ) print('imports OK')" Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Task 3 of the W4.5 accuracy plan — make the new blockscale wrappers importable as `aiter.ops.flydsl.flydsl_moe_blockscale_stage1/2` so the fused_moe dispatcher can route DSV4's per_1x128 quant scheme to them. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Task 4 — extends `aiter/fused_moe.py:get_2stage_cfgs` to route blockscale kernels (DSV4's per_1x128 FP8/FP8 path) through new wrappers. Selection rule: substring `_blockscale_` in kernelName1/2 (NOT prefix match — actual names like `flydsl_moe1_afp8_wfp8_bf16_blockscale_t...` share the `flydsl_moe1_` prefix with the FP4 path). Adds two new wrappers beside the FP4 ones: - `_flydsl_blockscale_stage1_wrapper` → calls flydsl_moe_blockscale_stage1 - `_flydsl_blockscale_stage2_wrapper` → calls flydsl_moe_blockscale_stage2 Both forward kernel params (tile_m/n/k, waves_per_eu, b_nt) parsed from the kernel name registry; pass scale_block_k=128 explicitly (per_1x128 quant). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…#37 W4.5) Adds op_tests/test_flydsl_moe_blockscale.py covering the AITER FlyDSL blockscale MoE port at the small/medium upstream sweep shapes (E=8, inter=256) and the DSV4 prefill 12-token gate shape (B=12, model=7168, inter=3072, E=384, topk=6). Sweep results: - test_moe_blockscale_stage1_sweep[small-E8]: PASS - test_moe_blockscale_stage1_sweep[medium-E8]: PASS - test_moe_blockscale_stage2_sweep[small-E8]: PASS - test_moe_blockscale_stage2_sweep[medium-E8]: PASS - test_moe_blockscale_dsv4_shape[12-7168-3072-384-6]: XFAIL The DSV4 case is marked xfail because the FlyDSL preshuffle B-tensor stride descriptor (mfma_preshuffle_pipeline.make_preshuffle_b_layout) casts row strides to i32 before fx.make_layout, which overflows once E*2*inter*K exceeds 2**32 elements. Empirically: E= 96 ~4.23 GB elems -> 100.00% close (under 2**32) E=100 ~4.40 GB elems -> 93.29% close E=128 ~5.64 GB elems -> 78.04% close E=256 ~11.27 GB elems -> 43.45% close E=384 ~16.91 GB elems -> 25.56% close (DSV4) Resolving the DSV4 case requires a kernel-side i64 stride path through the preshuffle layout helpers, plus per-expert buffer descriptors to get around the AMD u32 NUM_RECORDS field (4 GB per descriptor). Tracked in sunway513/ATOM#37 W4.5 follow-up. The test file is structured so removing the xfail decorator after the kernel fix lands will validate it immediately -- DSV4 shares its body with the passing sweep cases. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The original DSV4 test used the no-TP shape (E=384, inter=3072, ~16.9 GB FP8) which triggers an i32 stride overflow in FlyDSL's `make_preshuffle_b_layout` (subagent's empirical bisect: clean cutoff at E*2*inter*K ≈ 2**32). That's a real kernel limitation but NOT the production case — ATOM silicon runs DSV4 with TP=8 + column-parallel sharding on w1's inter_dim, so per-rank w1 shape is [E=384, 2*768=1536, 7168] = 2.1 GB FP8 (4.22 G elements, just under 2**32). This commit: - Renames the original test to `test_moe_blockscale_dsv4_no_tp_shape` and keeps it as xfail — useful regression monitor for when FlyDSL upstream gains i64-aware preshuffle helpers. - Adds `test_moe_blockscale_dsv4_shape` using `inter_dim_per_rank=384` (= 3072/8) — the actual kernel input on TP=8 silicon. PASSES. Result: 5 passed, 1 xfailed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
#37 W4.5) Task 6 — adds aiter/configs/model_configs/dsv4_fp8_blockscale_tuned_fmoe.csv with 12 rows for the DSV4 MoE per-rank shape (TP=8 column-parallel sharded: model_dim=7168, inter_dim_per_rank=384, E=384, topk=6, FP8/FP8, per_1x128) covering token counts {1, 2, 4, 8, 12, 16, 32, 64, 128, 256, 512, 1024}. Tile choices follow user memory project_aiterforge_fp8.md (tile_m=32 for ≤128 tokens with waves_per_eu=2; tile_m=64 for >128 tokens). Kernel names use the new blockscale variant (flydsl_moe{1,2}_afp8_wfp8_bf16_blockscale_t...). All 6 unique kernel names verified to resolve through `get_flydsl_blockscale_kernel_params` (Task 6.5 will exercise them numerically per token count). To use this config at runtime: AITER_CONFIG_FMOE=/path/to/dsv4_fp8_blockscale_tuned_fmoe.csv Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…4.5) Tasks 6+6.5 — adds CSV-coverage tests and corrects the DSV4 tuned config based on the empirical numerical findings. Test additions (op_tests/test_flydsl_moe_blockscale.py): - test_dsv4_csv_kernel_names_resolve: every kernelName1/2 row in the DSV4 CSV must contain `_blockscale_` substring AND resolve through `get_flydsl_blockscale_kernel_params`. Catches dispatcher fall-through. - test_dsv4_csv_pair_correctness[1, 12, 128, 512, 1024]: for each parametrized token count, run the CSV row's (kernelName1, kernelName2) pair end-to-end and assert numerical correctness vs torch reference. Failure here BLOCKS silicon validation (Task 7+). CSV correction (aiter/configs/model_configs/dsv4_fp8_blockscale_tuned_fmoe.csv): - Discovered: stage2 with tile_k=256 produces ~3% close vs torch ref on DSV4 inter_dim_per_rank=384 shape (regardless of M). tile_k=128 passes cleanly. Root cause: inter_dim_per_rank=384 = 3 * scale_block_k=128, so tile_k MUST equal scale_block_k for stage2 to align reduction with the per-1x128 scale stride. - Stage1 is fine with tile_k=256 (no reduction over inter_dim — output per-block). - Updated all CSV rows: stage1 keeps tile_k=128/256 by token count (perf), stage2 hard-coded to tile_k=128 (correctness). Test results: 11 passed, 1 xfailed (the no-TP DSV4 shape, separate issue). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Silicon trace with AITER_FMOE_DEBUG_LOOKUP=1 revealed DSV4 dispatches
MoE through aiter with FP4/FP4 per_1x32 (NOT FP8/per_1x128 as
config.json's weight_block_size suggested). ATOM's quant_v4 layer
rewrites the dispatch dtype before reaching aiter.
Result before: 24 LOOKUP MISS, all-punctuation gibberish output.
Result after: 24 LOOKUP HIT, FlyDSL FP4 kernels firing on stage1,
CK FP4 kernels on stage2.
Single-mode (W3 baseline, USE_W4_PATH=0) silicon now produces real
Chinese tokens vs the previous "〖,〖" gibberish.
CSV adapted from kimik2_fp4_tuned_fmoe.csv (same dims 7168/512/385,
same FlyDSL FP4 kernel set) with topk 9→6 for DSV4. Stage1 uses
flydsl_moe1_afp4_wfp4_bf16_*; stage2 uses moe_ck2stages_gemm2_*FP4X2*B16
which are already pre-built in the bind-mounted aiter.
The blockscale FP8 kernel port (commits 8fedbc1..4ca7936) remains
in-tree as future-proofing if DSV4 ever switches to per_1x128.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
ruff
Local variable vec2_i32 is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 1962 in e450e4d
Local variable n_idx_s2 is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2190 in e450e4d
Local variable expert_off is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2683 in e450e4d
Local variable s2 is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2726 in e450e4d
Local variable gpu_arch is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2967 in e450e4d
Local variable DYN is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2968 in e450e4d
Local variable USE_NONTEMPORAL is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2973 in e450e4d
Local variable VEC_ALIGN is assigned to but never used
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2974 in e450e4d
Do not assign a lambda expression, use a def
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2988 in e450e4d
Do not assign a lambda expression, use a def
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2989 in e450e4d
Do not assign a lambda expression, use a def
aiter/aiter/ops/flydsl/moe_blockscale_kernels.py
Line 2990 in e450e4d
Module level import not at top of file
aiter/op_tests/test_flydsl_moe_blockscale.py
Line 613 in e450e4d
Comment on lines
+254
to
+258
| out_mlir = lambda: ( | ||
| (lambda ty: ty() if callable(ty) else ty)( | ||
| T.f16 if out_dtype == "f16" else T.bf16 | ||
| ) | ||
| ) |
There was a problem hiding this comment.
Do not assign a lambda expression, use a def
Suggested change
| out_mlir = lambda: ( | |
| (lambda ty: ty() if callable(ty) else ty)( | |
| T.f16 if out_dtype == "f16" else T.bf16 | |
| ) | |
| ) | |
| def out_mlir(): | |
| return ( | |
| (lambda ty: ty() if callable(ty) else ty)( | |
| T.f16 if out_dtype == "f16" else T.bf16 | |
| ) | |
| ) |
| sb_per_tile_s1 = tile_k // scale_block_k # scale blocks per tile (in K dim) | ||
| ku_per_sb_s1 = scale_block_k // 64 # K64-steps per scale block = 2 | ||
| nblk_k_w1 = model_dim // scale_block_k # K-blocks in W1 (=scale_k) | ||
| nblk_n_w1 = (2 * inter_dim) // 128 # N-blocks in W1 (ScaleBlockN=128) |
Comment on lines
+299
to
+300
| size_out = DYN | ||
| size_x = DYN |
Comment on lines
+300
to
+301
| size_x = DYN | ||
| # W is packed int4 for W4A8: 2 values per byte. |
| size_out = DYN | ||
| size_x = DYN | ||
| # W is packed int4 for W4A8: 2 values per byte. | ||
| size_w = ( |
| topk_idx = fx.Index(topk) | ||
| m_in = tokens_in * topk_idx | ||
| m_i32_v = arith.index_cast(T.i32, m_in) | ||
| layout_x = fx.make_layout((m_i32_v, k_i32_v), stride=(k_i32_v, 1)) |
| elem_bytes=elem_bytes, | ||
| ) | ||
| layout_b = b_layout.layout_b | ||
| c_k0 = (k_in * arith.index(int(elem_bytes))) // fx.Index(64) |
| k_blocks16 = arith.index(tile_k_bytes // 16) | ||
| layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) | ||
| layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) | ||
| layout_lin_rowcol = fx.make_layout((tile_m, tile_k), stride=(tile_k, 1)) |
|
|
||
| c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4) | ||
| c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4) | ||
| layout_x_div4 = fx.make_layout( |
| chunk_i32=chunk_i32, | ||
| ) | ||
|
|
||
| vec1_i32 = T.vec(1, T.i32) |
10 tasks
Owner
Author
Final closure summary (#37 W4.5 MoE half — shipped)Status: ready to merge. The MoE-routing half of #37 is closed by this PR. What landed
Silicon evidence (MI355X 8x, TP=8)
MoE numerics are sound. The same FlyDSL FP4 kernels that produce coherent text under W3 are firing under W4 — the gibberish in W4 was bisected to ATOM-side gsm8k baseline (limit=20 num_concurrent=1, W3+FP4)flexible-extract: 0.30 ± 0.105 | strict-match: 0.30 ± 0.105 Proves the FP4 routing fix doesn't regress the W3 baseline (previous behavior was effectively 0% — gibberish or crash). What this PR does NOT close
The blockscale FP8 kernel port (most of the LOC in this PR) is future-proofing. DSV4 silicon dispatches FP4/per_1x32, so the blockscale path isn't on the silicon hot path today, but it's tested and ready if any DSV-line model later switches to per_1x128. 🤖 Generated with Claude Code |
Owner
Author
Final closure (with W4 silicon update)This PR is ready to merge — the MoE-routing half of #37 is closed. ATOM-side W4-path bisection landed three follow-up fixes; published gsm8k baseline now available. MoE evidence (unchanged from prior comment)aiter LOOKUP: 24 HIT / 0 MISS under W3 + FP4 CSV. Same kernels fire under W4 + FP4 — MoE numerics are independently verified to be sound. Final gsm8k baseline (W3 + FP4, limit=20, num_concurrent=1)
Proves the FP4 routing fix doesn't regress baseline. Previous run was effectively 0 (gibberish/crash). W4-mode (USE_W4_PATH=1) status — out of scope here, fully tracked in ATOM PR #59W4-path fixes shipped on the ATOM side ( 🤖 Generated with Claude Code |
Owner
Author
Sprint 4 final state — MoE half closed; ATOM W4 path also closed; reference gap remainsaiter side (this PR) — ready to merge. ATOM PR sunway513/ATOM#59 has now landed all three W4-path bug fixes (commits
Final accuracy data (limit=20 num_concurrent=1):
Real correctness gap remains between ATOM and SGLang reference. CIs do not overlap (SGLang lower 0.92 > ATOM upper 0.91 even at n=100). Hypotheses being investigated:
These investigations are upstream of the MoE / FlyDSL work in this aiter PR. MoE numerics are independently verified (24 LOOKUP HIT, kernel outputs match across W3 path). The accuracy gap is not in this PR's scope. 🤖 Generated with Claude Code |
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
20 participants
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.
Summary
Closes the MoE routing half of sunway513/ATOM#37 (DSV4 W4.5 accuracy regression).
aiter/configs/model_configs/dsv4_fp4_tuned_fmoe.csvregisters 16 tuned rows so DSV4's MoE lookup(7168, 512, 385, topk=6, FP4/FP4 per_1x32)resolves to FlyDSL FP4 stage1 + CK FP4 stage2 instead of falling through to an unmatched CK MoE backend.aiter/ops/flydsl/moe_blockscale_kernels.py(Apache-2.0 from FlyDSL upstream) brings the FP8/per_1x128 blockscale 2-stage GEMMs into aiter, with dispatcher routing inaiter/fused_moe.pyand a numerical correctness test inop_tests/test_flydsl_moe_blockscale.py. Not on the silicon hot path today (DSV4 dispatches FP4), kept in-tree if DSV4 or another model later switches to per_1x128.How the FP4 path was discovered
Silicon trace with
AITER_FMOE_DEBUG_LOOKUP=1exposed the actual lookup key:ATOM's quant_v4 layer rewrites
config.json:weight_block_size=[128,128](FP8 per_1x128) into FP4 per_1x32 before reaching aiter, which the original plan v1/v2 missed. Plan revision log + diagnosis is in the ATOM-side companion PR.Silicon evidence (MI355X 8x, TP=8)
_forward_w4bugThe W4-path collapse reproduces at conc=1, isolating it from MoE. Filed as follow-up sub-issue under sunway513/ATOM#37.
Test plan
op_tests/test_flydsl_moe_blockscale.py— 5 PASS / 1 XFAIL (no-TP shape; per-rank TP=8 shape passes)dsv4_fp4_tuned_fmoe.csvresolves to a registered kernelCross-repo PR
ATOM PR: sunway513/ATOM (branch
plan/dsv4-w45-flydsl-blockscale-moe) ships plan v1→v3 revision log + Evidence M with silicon JSON artifacts.🤖 Generated with Claude Code