obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-12 00:38:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Model] Add Inkling compressed-tensors dynamic FP8 support (#48876)

Browse Source 
Signed-off-by: mgoin <[email protected]>
Co-authored-by: mgoin <[email protected]>
Co-authored-by: OpenAI Codex <[email protected]>
This commit is contained in:
Krishna Teja Chitty-Venkata
2026-07-28 19:27:01 -07:00
committed by GitHub
co-authored by mgoin OpenAI Codex
parent 54ab69b14e
commit 17a74b745b
2 changed files with 29 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/inkling/test_moe_weight_layout.py
+26
View File
@@ -187,6 +187,32 @@ def test_moe_loads_compressed_tensors_global_scale(
assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"]
@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)])
def test_moe_loads_channelwise_scale_for_tp(
projection: str, checkpoint_rows: int
) -> None:
param = torch.nn.Parameter(torch.empty(2, 4, 1))
experts = SimpleNamespace(
**{f"{projection}_weight_scale": param},
moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)),
)
layer = SimpleNamespace(
experts=SimpleNamespace(routed_experts=experts),
_local_expert_slots=lambda: {0: 0, 2: 1},
)
checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1)
loaded = moe.InklingMoE.load_expert_weight(
layer, f"experts.{projection}_weight_scale", checkpoint_scale
)
expected = checkpoint_scale[[0, 2]]
if projection == "w13":
expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2)
torch.testing.assert_close(param, expected.float())
assert loaded == [f"experts.routed_experts.{projection}_weight_scale"]
def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None:
monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2)
monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1)
vllm/models/inkling/nvidia/moe.py
+3
View File
@@ -589,6 +589,9 @@ class InklingMoE(nn.Module):
param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to(
param.device
)
elif key == "w2_weight_scale" and weight.shape[-1] == 1:
# Per-output-channel scales are replicated across TP ranks.
param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype)
elif key.startswith("w13"):
# Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the
# fused param layout is [w1(gate); w3(up)]. The TP-local rows form