[ROCm][Kimi-K3] Enable gfx942 serving - #50319
Conversation
|
I have read the DCO document and hereby sign off past commits made by me. |
| # gfx942 has no native MXFP4 matmul. Convert the weights to groupwise | ||
| # int4 once at load time for AITER's existing bf16 x int4 FlyDSL path. | ||
| from aiter import dtypes as aiter_dtypes | ||
| from aiter.ops.quant import per_1x32_i4_quant | ||
| from aiter.ops.shuffle import ( | ||
| pack_int8_to_packed_int4, | ||
| shuffle_scale_for_int4, | ||
| shuffle_weight, | ||
| ) | ||
| from aiter.utility import fp4_utils | ||
|
|
||
| fp4_dtype = torch.float4_e2m1fn_x2 | ||
| e8m0_dtype = torch.float8_e8m0fnu | ||
|
|
||
| def convert( | ||
| weight: torch.nn.Parameter, | ||
| scale: torch.nn.Parameter, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| weight_f32 = fp4_utils.mxfp4_to_f32(weight.data.view(fp4_dtype)) | ||
| scale_f32 = fp4_utils.e8m0_to_f32(scale.data.view(e8m0_dtype)) | ||
| num_experts, output_size, input_size = weight_f32.shape | ||
| weight_bf16 = ( | ||
| ( | ||
| weight_f32.view(num_experts, output_size, input_size // 32, 32) | ||
| * scale_f32.view(num_experts, output_size, input_size // 32, 1) | ||
| ) | ||
| .view(num_experts, output_size, input_size) | ||
| .to(torch.bfloat16) | ||
| ) | ||
| del weight_f32, scale_f32 | ||
|
|
||
| weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16) | ||
| del weight_bf16 | ||
| weight_int4 = weight_int4.view(aiter_dtypes.i4x2).view( | ||
| num_experts, output_size, input_size | ||
| ) | ||
| weight_packed = pack_int8_to_packed_int4( | ||
| shuffle_weight(weight_int4.view(aiter_dtypes.i8), (16, 16)) | ||
| ) | ||
| weight_packed = weight_packed.view( | ||
| num_experts, output_size, input_size // 2 | ||
| ).view(aiter_dtypes.i4x2) | ||
| weight_scale = ( | ||
| shuffle_scale_for_int4(weight_scale, group_size=32) | ||
| .view(-1) | ||
| .contiguous() | ||
| ) | ||
| return weight_packed, weight_scale | ||
|
|
||
| w13, w13_scale = convert(layer.w13_weight, layer.w13_weight_scale) | ||
| w2, w2_scale = convert(layer.w2_weight, layer.w2_weight_scale) | ||
| replace_parameter(layer, "w13_weight", w13) | ||
| replace_parameter(layer, "w2_weight", w2) | ||
| replace_parameter(layer, "w13_weight_scale", w13_scale) | ||
| replace_parameter(layer, "w2_weight_scale", w2_scale) | ||
| layer.w13_weight.is_shuffled = True | ||
| layer.w2_weight.is_shuffled = True | ||
|
|
||
| # The router emits fp32 logits, so aiter's biased_grouped_topk upcasts | ||
| # the bf16 correction bias to match on every step -- a per-layer, | ||
| # per-token bf16->fp32 copy kernel. Keep the tiny [num_experts] bias in | ||
| # fp32 so the runtime cast becomes a no-op. | ||
| correction_bias = getattr(layer, "e_score_correction_bias", None) | ||
| if correction_bias is not None: | ||
| correction_bias.data = correction_bias.data.to(torch.float32) |
There was a problem hiding this comment.
This conversion logic should make use of convert_gpt_oss_weight_to_mxfp4_moe_kernel_format (bad function name):
vllm/vllm/model_executor/layers/quantization/mxfp4.py
Lines 329 to 342 in e04a30a
Fine to do that in an other PR as #50000 is already not really respecting that..
There was a problem hiding this comment.
nevermind, since you requantize, convert_gpt_oss_weight_to_mxfp4_moe_kernel_format should not be used.
Still, the mxfp4.py codebase will needs cleanup as discussed (maybe in a followup PR, but preferably here).
Using:
- Dequantization from
mxfp4.py vllm/model_executor/layers/quantization/online/int4.py
would be nicer. You get the benefit of being able to dispatch to any INT4 MOE backend, and dissociating the weight loading and modeling/kernel dispatch logic. The current convert logic belongs as an INT4 MOE backend, not really to mxfp4.py.
Having a w-int4 model still use an mxfp4_backend: Mxfp4MoeBackend (that is, here Mxfp4MoeBackend.AITER_MXFP4_BF16 that priorities the centralized AiterExperts:
vllm/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
Lines 224 to 232 in 5df9999
|
Is this validated with 1 node (8 MI300) or at least requires 2 nodes (16 MI300)? |
There was a problem hiding this comment.
IMO there should be an extension in AITER or somewhere for mxfp4-bf16 mixed-precision MOE on gfx942 similar to: https://github.com/ROCm/aiter/blob/1aeb91bad0f2e71140834acbdf187ae9ca39112d/aiter/ops/flydsl/moe_kernels.py#L554 (ROCm/aiter#2863, flydsl int4-bf16). There is ongoing: ROCm/aiter#4398
Converting to INT4 may not be accuracy-free (or at least should be verified as well), and I am not sure whether there is a latency benefit from using int4-bf16 instead of mxfp4-bf16.
Existing implems:
- triton path:
moe_gemm_a16w4into: https://github.com/ROCm/aiter/blob/78e45124a09c0c0e8b831501429f938ffaccc1cc/aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a16w4.py#L305 should work on gfx942 with w-mxfp4-a-bf16 - flydsl path: https://github.com/ROCm/aiter/blob/1aeb91bad0f2e71140834acbdf187ae9ca39112d/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py#L1425. Unfortunately this only run on gfx950.
At least, from the API standpoint, silent conversion from mxfp4 to int4 is not something that exists currently in vLLM.
The triton path can be tested on gfx942 with: python3 op_tests/op_benchmarks/triton/bench_moe_gemm_a16w4.py --shape 7168 384 --experts 896 16. Seems like there is no tuning infra for it at the moment, though.
| if on_gfx942(): | ||
| self._setup_kernel_k3_situ_gfx942(layer) | ||
| return |
There was a problem hiding this comment.
This silently converts weights from mxfp4 -> int4, which is not an existing pattern in vLLM
Closest pattern I see is online re-quantization, see #48427
There was a problem hiding this comment.
On load time quantization is actually better for performance than quantization at compute. That is why I would like to keep this quantization at load time. It does not hurt accuracy as we can see from the tests.
There was a problem hiding this comment.
Yes, I get it, what you are looking to do is to override the weight_quant_key. I am saying that doing so silently is probably not the way to go; at least long term, maybe as a quick unblocker merging this would be OK, up to maintainers.
Typical API would rather be to use: vllm serve moonshotai/Kimi-K3 --quantization-config.moe.weight int4.
Or enforce through vllm serve --moe-backend aiter_int4_bf16 (see #50000 (comment)), which does not exist at the moment as only aiter is usable as MOE backend (and I think vllm's intent is rather to rely on quant key override rather than extend aiter_* MOE backends shorthands)
See https://docs.vllm.ai/en/stable/features/quantization/online/#activation-overrides-on-already-quantized-checkpoints for activation. The API does not exist at the moment for weight quantization. Similar API is proposed in #48427.
My point is also that if a mixed-precision mxfp4-bf16 MOE kernel is available for gfx942, it'd be a safer/more simple option to use that and avoid conversion altogether.
There was a problem hiding this comment.
Activation bug: while investigating your comment I noticed that the int4 path computes the wrong function. compile_flydsl_moe_stage1 forwards act, situ_beta and situ_linear_beta to every branch except b_dtype == "int4", and compile_moe_gemm1 has no act parameter, so y = silu(vg) * vu is hardcoded at both stage1 epilogue sites. K3 asks for SiTUv2. I confirmed it on gfx942 by compiling the same shape twice: act="silu" and act="situv2" return the identical executable on the int4 path, and two different ones on the mxfp4 path. Cosine against a SiTUv2 reference is 0.961917, against SiLU it is 0.999976.
Native mxfp4: I agree it is the right target and it removes the API question entirely, since the mxfp4 branch already implements SiTUv2. It does not compile on gfx942 today. It hits LLVM ERROR: Do not know how to expand this operator's operand! because the FlyDSL heuristic fallback emits a 128 bit llvm.amdgcn.raw.ptr.buffer.load.lds and CDNA3 only does 4 byte global to LDS. ROCm/aiter#3926 adds the software decode that fixes this. Measured on 8x MI325X TP8 at 8192 in and 1024 out it is throughput neutral at -0.14 percent, gives 19.1 percent more KV cache, and cuts load conversion from 112.8 s to 8.7 s. That PR is an open draft with no movement since 2026-07-13.
Triton moe_gemm_a16w4: swiglu only in v0.1.19. No situ in aiter/ops/triton/moe/moe_op_gemm_a16w4.py or its _triton_kernels counterpart, so it needs SiTUv2 written before it can serve K3.
Current plan: fix the activation in AITER first, then drop the silent override here and gate the conversion behind an explicit flag. I will split out the MLA head padding and the sampler fallback so they can land separately. Native mxfp4 is the follow-up once #3926 has an owner.
|
This pull request has merge conflicts that must be resolved before it can be |
This requires two nodes on MI300X, while on MI325X it can run on a single node. |
We have not been giving the image built on this yet publicly, even if it has been validated. We are working on getting this image up in the https://hub.docker.com/r/vllm/vllm-openai-rocm repository. It will then run PP2xTP8 on two MI300X nodes. |
Is one MI300x8 possible? |
85c84a8 to
c22d0c7
Compare
|
Rebased onto current main and reworked the gfx942 path. Activation fix: ROCm/aiter#4471 adds SiTUv2 to the packed-int4 FlyDSL stage1 epilogue. Without it that path drops the requested activation and hardcodes SiLU, so K3 serves fluent text while computing the wrong function. This PR now refuses to load against an AITER that predates it rather than repeating the failure silently. Quantization API: the MXFP4 to int4 requantization no longer happens because the hardware matched. It is opt-in through Rebase notes: current main gained a native MXFP4 A16W4 path for K3 on gfx950 and the MLA small-head Gluon decode, so my MLA head padding commit and the registry commit are both dropped as redundant. The diff is down to 7 files. Still open: the accuracy numbers in the description were measured on builds that computed SiLU, so they do not validate the activation. I will rerun gsm8k once #4471 lands and update the table. Native MXFP4 on gfx942 through ROCm/aiter#3926 stays the follow-up. |
|
This pull request has merge conflicts that must be resolved before it can be |
gfx942 (MI325X, MI300X) has no native MXFP4 matmul, so the day-0 Kimi-K3 expert path falls through to code that dies in LLVM codegen. Convert the MXFP4 expert weights to int4 with groupwise bf16 scales at load time and let the existing FlyDSL SiTU stage1 kernel consume them. The gfx950 native MXFP4 path is left untouched. Two supporting changes are needed for the same reason. The AITER MoE expert backend refuses kMxfp4Static outside gfx950 and does not list the SiTU activation as supported. The AITER top-k/top-p sampler is a gfx950-only prebuilt that segfaults on gfx942, so sampling falls back to the native torch implementation there. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 2253680)
Until now the only way to get a runnable gfx942 image was to build from Dockerfile.rocm and then apply an AITER upgrade by hand, which lived in a scratch directory and was not reproducible by anyone else. This commits that step. It swaps AITER rather than rebuilding the base, because Dockerfile.rocm_base pins torch, triton, flash-attention and MORI to the same commits the published base images already use. Building those again costs hours and changes nothing. The base pin is still corrected for anyone who does want a full from-source base. The base image is a required build arg so no registry-specific tag is baked into the file. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit d714094)
The fork branch existed to carry gfx942 tuned GEMM configs. Those configs are measured and correct, 110 shapes with no failures, but they do not move end to end throughput on Kimi-K3: 30.99, 61.87, 109.32 and 170.25 output tokens per second at concurrency 1, 2, 4 and 8, against 32.0, 61.8, 109.5 and 169.5 for the current release image. The dense bf16 GEMM path is not the bottleneck for this workload, so there is nothing to ship. Upstream tag v0.1.19 is commit 3135022616, which is exactly what the fork branch was based on. Pinning the tag keeps the build free of any dependency on a personal fork. The tuned rows stay on the fork branch for reference. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 8928b23)
The MoE router gate is built with out_dtype=fp32 (grouped_topk needs fp32 logits), but the fused fp32-output GEMM tiers are gated on is_cuda(), so on ROCm the gate falls back to a bf16 GEMM plus a separate output.to(fp32), emitting a bf16→fp32 copy kernel before grouped_topk every layer. hipBLASLt supports the same torch.mm out_dtype epilogue, so enable the fused tier on ROCm (bias-free, bf16 weight, fp32 out). CUDA behavior is unchanged. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Matvei Pashkovskii <mpashkov@amd.com> (cherry picked from commit ad8ec83) (cherry picked from commit d364e99)
The router gate emits fp32 logits, so aiter's biased_grouped_topk upcasts the bf16 e_score_correction_bias to match on every step, adding a per-layer, per-token bf16->fp32 copy kernel before grouped_topk. Cast the tiny [num_experts] bias to fp32 once at load so the runtime cast is a no-op. Mirrors the DeepSeek-V2 ROCm aiter path. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Matvei Pashkovskii <mpashkov@amd.com> (cherry picked from commit 05d85e2) (cherry picked from commit 34c98a7)
Bind the optional correction bias before converting it to fp32 so static type checking can prove the parameter is present. Keep the ported weight conversion formatted consistently with current main. Co-authored-by: Cursor <cursoragent@cursor.com>
The gfx942 path requantized Kimi-K3's MXFP4 experts to groupwise int4 whenever the hardware matched, without the user asking. The conversion is lossy, so it is now opt-in through --quantization-config.moe.weight int4 and gfx942 keeps the native MXFP4 path otherwise. Also refuse to load when the installed AITER predates ROCm/aiter#4471. Before that fix the packed-int4 stage1 dropped the requested activation and hardcoded SiLU, so Kimi-K3 served fluent text while computing SiLU instead of the SiTUv2 its config asks for. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
The gfx942 Kimi-K3 path asks for groupwise int4 MoE weights through --quantization-config.moe.weight, but QUANT_KEY_NAMES had no name for that scheme, so the flag was rejected before the path could be selected. Register the existing kInt4Static32 key under int4_per_group_32 and match on the parsed QuantKey rather than a raw string. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
Fold the residual add and output RMSNorm into the attention-residual Triton kernel, removing three launches per decoder layer without scratch storage. MI325X TP=8: output throughput improved 2.88%, TPOT improved 2.96%, and mean TTFT improved by 35 ms. All requests completed without GPU faults. Add ROCm coverage for fused decode and prefill paths with and without the residual add. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> (cherry picked from commit f709f63)
KimiRoutedOutputTransform is replicated and fed a rank-identical input, so every rank recomputes the same up projection. It is linear, so each rank can contract only its own 1/tp latent slice; the partials are summed by the all-reduce the combined MoE output already performs, at no extra cost. MI325X TP=8, 8192-token prefill: 4.83x on the projection, mean TTFT 2840 -> 2678 ms, TPOT unchanged. Gated on VLLM_KIMI_ROW_PARALLEL_LATENT_UP_PROJ (token threshold, default 0 = off); set above the decode batch size, decode stays bitwise identical. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> (cherry picked from commit 2619848)
c9d3d07 to
38db061
Compare
|
Test results on a build from this branch with the activation fix applied. gsm8k, 1319 samples, 5-shot, 8x MI325X TP8 with expert parallel: 96.06% ± 0.54 flexible-extract and 96.06% strict-match. The previous numbers in the description were 95.68% and 95.60%, measured on a build that computed SiLU. Serving is healthy, KV cache 640,942 tokens, Important caveat on the AITER pin. These numbers were produced against ROCm/aiter#4471, not the tagged v0.1.19 that this PR currently pins. v0.1.19 does not contain the SiTUv2 fix for the packed-int4 stage1 epilogue, so building this PR as pinned reproduces the SiLU behaviour. The pin needs to move to a release that includes #4471 before these numbers are reproducible from the branch alone. Also registered |
mxfp4_to_f32 splits packed nibbles with repeat_interleave and then gathers through an f32 lookup table, so the working tensor is 8x the packed weight before per_1x32_i4_quant shrinks it again. Materializing that for a whole expert tensor peaks above 20 GiB per rank and fails once the weights are resident, which is what pure tensor parallel hits since it holds all experts per rank. Convert 8 experts at a time and free each slice, bounding the transient without changing the result. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
38db061 to
9f5d657
Compare
The ROCm branch of mxfp4_round_up_hidden_size_and_intermediate_size() rounds the per-partition intermediate size up to 256. Kimi-K3 has moe_intermediate_size 3072, so a TP8 shard is 384 and gets rounded to 512. That inflates every w13 tensor from (896, 768, 1792) to (896, 1024, 1792), a 33 percent increase on all 92 MoE layers, which is about 38 GiB per rank. Pure TP8 then sits at 248.69 GiB resident with 2.12 GiB free and dies during the int4 conversion. The round-up is not needed here. AITER's resolve_flydsl_stage1_tile_n() already downgrades tile_n from 256 to 128 for a non-256-aligned inter_dim, and 128 divides 384 exactly. Mxfp4MoEMethod already skipped the round-up for the gfx950 SiTU path, so extend the same condition to the gfx942 int4 path. Measured on 8 MI325X at TP8 without expert parallel, max-model-len 16384, gpu-memory-utilization 0.97: resident drops from 248.69 GiB to 192.51 GiB, which matches the expert-parallel figure exactly, free rises from 2.12 GiB to 58.61 GiB, all 96 shards load with zero out-of-memory workers, and the KV cache is 948,305 tokens.
Is there TPS number for single concurrency? |
…2 branch Grafts the 22-commit MoRIIO delta (multi-pod PD disaggregation, DP-EP-16 bring-up, router-authoritative KV-notify DP routing, sampler warm-up fixes) onto the maeehart Kimi-K3 gfx942 model support (PR vllm-project#50319). Clean merge, no conflicts. K3 model tree (vllm/models/kimi_k3) intact. Known gap: MoRIIO connector transfers attention KV blocks only; K3 hybrid KDA layers (~69/93) hold mamba conv+recurrent state not yet carried across the P/D boundary. Colocated EP works; true 1P/1D disagg needs connector state-transfer extension (follow-up). Co-Authored-By: Claude <noreply@anthropic.com>
| from vllm.platforms.rocm import on_gfx942, on_gfx950, on_gfx1250 | ||
|
|
||
| if not on_gfx950() or on_gfx1250(): | ||
| if not (on_gfx950() or on_gfx942()) or on_gfx1250(): |
There was a problem hiding this comment.
Don't think it should necessarily block this PR, but AiterExperts handling every other so weight/activation quant key (dtype), activation, gfx, etc. is IMO risking breaking the MOE oracles and/or wrongfully selecting it in a case that is NOT supported (e.g. AiterExperts + gfx942 + (kMxfp4Static, kMxfp4Dynamic))
If the only change is in supported activation / quant key / _supports_current_device, imo we should have AiterInt4Experts, AiterMxfp4Experts, etc. similar to:
possibly inheriting from AiterExperts if needed
| # gfx942 has no native MXFP4 matmul. Convert the weights to groupwise | ||
| # int4 once at load time for AITER's existing bf16 x int4 FlyDSL path. | ||
| from aiter import dtypes as aiter_dtypes | ||
| from aiter.ops.quant import per_1x32_i4_quant | ||
| from aiter.ops.shuffle import ( | ||
| pack_int8_to_packed_int4, | ||
| shuffle_scale_for_int4, | ||
| shuffle_weight, | ||
| ) | ||
| from aiter.utility import fp4_utils | ||
|
|
||
| fp4_dtype = torch.float4_e2m1fn_x2 | ||
| e8m0_dtype = torch.float8_e8m0fnu | ||
|
|
||
| def convert( | ||
| weight: torch.nn.Parameter, | ||
| scale: torch.nn.Parameter, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| weight_f32 = fp4_utils.mxfp4_to_f32(weight.data.view(fp4_dtype)) | ||
| scale_f32 = fp4_utils.e8m0_to_f32(scale.data.view(e8m0_dtype)) | ||
| num_experts, output_size, input_size = weight_f32.shape | ||
| weight_bf16 = ( | ||
| ( | ||
| weight_f32.view(num_experts, output_size, input_size // 32, 32) | ||
| * scale_f32.view(num_experts, output_size, input_size // 32, 1) | ||
| ) | ||
| .view(num_experts, output_size, input_size) | ||
| .to(torch.bfloat16) | ||
| ) | ||
| del weight_f32, scale_f32 | ||
|
|
||
| weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16) | ||
| del weight_bf16 | ||
| weight_int4 = weight_int4.view(aiter_dtypes.i4x2).view( | ||
| num_experts, output_size, input_size | ||
| ) | ||
| weight_packed = pack_int8_to_packed_int4( | ||
| shuffle_weight(weight_int4.view(aiter_dtypes.i8), (16, 16)) | ||
| ) | ||
| weight_packed = weight_packed.view( | ||
| num_experts, output_size, input_size // 2 | ||
| ).view(aiter_dtypes.i4x2) | ||
| weight_scale = ( | ||
| shuffle_scale_for_int4(weight_scale, group_size=32) | ||
| .view(-1) | ||
| .contiguous() | ||
| ) | ||
| return weight_packed, weight_scale | ||
|
|
||
| w13, w13_scale = convert(layer.w13_weight, layer.w13_weight_scale) | ||
| w2, w2_scale = convert(layer.w2_weight, layer.w2_weight_scale) | ||
| replace_parameter(layer, "w13_weight", w13) | ||
| replace_parameter(layer, "w2_weight", w2) | ||
| replace_parameter(layer, "w13_weight_scale", w13_scale) | ||
| replace_parameter(layer, "w2_weight_scale", w2_scale) | ||
| layer.w13_weight.is_shuffled = True | ||
| layer.w2_weight.is_shuffled = True | ||
|
|
||
| # The router emits fp32 logits, so aiter's biased_grouped_topk upcasts | ||
| # the bf16 correction bias to match on every step -- a per-layer, | ||
| # per-token bf16->fp32 copy kernel. Keep the tiny [num_experts] bias in | ||
| # fp32 so the runtime cast becomes a no-op. | ||
| correction_bias = getattr(layer, "e_score_correction_bias", None) | ||
| if correction_bias is not None: | ||
| correction_bias.data = correction_bias.data.to(torch.float32) |
There was a problem hiding this comment.
nevermind, since you requantize, convert_gpt_oss_weight_to_mxfp4_moe_kernel_format should not be used.
Still, the mxfp4.py codebase will needs cleanup as discussed (maybe in a followup PR, but preferably here).
Using:
- Dequantization from
mxfp4.py vllm/model_executor/layers/quantization/online/int4.py
would be nicer. You get the benefit of being able to dispatch to any INT4 MOE backend, and dissociating the weight loading and modeling/kernel dispatch logic. The current convert logic belongs as an INT4 MOE backend, not really to mxfp4.py.
Having a w-int4 model still use an mxfp4_backend: Mxfp4MoeBackend (that is, here Mxfp4MoeBackend.AITER_MXFP4_BF16 that priorities the centralized AiterExperts:
vllm/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
Lines 224 to 232 in 5df9999
|
Closing this in favour of a smaller opt-in optimization PR. Two things changed since I opened it. #50516 landed the lossless emulation fallback, so gfx942 already constructs and serves on main and this PR is no longer what enables it. #50761 landed the fp32 correction bias, which makes the post load cast here redundant. More importantly, #50817 reaches the same goal through AITER's Triton The parts of this PR that are worth keeping are independent of the MoE backend choice, so I will send them separately rather than as one 13 commit change: the attention residual and RMSNorm fusion, the chunked expert conversion that fixes the TP8 out of memory at load, and the MLA small head arch gate. The new PR will be opt in behind a flag. |
|
@maeehart Thanks, is it based on TP=8 and PP=2? (for single concurrency 2.23 tok/s) |
|
No. The 2.23 tok/s result was measured on a single 8x MI325X node with TP=8 and PP=1. |
Getting 81 tok/s for 1 concurrency, with GSM8K of scores 97%: https://github.com/microsoft/Tutel#steps-for-kimi-k3glm-5x-claude-code-mode |
6 participants
Summary
Scope of the PR.
PR #50089 added the Kimi-K3 model and kernels. This PR is the ROCm gfx942 enablement follow-up.
Validation
Test plan
mainafter PR [Model] Add Kimi K3 support: model files and kernels [1/N] #50089 merged.AI assistance
AI tools assisted with implementation, conflict resolution, and drafting. I reviewed the changed code and validation results.
Made with Cursor