Skip to content

fix(moe): auto-degrade FP8 block_n/k from 128 to 64 on alignment mismatch (#442)#1488

Open
haowu1234 wants to merge 1 commit into
ROCm:mainfrom
haowu1234:codex/fix-fp8-moe-block-alignment
Open

fix(moe): auto-degrade FP8 block_n/k from 128 to 64 on alignment mismatch (#442)#1488
haowu1234 wants to merge 1 commit into
ROCm:mainfrom
haowu1234:codex/fix-fp8-moe-block-alignment

Conversation

@haowu1234

Copy link
Copy Markdown

Problem

Models with intermediate_size=1536 (MiniMax M2.5, Qwen3-235B) at TP=8 produce intermediate_size_per_partition=192, which is not divisible by block_n=128 (the default for per_1x128 block quantisation). Both CompressedTensorsFp8MoEMethod and Fp8MoEMethod raise ValueError, forcing users to fall back to TP=4 and waste 4 GPUs.

The existing error:

ValueError: The output_size of gate's and up's weight = 192 is not divisible by weight quantization block_n = 128.

Additionally, after fixing block_n, block_k=128 also fails at TP=8 (192 % 128 != 0), which was masked by the earlier block_n error.

Fix (Option B from #442)

When intermediate_size_per_partition is not divisible by 128 but is divisible by 64, auto-degrade block_n and/or block_k to 64 and emit a logger.warning. The ValueError is preserved when alignment is truly impossible (e.g. 190 % 64 != 0).

Changes

  • CompressedTensorsFp8MoEMethod.create_weights: auto-degrade self.block_n and self.block_k before alignment checks.
  • Fp8MoEMethod.__init__: initialise self.block_n/self.block_k instance attributes (were local variables in create_weights only).
  • Fp8MoEMethod.create_weights: auto-degrade both block sizes using instance attributes.
  • Fp8MoEMethod.get_fused_moe_quant_config: use [self.block_n, self.block_k] instead of hardcoded [128, 128] that would mismatch the degraded scale shapes.
  • tests/test_fp8_moe_block_align.py: 13 unit tests covering degrade path, aligned path, impossible alignment, per_1x32 no-op, scale shape correctness, block_k degradation at TP>1, and full TP=8/1536 end-to-end scenarios.

Validation

  • 13/13 tests passed on MI300X (gfx942, ROCm 7.1, Python 3.12, aiter 0.1.1.dev1)
  • scale shapes verified: block_n=block_k=64w13_scale=(E, 6, 112), w2_scale=(E, 112, 3) for 192 intermediate at TP=8
  • existing TP=4 / aligned paths unchanged

Refs: #442

…nment mismatch

Models with intermediate_size=1536 at TP=8 produce
intermediate_size_per_partition=192, which is not divisible by
block_n=128 or block_k=128 (the defaults for per_1x128 block
quantisation).  Both CompressedTensorsFp8MoEMethod and
Fp8MoEMethod raise ValueError, forcing users to fall back to
TP=4 and waste 4 GPUs (MiniMax M2.5, Qwen3-235B, etc.).

Fix (Option B from ROCm#442): When intermediate_size_per_partition
is not divisible by 128 but is divisible by 64, auto-degrade
block_n and/or block_k to 64 and emit a logger warning.  The
ValueErrors are preserved when alignment is still impossible.

Changes:
- CompressedTensorsFp8MoEMethod.create_weights: auto-degrade
  self.block_n and self.block_k before alignment checks.
- Fp8MoEMethod: initialise self.block_n/self.block_k in __init__
  (was local variables in create_weights only); auto-degrade in
  create_weights; use them in get_fused_moe_quant_config instead
  of the hardcoded [128,128] that would mismatch the degraded
  scale shapes.
- tests/test_fp8_moe_block_align.py: 13 unit tests covering
  degrade path, aligned path, impossible alignment, per_1x32
  no-op, scale shape correctness, block_k degradation at TP>1,
  and full TP=8/1536 end-to-end scenarios.

Validated: 13/13 passed on MI300X (gfx942, ROCm 7.1, Python 3.12).

Refs: ROCm#442
@zufayu zufayu requested a review from ZhangLirong-amd July 7, 2026 02:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant