[Bug]: compressed-tensors W4A16 MoE: weight_scale not sharded along K under tensor parallelism, kernel computes wrong group_size #41511

Your current environment

• vLLM commit: 0ac3de079 = jasl/ds4-sm120 base 428e08e + neuralmagic/kylesayrs/deepseek-ct (PR [CompressedTensors] DeepSeek4 CT Quantization Support #41276) f910a73a cherry-picked + a packed_modules_mapping class-attr patch on DeepseekV4ForCausalLM (which the PR references at line 205 but doesn't define on the class)
• The 4 commits between live tree base and current jasl/ds4-sm120 HEAD 68901da (06e11f8, 2148a6e, 2e06cbd, 68901da) do not modify vllm/model_executor/layers/quantization/compressed_tensors/. Verified empty: git log 06e11f8..68901da -- vllm/model_executor/layers/quantization/compressed_tensors/. Behavior on 68901da is the same by construction for the affected code path
• PyTorch: 2.11.0+cu130
• CUDA: 13.0
• GPU: 8× NVIDIA H200 (SM 9.0), 141 GB HBM3e each
• Model: pastapaul/DeepSeek-V4-Flash-W4A16-FP8 (FP8_BLOCK attention + W4A16 GPTQ routed experts, group_size=128, produced via llm-compressor PR #2647
  branch kylesayrs/transformers-v5, attention quant mirrors RedHatAI/DeepSeek-V4-Flash-NVFP4-FP8 topology)

🐛 Describe the bug

When serving a compressed-tensors W4A16 MoE model under tensor parallelism, the Marlin MoE kernel moe_wna16_marlin_gemm raises:

RuntimeError: Invalid thread config: thread_m_blocks = 4, thread_k = -1,
thread_n = -1, num_threads = -1 for MKN = [49152, 256, 4096], num_bits = 4,
group_size = 16, has_act_order = 0, is_k_full = 0, has_zp = 0,
is_zp_float = 0, max_shared_mem = 232448

Captured live at TP=8 (2026-05-02 16:05:31 UTC, all 8 worker ranks fired identical error during gpu_worker.profile_run → _dummy_run → model(...) → MoE forward → moe_wna16_marlin_gemm).

Note that the kernel selector logs the correct group_size from config at engine startup:

INFO compressed_tensors_moe_wna16_marlin.py:87] Using Marlin backend for WNA16 MoE (group_size=128, num_bits=4)

…but at kernel-dispatch time the GEMM call sees group_size = 16. That's a 1×8 mismatch under TP=8, exactly the K-shard ratio.

Empirical observations across TP sizes

The TP=2 case is the smoking gun: the kernel believes group_size=64, but the model produces correct output anyway because reading the same scale tensor at finer granularity than its actual quantization grouping is a no-op (each scale value is repeated across smaller windows). This proves the scales are correct on disk and the bug is purely in how vLLM passes them to the kernel under TP.

Root cause

weight_scale is not sharded along K when the routed expert weight is sharded along K under TP. num_groups is derived at runtime from the unsharded scale's shape:

group_size = K_per_rank / num_groups = K_per_rank / 16

Yields:

• TP=1: 2048 / 16 = 128 ← correct
• TP=2: 1024 / 16 = 64 ← off by 2×, kernel template exists, no crash
• TP=4: 512 / 16 = 32 ← off by 4×, may or may not have template
• TP=8: 256 / 16 = 16 ← off by 8×, no template, hard crash

The weight tensor is correctly sharded:

• Full down_proj.weight: (out_dim, K=2048)
• Per-rank shard at TP=8: (out_dim, K_shard=256)

The scale tensor is not sharded along K:

• Saved down_proj.weight_scale: (out_dim, num_groups=16) (verified on-disk; see diagnostic below)
• Per-rank scale: (out_dim, num_groups=16) ← unchanged across ranks

When compressed_tensors_moe_wna16_marlin calls Marlin, it passes the unsharded scale alongside the sharded weight. Marlin's moe_wna16_marlin_gemm validation derives group_size = size_k / num_groups where size_k is the per-rank K. At TP=8, group_size collapses to 16, which is below MIN_THREAD_K = 128, so is_valid_config rejects every entry in large_batch_thread_configs[] and the kernel selector returns MarlinDefault, yielding the cryptic thread_k=-1 error.

Reproduction

Recipe (mirrors RedHatAI/DeepSeek-V4-Flash-NVFP4-FP8 attention group + W4A16 GPTQ on routed experts):

from llmcompressor.modifiers.quantization import (
    GPTQModifier, QuantizationModifier
)
from compressed_tensors.quantization import (
    QuantizationScheme, QuantizationArgs
)

# Group 1: FP8_BLOCK attention (data-free, matches RedHat exactly)
fp8_attn = QuantizationScheme(
    targets=[
        "re:model.*attn.*(wkv|wo_a|wo_b|wq_a|wq_b)$",
        "re:model.*attn\\.compressor.*(wgate|wkv)$",
    ],
    weights=QuantizationArgs(
        num_bits=8, type="float", symmetric=True,
        strategy="block", block_structure=[128, 128],
        observer="memoryless_minmax",
    ),
    input_activations=QuantizationArgs(
        num_bits=8, type="float", symmetric=True,
        group_size=128, strategy="group", dynamic=True,
    ),
)

# Group 2: W4A16 GPTQ on routed experts (group_size=128)
w4a16_experts = QuantizationScheme(
    targets=["re:model.*mlp.*(gate|up|down)_proj$"],
    weights=QuantizationArgs(
        num_bits=4, type="int", symmetric=True,
        strategy="group", group_size=128,
    ),
)

recipe = [
    QuantizationModifier(
        config_groups={"attention": fp8_attn},
        ignore=["lm_head", "re:.*shared_experts.*"],
    ),
    GPTQModifier(
        config_groups={"experts": w4a16_experts},
        ignore=["lm_head", "re:.*shared_experts.*", "re:.*attn.*"],
        dampening_frac=0.1,
        sequential_targets=["DeepseekV4DecoderLayer"],
    ),
]

Serve commands to reproduce:

# Crashes (confirmed 2026-05-02)
vllm serve <model_dir> --tensor-parallel-size 8 \
  --kv-cache-dtype fp8 --block-size 256 --max-model-len 16384 \
  --gpu-memory-utilization 0.85 --enforce-eager \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --enable-auto-tool-choice --trust-remote-code

# Works (confirmed, full smoke + harness gates pass)
vllm serve <model_dir> --tensor-parallel-size 2 \
  --kv-cache-dtype fp8 --block-size 256 --max-model-len 16384 \
  --gpu-memory-utilization 0.85 --enforce-eager \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --enable-auto-tool-choice --trust-remote-code

Verification of saved scale shape (proof scales are correct on disk)

from safetensors import safe_open
with safe_open("model-00001-of-00004.safetensors", framework="pt") as sf:
    t = sf.get_tensor(
        "model.layers.0.ffn.experts.0.down_proj.weight_scale"
    )
    print(t.shape)  # → (4096, 16) — group_size=128 over K=2048 = 16 groups ✓

Suggested fix location

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py around process_weights_after_loading (lines ~215–225 in current HEAD), where the routed expert weight is sharded along K at TP boundaries but weight_scale is left at full-tensor group count.

In the weight loading / create_weights path, when the routed expert weight is sharded along K at TP boundaries, the weight_scale parameter must also shard along K:

# Sharding logic (current code is missing this for weight_scale):
num_groups_per_shard = num_groups // tp_size
scale_shard = weight_scale[:, k_group_start:k_group_end]

For comparison, the FP8 W8A8 MoE path at compressed_tensors_moe_w8a8_fp8.py correctly shards weight_scale_inv along K under TP for block-quantized scales. The W4A16 grouped path appears to skip the equivalent sharding.

Impact

This blocks all compressed-tensors W4A16 MoE deployments under TP > 2 on every GPU architecture, not just SM12x. Specifically affects:

• Intel/DeepSeek-V4-Flash-W4A16-AutoRound (Intel's published model card notes "vLLM and SGLang is not supported currently" — same root cause class)
• All future GPTQ/AWQ W4A16 quantizations of MoE models targeting vLLM TP>2

Note that TP=2 deployments work because the kernel's wrong-believed group_size=64 happens to have a template instantiation; output is mathematically correct because per-group scales are constant within each subdivision. So TP=2 is a viable workaround until the fix lands, but the kernel dispatch is sub-optimal (smaller-than-necessary tiles).

Related issues / PRs

• [Bug]: Marlin MoE kernel fails with MXFP4-quantized GPT-OSS 20B - Invalid thread config for non-aligned dimensions (K=2880, N=2880) #38022 / PR [Bugfix] Add dimension alignment check to Marlin MoE kernel selection #38222: separate Marlin MoE shape-alignment bug for non-128-aligned hidden dims (different root cause; same symptom class).
• PR [CompressedTensors] DeepSeek4 CT Quantization Support #41276 (kylesayrs compressed-tensors V4 support): handles attention scale_fmt and Linear-wrap fixes; does not touch MoE TP scale sharding.
• PR [DSv4][Nvidia] SM12x DeepSeek V4 support #40991 (jasl SM12x V4): adds SM12x kernel fallbacks; does not touch compressed-tensors MoE quantization paths. Confirmed bug is independent of jasl's recent indexer cache fix (68901da) and FP8 einsum full-group fix (2148a6e) — git log 06e11f8..68901da -- vllm/model_executor/layers/quantization/compressed_tensors/ is empty.

Repro artifacts

Full integration findings, scripts, and reproduction details:
https://github.com/pasta-paul/dsv4-flash-awq-w4a16

Specifically:

• Recipe: scripts/quantize_v4_w4a16.py
• Diagnostic that confirmed scales-correct-on-disk: scripts/inspect_scales.py
• 5 prior upstream gaps documented: findings/kylesayrs-pr-41276-integration.md
• This bug detail: findings/upstream-issue-marlin-tp-sharding.md

Before submitting a new issue

• Searched for relevant issues
• Verified bug fires on 0ac3de079 = jasl 428e08e + PR [CompressedTensors] DeepSeek4 CT Quantization Support #41276 cherry-pick (Python content equivalent to current 68901da HEAD for the affected code path; intervening commits do not touch compressed_tensors_moe_wna16_marlin.py)
• Confirmed independent of [Bug]: Marlin MoE kernel fails with MXFP4-quantized GPT-OSS 20B - Invalid thread config for non-aligned dimensions (K=2880, N=2880) #38022 (different shape constraint, same kernel symptom class)
• Saved checkpoint scales verified correct on disk
• TP=2 confirms output is mathematically correct (workaround) — full chat-smoke + toolcall15 harness pass
• TP=8 confirms bug fires deterministically across all 8 ranks