Enable triton kernels matmul tests

[wdziurdz]

Fixes #5074

[whitneywhtsang]

Is this for testing? If it is not ready for review, please convert to draft.

[wdziurdz]

The last eight tests could be fixed by this PR: #5128.
Example of one failed test:

AssertionError: ref_y_scale: 0.004773152060806751, tri_y_scale: 0.005022321827709675
  assert tensor(False, device='xpu:0')
   +  where tensor(False, device='xpu:0') = <built-in method all of type object at 0x7f4175d82400>(tensor([0.0002], device='xpu:0', grad_fn=<AbsBackward0>) < 1e-10)
   +    where <built-in method all of type object at 0x7f4175d82400> = torch.all
   +    and   tensor([0.0002], device='xpu:0', grad_fn=<AbsBackward0>) = <built-in method abs of Tensor object at 0x7f406d589260>()
   +      where <built-in method abs of Tensor object at 0x7f406d589260> = (tensor(0.0048, device='xpu:0', grad_fn=<DivBackward0>) - tensor([0.0050], device='xpu:0')).abs

With this rounding, for the test parameters below we could to pass failed tests:

    sep_scatter = mode == "ragged" and do_scatter and n_expts_act > 1 and split_k == 1
    if is_xpu():  # additional rounding on XPU for references values
        sep_scatter = sep_scatter or (do_scatter and not fused_scatter and n_expts_tot > 1 and split_k == 1 and act_dtype_str == "float8_e4m3fn")
    y_scale = flex.out_data.expected_scale if act_is_float8 else 1

    def round_x(x, idx):
        return x.to(act_dtype).to(torch.float32) if sep_gather else x

    round_y = lambda y: (y / y_scale).to(act_dtype).to(torch.float32) * y_scale if sep_scatter else y
    ref_y = matmul_ogs_torch(x_ref, w_ref, bias_ref,  #
                             rdata, gindx, sindx, round_x=round_x, round_y=round_y, gammas=gs1_ref,
                             inner_routing_data=inner_routing_data, device=device)

Without this fix, there’s look on small precision difference. It’s sufficient to pass all assertions, but the test fails when comparing the actual_scale for these parameters.

Pls looks on this @etiotto @whitneywhtsang

[etiotto]

Given that there is a proposed fix upstream we should wait for that change to land.

[wdziurdz]

Waiting for the cherry-pick commit from upstream: triton-lang/triton@d8119d5

[anmyachev]

LGTM! I've used your changes in #5598. They helped a lot!