Repository navigation
add add_transpose implement - #4
Merged
Merged
Conversation
amd-youchen
pushed a commit
to amd-youchen/aiter
that referenced
this pull request
Dec 3, 2025
* Update trtllm_all_reduce_fusion.cu * add single load policy
poyenc
added a commit
to poyenc/aiter
that referenced
this pull request
Feb 4, 2026
- Update Issue ROCm#2 with complete root cause analysis: - Bug occurs when (seqlen_k % 16) is in range [5, 11] - K positions 4-7, 20-23, 36-39, etc. not accumulated - Verified P is correct (K[pos]=0 changes output) - Verified V is not accumulated (V[K=4]=0 produces no change) - Add Issue ROCm#3 for separate causal masking bug (seqlen≥256) - Renumber V Tile Transpose bug to Issue ROCm#4 - Add complete test results for seqlen_k 1-64 - Include pattern diagram and verification test commands Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
anhminhnguyenhoang
added a commit
that referenced
this pull request
Jun 9, 2026
… fix Refresh the optimization report and benchmark repro to the current kernel (16x16x32 atom + XCD remap + bounded-A fix): - §2 rewritten with the two-method comparison (do_bench cold-L2 + rocprofv3 hot-L2, both autotune-excluded): FlyDSL wins 3 of 4 shapes (~5-6% ahead on both D=512 cells, tie-to-slightly-behind on B1024_D256 depending on L2 warmth). - Added method lessons #3 (cold vs hot L2) and #4 (skew finding reversed on the current kernel) and reproduction subsections 4.5 (do_bench) / 4.6 (skew remap). - §6 next-opportunities: XCD remap + gate marked done; fp8 weights flagged as the only decisive lever; confirmed dead-ends listed. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
valarLip
added a commit
that referenced
this pull request
Jul 27, 2026
…igs (#4397) * [dev] Pr/a8w4 situv2 (#4) * add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [dev] Pr/fix k dimension (#5) * add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. * Fix GUI shuffle_scale k_pad for non-256-aligned MoE K and auto stage2 tile_k. Pad w2/w1 GUI e8m0 scales to k_groups multiple of 8 (DSV4 inter=640), unify stage2 tile_k selection in flydsl_moe_stage2 and fused_moe, and add staged a8w4 regression tests across inter/model K sweeps. * test: integrate a8w4 K-dimension regressions into existing test suites. Move shuffle pad coverage into test_quant_mxfp4 and FlyDSL GUI stage2/e2e into test_flydsl_moe_a8w4; drop standalone op_tests files for repo convention. * test: move flydsl a8w4 GUI regressions to op_tests/flydsl_tests. Relocate test_flydsl_moe_a8w4 under op_tests/flydsl_tests and align with pytest conventions used by other FlyDSL op tests. --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [dev] Feat/flydsl moe a16wfp4 (#6) * add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. * Fix GUI shuffle_scale k_pad for non-256-aligned MoE K and auto stage2 tile_k. Pad w2/w1 GUI e8m0 scales to k_groups multiple of 8 (DSV4 inter=640), unify stage2 tile_k selection in flydsl_moe_stage2 and fused_moe, and add staged a8w4 regression tests across inter/model K sweeps. * test: integrate a8w4 K-dimension regressions into existing test suites. Move shuffle pad coverage into test_quant_mxfp4 and FlyDSL GUI stage2/e2e into test_flydsl_moe_a8w4; drop standalone op_tests files for repo convention. * test: move flydsl a8w4 GUI regressions to op_tests/flydsl_tests. Relocate test_flydsl_moe_a8w4 under op_tests/flydsl_tests and align with pytest conventions used by other FlyDSL op tests. * Add FlyDSL MoE a16wfp4 (bf16 x mxfp4) kernels on gfx950. Port a16w4 stage1/stage2 into mixed_moe_gemm_2stage as dedicated _a16w4 builders while keeping HEAD fp8/fp4 compile paths unchanged. Wire dispatch in moe_kernels and add op_tests/flydsl_tests regression for stage1, stage2, and e2e against torch_moe references. * Fix ruff/black lint issues in a16wfp4 PR files. Remove unused _fly imports, move lds_space/types imports to module top, add noqa for test sys.path bootstrap, and apply black formatting. * Remove dead duplicate waves_per_eu None checks in a16w4 stage2. The parameter defaults to int and callers always pass an int after dispatch. * Fix a16w4 stage2: skip second compute when total_tiles==1. The CK-style ping-pong pipeline assumed total_tiles>=2. When inter_dim==tile_k (e.g. inter_dim=256, tile_k=256) total_tiles=1, the HEAD pre-fetched an out-of-bounds A tile into lds_ping, and the TAIL then computed with that garbage data, producing cos~0.71 vs reference. Fix: const_expr(total_tiles==1) guard in the TAIL block--skip the second compute_tile and the ping-buffer load entirely. Also tighten test _check_result: add cosine+rel_L2 as primary gate (cos>0.999) so future stage2 bugs cannot be masked by atol=1.0 on small-magnitude outputs. * Add pipeline guards for a16w4 stage1 total_tiles==1 and odd-tile validation. The stage1 ping-pong pipeline has the same total_tiles>=2 assumption as stage2. Add: - total_tiles==1 guard in stage1 TAIL (same pattern as stage2 fix) - Validation: reject odd total_tiles and K < tile_k at compile time for both stage1 and stage2 a16w4 paths These configs cannot occur with real LLM shapes (model_dim and inter_dim are always multiples of 256), but the guards prevent silent corruption on contrived inputs. * Add SiTUv2 activation to a16w4 (bf16 x mxfp4) stage1. Port the SiTUv2 activation from pr/a8w4-situv2 into the a16w4 stage1 kernel (both direct and cshuffle epilogues): situ_g = beta * tanh(gate/beta) * sigmoid(gate) up_scaled = linear_beta * tanh(up/linear_beta) with gate/up clamped to <=7 / [-7,7] before activation, matching the a8w4 situv2 numerics exactly. Thread situ_beta/situ_linear_beta through compile_flydsl_moe_stage1 and flydsl_moe_stage1; the per-beta cache tag keeps distinct binaries. Verified vs an inline torch SiTUv2 reference (mxfp4 dequant GEMM + clamp + situv2): cos>=0.99996 across beta in {(1,1),(0.5,2),(1.5,0.8)}, shapes (512/256, 3072/256), tile_n {128,256}. silu/swiglu unchanged. * Strip dead generic pipeline from a16w4 stage1 kernel (~2260 lines). compile_mixed_moe_gemm1_a16w4 was created by copying the generic fp8/fp4 builder and gating with is_a16w4_stage1. The live a16w4 path is fully contained in `if const_expr(is_a16w4_stage1): ... return`; everything after that return was an unreachable copy of the generic f8f6f4 kernel body (referenced generic-only vars like _lds_tid_offset_pong / _pipe_* that are never defined on the a16w4 path -- proven dead since a16w4 tests pass without them). Remove the dead generic kernel body, the generic-only setup blocks (postlude pipeline schedule, not-is_a16w4 guards), and the now-unused setup vars. Purely dead-code deletion: a16w4 stage1/stage2/e2e/situv2 + a4w4 all still pass with identical cosine. * Strip dead generic pipeline from a16w4 stage2 kernel (~1140 lines). Same shape as the stage1 cleanup: compile_mixed_moe_gemm2_a16w4's live a16w4 path is fully contained in `if const_expr(is_a16w4): ... return` inside _moe_gemm2_then_body; everything after that return was an unreachable copy of the generic f8f6f4 stage2 body (f8f6f4 MFMA, generic-only scale layouts). Proven dead: a16w4 stage2/e2e/situv2 + a4w4 all pass unchanged after removal. Remove the dead body plus the now-unused generic setup vars (pack_N/pack_K, cbsz/blgp, generic scale layouts, sx_rsrc sentinel, etc.). Purely dead-code deletion. * feat(moe): route a16w4 (bf16 x mxfp4) SiTUv2 through fused_moe Wire the mixed_moe a16w4 kernel into the fused_moe 2-stage path for SiTUv2: get_2stage_cfgs now matches bf16/fp16 x fp4 when activation is SiTUv2 (which uniquely identifies the a16w4 kernel, so GPT-OSS / legacy bf16-Swiglu keep their CK-Tile routing), maps _a_type to bf16/fp16, and infers q_dtype_a=bf16 for SiTUv2+separated so the activation stays bf16 (no fp4 quant). The stage1/ stage2 activation-quant gates also accept SiTUv2. Tests: a16w4 SiTUv2 via fused_moe (cos=1.0 vs torch ref) and an a8w4 SiTUv2 inter=640 vec4 case exercising the fix-k non-256 K-tiling. * feat(moe): support non-256 inter_dim for a16w4 (a8w4 parity) Bring the FlyDSL a16w4 (bf16 x mxfp4) MoE kernels to a8w4 parity so inter_dim need not be a multiple of 256 (e.g. DSV4 inter=640, and arbitrary values like 384 via inter_dim_pad), across stage1/stage2/E2E. Kernel (mixed_moe_gemm_2stage.py): - stage2: allow inter_dim % tile_k == 0 (drop even-only rule), add odd_k_tiles 1-tile tail path, 256-pad scale reads (#3476). - stage1: same odd_k_tiles K-loop + 256-padded w1 scale layout; fix gx launch grid (tile2_pad + full-N when inter_dim_pad > 0). Caller (moe_kernels.py): - pick/resolve_flydsl_stage1_tile_n: use tile_n=128 when inter%256!=0 (fixes stage1 error on cols 0..inter%256 under tile_n=256). - zero-init stage1 out when inter_dim_pad > 0. Tests (test_flydsl_moe_a16wfp4.py): - non-256 regressions (256/384/640) for stage1/stage2/E2E + pick_* tests. - _generate_a16wfp4_data gains activation/situ_beta/situ_linear_beta. - --perf sweep mode (correctness + latency/TFLOPs). * test(moe): fold situv2 tests into dtype files, drop standalone situv2 file Consolidate test_flydsl_moe_situv2.py into the per-dtype test files (matching how test_flydsl_moe_a16wfp4.py already embeds its situv2 tests): - test_flydsl_moe_a8w4.py: add the a8w4 SiTUv2 vec4 stage1 sweep + the host-only situv2 reference test; reuse the file's existing _check_close helper (fp32 cast for the bf16-ref vs f16-out case). Add shuffle_weight / e8m0_shuffle imports. - Delete test_flydsl_moe_situv2.py. Its broken-under-pytest a4w4/a8w4 main()-runner is dropped; a8w4 situv2 is now real parametrized coverage. a16wfp4 situv2 tests already lived in their dtype file and are unchanged. * fix conflict * add timing args --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mh <mh@local> * fix(moe): a8w4 non-256 inter_dim stage1 tile_n + situv2 E2E regression tests Squashed fix + tests (PR#3 commits 50729303, 2f33bd40, 1f0c6485, 46e3e056). flydsl_moe_stage1 only downgraded tile_n (256->128) for non-256-aligned inter_dim on a16w4 (bf16 x mxfp4); a8w4 (fp8 x mxfp4) kept tile_n=256, which in separated gate_mode over-runs the gate/up (N) axis for non-256 inter_dim: ~30% wrong E2E output or GPU memfault at inter=384/640. Extend the resolve to a8w4 (b_dtype fp4/mxfp4, a_dtype in {bf16,fp8}); a4w4 untouched. Callers keep passing tile_n=256; the kernel resolves internally. Tests (numeric vs torch reference, no-pad path, 128-multiple inter_dim): - test_flydsl_e2e_a8w4_situv2 (separated, the production/customer path) - test_flydsl_e2e_a16wfp4_situv2 (separated + interleave) - test_flydsl_e2e_a8w4_gui extended to inter=256/384/640 (interleave, swiglu) Verified on gfx950 (full rebuild): a8w4 20 / a16wfp4 27 / a4w4 4 passed; ruff + black clean. * refactor(moe): observable non-256 tile downgrade + a8w4 interleave situv2 E2E Squashed (PR#3 c827ddbb + d381aae7). - resolve_flydsl_stage1_tile_n / resolve_flydsl_stage2_tile_k: keep the silent auto-downgrade of a non-dividing tile (256->128) for non-256 inter_dim, but make it observable -- full docstrings noting tile=256 is NOT tunable for such shapes, plus a one-time (deduped) logger.warning on override. - test_flydsl_e2e_a8w4_situv2: enable interleave (a16w4-style shuffle_weight_a16w4 recipe) alongside separated, over 128-multiple inter_dim. Verified all activations (silu/swiglu/situv2) x both gate_modes x non-256 give E2E 0.0000. * Add interleave_gate_up_rows + moe_shuffle_weight to ops/shuffle.py Needed by atom-k3 (rocm/atom HEAD imports them from aiter.ops.shuffle); ported verbatim from aiter main. Keeps the single aiter-k3 branch usable by vllm-k3, sglang-k3 and atom-k3. * [kimi-K3] extend conv2d support to gfx1250 * [Kimi-K3] pin _MIN_FLYDSL_VERSION to 0.2.2 (K3 MoE kernels need loc= API removed in 0.2.4) * [Kimi-K3] fix MoE A16W4 for flydsl 0.2.4 and bump _MIN_FLYDSL_VERSION flydsl 0.2.4 removed the loc= parameter from rocdl._split_mfma_operands. The A16W4 BF16 K32 MFMA helper in mixed_moe_gemm_2stage.py called it as _split_mfma(operands, loc=loc), which raises TypeError under 0.2.4. Drop loc= from both (identical) helper sites; _split_mfma_operands only unwraps operand Values and reads int flags, so it never needed loc. The loc/ip are still threaded to the real op builder (_mfma_k32_raw), so location tracking is preserved. Bump _MIN_FLYDSL_VERSION to 0.2.4 and drop the now-stale 0.2.2 pin note. Verified on gfx950 with flydsl 0.2.4 via test_flydsl_moe_a16wfp4.py (stage1, stage2, e2e all pass, cos>=0.99999). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [Kimi-K3] add fmoe tuned config * [Kimi-K3] add bf16 dense GEMM tuned config * [Kimi-K3] extend a16w4 fmoe tuned config inter_dim=384 * fix(kimi-k3): support A4W4 SiTUv2 on gfx950 and gfx1250 Keep Kimi-K3 SiTUv2 on the A4W4 FlyDSL path, add gfx1250 grouped-MoE activation and stability support, preserve ATOM's legacy A4W4 call, and cover the gfx1250 path with focused tests. Co-authored-by: Cursor <cursoragent@cursor.com> * [Kimi-K3] fix flydsl aot build failed * Gluon reduce guard * add a8w4 fmoe tune config * Enable a8w4 SiTUv2 MoE via AITER_SITUV2_A8W4 Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> * Support row-strided inputs in grouped TopK * chore(gfx1250): drop CK enable patch — not needed for Kimi-K3 patches/ck_gfx1250_enable.patch only served to build native CK kernels (quant/cache/rmsnorm/moe/sample/custom_all_reduce) on gfx1250. Verified that with ENABLE_CK=0 (no patch, triton/hip fallbacks) Kimi-K3 serves correctly on gfx1250 x4 (tp4): full gsm8k 1319 = 0.956 flexible/strict, vs 0.9613 with the patched native-CK path — within stderr (+/-0.0056). Dropping the patch keeps cleanup/k3-minimal minimal for merge into k3-for-amd; gfx1250 runs ENABLE_CK=0 (cost: ~13% decode vs native CK, no accuracy loss). * chore(gfx1250): drop unified_attention Gluon guard — K3 doesn't use it Kimi-K3 runs full-attn via MLAAttention (MLA-latent) and KDA via fla, so it never dispatches to unified_attention (verified 0 calls across a full gsm8k run). The non-power-of-two head_size / Gluon-reduce guards were added for K3's old MHA-via-unified_attention path, which ATOM has now removed. Reverting unified_attention.py to k3-for-amd keeps cleanup/k3-minimal focused on what K3 actually needs. Full gsm8k 1319 with this + the dropped CK patch = 0.955. * Support row-strided inputs in grouped TopK opt-sort * rm test_flydsl_moe_situv2.py: SiTUv2 coverage lives in test_flydsl_moe_a8w4.py Per PR #6 review: the SiTUv2 stage1 variants (a4w4/a8w4) are already covered by test_flydsl_moe_a8w4.py (cases were migrated there), so drop the standalone file. * test: fold row-strided biased_grouped_topk case into test_moeTopkSoftmax.py Per PR review: stride is a simple feature — extend the existing test_biased_grouped_topk with a dense-vs-strided equality check (gating_output copied into a padded, non-contiguous row-strided view) instead of a dedicated test file. Removes op_tests/test_biased_grouped_topk_strided.py. Verified on gfx950: topk_ids/topk_weights [dense vs strided] pass. * test: simplify strided coverage — make gating_output itself row-strided Instead of a separate dense-vs-strided check block, create gating_output as a non-contiguous row-strided view (slice of a padded backing buffer) so the existing biased_grouped_topk_hip checks exercise the strided path directly. moe_fused_gate still gets a dense copy (it does not accept strided input). The randn stream is kept identical to the dense layout so later tests in the file see the same random inputs. Verified on gfx950: full file run, zero failures. * test: add K3 row-strided biased_grouped_topk case in main Per review discussion: keep test_biased_grouped_topk's dense default unchanged; the function now accepts an optional gating_output, and main gains one case passing the K3 fused MoE-front router layout — logits as a non-contiguous row-strided slice of the fused [gate_up|experts|routed] buffer (896 experts, topk=16). The dedicated strided test file stays removed. Verified on gfx950: full file run, zero failures (err_aiter = 0). * style: apply black formatting to gfx1250 flydsl files Match the repo's pre-checks CI (black via psf/black@stable, ruff check): reformat grouped_moe_gfx1250.py, gemm_mxscale_gfx1250.py and moe_grouped_gemm_mxscale_gfx1250.py. Ruff reports no issues on the branch. No functional change. * add gfx1250 bf16 tuned config * [Kimi-K3] fix ruff findings reported by the CI style check The pre-checks ruff job installs ruff unpinned (`pip3 install ruff`) and the repo carries no ruff config, so the enforced rule set follows ruff's defaults. ruff 0.16.0 widened those defaults well beyond E4/E7/E9/F, which is what this branch tripped over. Fix the 41 findings reviewdog flagged inside the PR diff: - silu_and_mul_fq.py (B023): bind the per-iteration SSA values as default args in _fmin / _sigmoid_s / _situv2_elem. The closures are deliberately redefined per unrolled iter_idx because the arith.constant ops must be emitted at the current insertion point, so hoisting them out of the loop is not an option. - moe_kernels.py: narrow the logger-import guard to ImportError (BLE001) and switch the flagged annotations to PEP 604 / builtin generics (UP006, UP045). - flydsl moe tests + test_moe_2stage.py: sort the import blocks (I001), drop the now-unused `# noqa: E402` markers (RUF100), rewrite dict() calls as literals (C408), and mark the sweep-runner blanket catches with an explicit `# noqa: BLE001` plus the reason they are intentional. Only findings inside the PR diff are addressed; pre-existing findings elsewhere in the same files are left alone. black --check stays clean. * [Kimi-K3] fix the second batch of ruff findings from the CI style check reviewdog only posts a bounded number of results per run, so resolving the first 41 findings surfaced a second batch of 27 that were in the PR diff all along. Same root cause (unpinned ruff + no repo ruff config, 0.16.0 defaults). - fused_moe.py / grouped_moe_gfx1250.py / moe_grouped_gemm_mxscale_gfx1250.py (UP045): PEP 604 annotations. Converted whole parameter runs rather than the individual flagged lines, because fixing one line pulls its neighbours into the diff hunk and the findings cascade. - mfma_preshuffle_pipeline.py: drop the quotes from the ir.Value/ir.Type annotations (UP037) and sort __all__ (RUF022). - mixed_moe_gemm_2stage.py: sort the import blocks (I001), functools.cache for the maxsize=None caches (UP033), inline the directly-called lambda in out_mlir (PLC3002), merge the nested bf16-atomics guard (SIM102), and underscore-prefix the unpacked values that are never read (RUF059) -- the decode calls are kept since they emit IR. All rewrites are behaviour-preserving; black --check stays clean and the import name sets are unchanged. Findings outside the PR diff are still left alone. --------- Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: billishyahao <yahao.he@gmail.com> Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: billishyahao <bill.he@amd.com> Co-authored-by: mh <mh@local> Co-authored-by: Dewei Wang <Dewei.Wang@amd.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Felix Li <felixamd@163.com> Co-authored-by: XiaobingSuper <xiaobingzhangupc@gmail.com> Co-authored-by: root <root@smci355-ccs-aus-m12-33.cs-aus.dcgpu> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: RolaoDenthu <xinyisong0111@gmail.com> Co-authored-by: zejunchen-zejun <zejun.chen@amd.com> Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
coderfeli
added a commit
that referenced
this pull request
Aug 2, 2026
… kernels dir #2 (dead CSV code): after the fold, a16w4 gemm2 tiles are already carried by the get_2stage_cfgs kernelName2 (parsed into tile_n/tile_k/b_nt/waves_per_eu/xcd_swizzle by _flydsl_stage2_wrapper), so the stage2 impl re-reading the CSV via resolve_a16w4_gemm2_config was redundant. Switch the a16w4 stage2 branch to the passed tiles (same as a4w4/a8w4) and delete the a16w4 CSV family (resolve_a16w4_gemm2_config, pick_a16w4_config, _load_a16w4_csv, _decode_a16w4_kname, _default_a16w4_csv_path, kernelName regexes, FLYDSL_A16W4_TUNED_CSV env). csv_dispatch.py is now heuristic-only -> renamed tile_heuristic.py. #4: move act.py (SiLU/SiTUv2 helpers) out of moe_2stage_a16wmix/ to the shared aiter/ops/flydsl/kernels/act.py so other FlyDSL MoE kernels can reuse it. Validated (gfx950 kimik3): logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest 6/6; a16w4 e2e perf parity (tok4/8/256/16384 within noise; get_padded_M is identity at those tokens so kernelName tiles == the old CSV-resolved tiles). pyflakes clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 task
simondanielsson
added a commit
to simondanielsson/aiter
that referenced
this pull request
Sep 9, 2026
valarLip
pushed a commit
that referenced
this pull request
Sep 11, 2026
…el for chunked prefill (#3606) * fix(mla): always write partial output + LSE for prefill PS tiles The PS scheduler in v1_2_host.cuh's generate_work_info tagged tiles that a single threadgroup could absorb from block 0 with the sentinel partial_o_loc = -1, causing generate_reduce_info to exclude them from reduce_indptr / reduce_final_map / reduce_partial_map. As a result mla_reduce_v1 never wrote final_lse[qo_start:qo_end, :] for those tiles, leaving stale workspace bytes that downstream consumers (vLLM's chunked-context MLA prefill on MI350) read as garbage LSE and fed into merge_attn_states, causing a ~2-3pp gsm8k accuracy regression on DeepSeek-V3. Always assign a real partial_o_loc so every tile produces a partial output + partial LSE that reduce processes (num_partials==1 for what used to be the fast-path tiles). The kernel already supports this path -- only the host scheduler needed changing; no assembly edits. The if(work.partial_o_loc == -1) continue; guard in generate_reduce_info is left in as defensive coding. Mirrors PR #3542 ("[Triton-Gluon-MLA-GFX950] return_lse: full decode + merged fp32 lse"), which fixed the analogous bug in the triton-gluon decode kernel's NUM_KV_SPLITS==1 fast path. * fix: dont skip reduction of tiles with num_split==1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: increase max_partial now that we support split tiles Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: compute the work and partial based on actual kv size Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: return final_lse instead of partial attn_lse from kernel wrapper Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * test: add proper final lse validation test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * test: add non-causal chunk context to correctness test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * pass max_kvlen to correctness test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: simplify comments Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant comment Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant handling of partial_o_loc=-1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: simplify comment Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: indent Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: indentation Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: avoid running reduce when num_splits==0 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: add comment about fullgraphs Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: size without mxa kvlen Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: fix docstring Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: rename lse_ref in test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add need_lse flag This avoid avoid unnecessary reduction when lse is not needed. Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: simplify comments Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: comment in get_ps_metadata_info_v1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: restructure to minimzie diff Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: use qo_len in test when available instead of ctx_len Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: remove empty eol spaces Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set max_partials to one per qo-tile + splits Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set lower lse test tolerances Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: typo Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set lse error atol to 3e-2 to be above the fp8 rounding onise Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: fix black erro Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: fill final_lse buffer with -inf (#3) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant asserts (#6) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: include cassert in v1_2_host.cuh (#7) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add return_lse param and check that ensures final_lse is populated completely if return_lse (#3 & #4) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: align flydsl reduce kernel with the cu variant Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: indent py bindings Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: convert asserts to AITER_CHECK Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: address comments about tests Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add test cases to flydsl reduce to ensre parity Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: switch to int64 slot offsets Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add test of aiter.mla.mla_prefill_ps_fwd python wrapper Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> --------- Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
2 tasks
apinge
added a commit
to yicyang/aiter
that referenced
this pull request
Sep 23, 2026
1 task done
1 task done
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
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.
No description provided.