fix: separate L1/L2 scale buffers (different K_sf), fix assembly calls

2026-05-17 07:43:05 +00:00
parent b824b838a9
commit 37fecb588f
2 changed files with 36 additions and 20 deletions
--- a/tests/test_scale_assembly.py
+++ b/tests/test_scale_assembly.py
@@ -70,7 +70,10 @@ def test_scale_assembly():
        # Path 2: _assemble_scales_cudagraph_safe (GPU-only)
        expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE)
        expert_offsets[1:] = torch.tensor(tokens_per_expert, dtype=torch.int32).cumsum(0)
-        scale_a_cudagraph = runner._assemble_scales_cudagraph_safe(x_sf, expert_offsets)
+        scale_a_cudagraph = runner._assemble_scales_cudagraph_safe(
+            x_sf, expert_offsets,
+            runner._padded_x_sf_buf_l1, runner._per_expert_scale_bufs_l1
+        )

        # Compare
        # Note: shapes may differ due to padding, but the data in the