[Feature]: General LL GEMMs with PDL Support

[robertgshaw2-redhat]

🚀 The feature, motivation and pitch

We currently have two highly specialized low latency GEMMS that support PDL after AR

• dsv3_a_gemm:

  vllm/csrc/dsv3_fused_a_gemm.cu

  Line 46 in 6241521

  char const* env = std::getenv("TRTLLM_ENABLE_PDL");

• dsv3_router_gemm:

  vllm/csrc/moe/dsv3_router_gemm_utils.h

  Line 38 in 6241521

  const char* env = std::getenv("TRTLLM_ENABLE_PDL");

These:

• currently only support specific shapes
• only support bf16

We should create generalized versions of these low-latency PDL enabled GEMMS, which support:

• bf16
• fp8
• nvfp4
• arbitrary shapes

Alternatives

None

Additional context

PDL fix in flashinfer: flashinfer-ai/flashinfer#2887

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[benchislett]

There's also the gpt-oss router GEMM (buggy, see #38754): #37205

[benchislett]

Note also that PDL can be enabled (opt-in) for torch.compile, but this only applies to the triton kernels and not the GEMMs:
https://github.com/pytorch/pytorch/blob/6e98ebebaf4e631bf2371fa35935dfe8a5be8333/torch/_inductor/config.py#L1870-L1873

TORCHINDUCTOR_ENABLE_PDL=1

[jeejeelee]

Router GEMM for M25 :#38445

[benchislett]

@robertgshaw2-redhat @LopezCastroRoberto @LucasWilkinson we might get this for free with FlashInfer's TinyGEMM:
#39921