[Model] Add Inkling compressed-tensors dynamic FP8 support (#48876)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Krishna Teja Chitty-Venkata committed
Jul 28, 2026 at 22:27 UTC
17a74b745b459eb72ea9eea2945bc83b1c17e0a0
2 files changed
+29
tests/models/inkling/test_moe_weight_layout.py
+26
@@ -187,6 +187,32 @@ def test_moe_loads_compressed_tensors_global_scale(
187
assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"]
188
189
190
+@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)])
191
+def test_moe_loads_channelwise_scale_for_tp(
192
+ projection: str, checkpoint_rows: int
193
+) -> None:
194
+ param = torch.nn.Parameter(torch.empty(2, 4, 1))
195
+ experts = SimpleNamespace(
196
+ **{f"{projection}_weight_scale": param},
197
+ moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)),
198
+ )
199
+ layer = SimpleNamespace(
200
+ experts=SimpleNamespace(routed_experts=experts),
201
+ _local_expert_slots=lambda: {0: 0, 2: 1},
202
+ )
203
+ checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1)
204
+
205
+ loaded = moe.InklingMoE.load_expert_weight(
206
+ layer, f"experts.{projection}_weight_scale", checkpoint_scale
207
+ )
208
+
209
+ expected = checkpoint_scale[[0, 2]]
210
+ if projection == "w13":
211
+ expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2)
212
+ torch.testing.assert_close(param, expected.float())
213
+ assert loaded == [f"experts.routed_experts.{projection}_weight_scale"]
214
+
215
+
216
def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None:
217
monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2)
218
monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1)
vllm/models/inkling/nvidia/moe.py
+3
@@ -589,6 +589,9 @@ class InklingMoE(nn.Module):
589
param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to(
590
param.device
591
)
592
+ elif key == "w2_weight_scale" and weight.shape[-1] == 1:
593
+ # Per-output-channel scales are replicated across TP ranks.
594
+ param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype)
595
elif key.startswith("w13"):
596
# Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the
597
# fused param layout is [w1(gate); w3(up)]. The TP-local rows form