[Kernel] Initial Machete W4A8 support + Refactors

[LucasWilkinson]

• add machete kernel support for QQQ style w4a8 quantization (including fp8 activations)
• refactor machete dispatching logic
• refactor machete file generation
• reduce the number of prepack kernels generated (e.g. don't need separate kernels for uint4b8 and uint4 since we are just shuffling data around)

TODO (Future PR):

• end2end integration
• perf improvements (mostly hoping to land now for the refactor)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[ProExpertProg]

2 minor comments, looks good otherwise

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. @LucasWilkinson please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[varun-sundar-rabindranath]

nit : for big models, I have found it useful to have their realistic TPn counter-parts also (e.g. for the 70B case, add a 70B-TP4 case). That way we can just list that version in the 1GPU model benchmarking.

[LucasWilkinson]

You mean so you can list it as a string as opposed to using the --tp-sizes args?

[tlrmchlsmth]

What is the right shift for?

[LucasWilkinson]

updated the comment to be more verbose reads as:

      // Determines if to get from the signed or unsigned candidates
      //  move into bit position 0x4 of each nibble so when or'd with
      //  final_prmt_base it selects the correct candidate, when elements
      //  in final_prmt_base are >= 0x4, the negative candidate is selected
      //  (i.e. from NEG_INT8_REG{1}{2}), when elements are < 0x4, the positive
      //  candidate is selected (i.e. from POS_INT8_REG{1}{2})
      uint32_t sign = (src & 0x88888888) >> 1;

      // `sign` is OR'd with 0x31203120 to find the correct value in the LUT
      // (selects correct positive or negative candidate)
      const uint32_t final_prmt_base = 0x32103210;

      // Ignore sign bit when indexing into LUT, for each 4bit value
      //  we index into both the positive and negative candidates then use
      //  sign | final_prmt_base to select the correct candidate
      uint32_t lut_idx = (src & 0x77777777);

[LucasWilkinson]

nvm refactored to use lut_4bit_to_8bit_convert utility function, now the comment looks like

  // Determines if the value is in the top half of the LUT if set or
  //  (i.e. LUT[8:15]) in the bottom half (i.e. LUT[0:7]) if not set. Then move
  //  into bit position 0x4 of each nibble so when or'd with final_prmt_base it
  //  selects the correct candidate. When elements in final_prmt_base
  //  are >= 0x4, the high candidate is selected (i.e. LUT[8:15]), when elements
  //  are  < 0x4, the low candidate is selected (i.e. LUT[0:7])
  uint32_t high_bit = (src & 0x88888888) >> 1;

  // `high_bit` is OR'd with 0x31203120 to find the correct value in the LUT
  // (selects correct high or low candidate)
  const uint32_t final_prmt_base = 0x32103210;

  // Ignore the high bit when indexing into LUT, for each 4bit value
  //  we index into both the high and low candidates then use
  //  high_bit | final_prmt_base to select the correct candidate
  uint32_t lut_idx = (src & 0x77777777);

[tlrmchlsmth]

cool!

[tlrmchlsmth]

It looks like the int4 -> fp8 and int4 -> int8 converters the same except for these constants? If so, might be nice to factor these out. Not a big deal though, because it'd be kind of annoying to do.

[LucasWilkinson]

Good call, bit janky but refactored to lut_4bit_to_8bit_convert

[tlrmchlsmth]

Could you explain this change?

[LucasWilkinson]

Its just moved up (to be closer to SmemLayoutA) from below, where the following is deleted:

   using SmemLayoutACopy = decltype(tile_to_shape(
       SmemLayoutAtomARowMajor{},
       make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideA>(),
                     Step<_2, _1, _3>, Step<_1, _2, _3>>{}));

the

 conditional_t<::cutlass::gemm::detail::is_major<0, StrideA>(),
                     Step<_2, _1, _3>, Step<_1, _2, _3>>{}))

is removed since it was just cruft from the original PR thats not actually exercised (that was my bad)

[varun-sundar-rabindranath]

Reviewed the cutlass refactor part - LGTM!

[tlrmchlsmth]

Great work, LGTM!

[mgoin]

LGTM thanks for getting it green!