🚀 The feature, motivation and pitch

There are two undocumented flags that improve low-latency performance on Hopper and Blackwell:

• TRTLLM_ENABLE_PDL: Enable PDL in the DeepSeek fused QKV A Projection.
• TORCHINDUCTOR_ENABLE_PDL: Enable PDL in all torch.compile'd triton kernels.

PDL is a valuable feature, but unfortunately adds some slight host overhead which can hurt in the prefill phase. We want to enable this only during FULL cuda graph execution, but not during PIECEWISE or NONE graph modes. This adds a fair bit of complexity to the implementation.

I am not sure if this can be implemented simply by just decorating set_forward_context or similar, since I don't know if/how torch.compile caches the compiled kernels; I am concerned that the piecewise graph capture might pre-populate torch's caches with non-PDL kernels and those would be used during the FULL graph capture. I'm not sure this is the case, but it seems probable and I don't see a clear path around it.

@contextmanager
def maybe_enable_pdl_for_full_cudagraph(
    cudagraph_runtime_mode: CUDAGraphMode,
) -> Iterator[None]:
    # PDL reduces kernel launch overhead when kernels are captured into a
    # full CUDA graph on SM90+, but its host-side cost outweighs the win
    # outside of one. Toggle the env vars only while capturing a FULL
    # graph so PIECEWISE / non-graph paths aren't affected, and defer to
    # the user if they have already chosen an explicit setting.
    if (
        cudagraph_runtime_mode != CUDAGraphMode.FULL
        or not current_platform.is_cuda()
        or not current_platform.has_device_capability(90)
        or "TRTLLM_ENABLE_PDL" in os.environ
        or "TORCHINDUCTOR_ENABLE_PDL" in os.environ
    ):
        yield
        return

    os.environ["TRTLLM_ENABLE_PDL"] = "1"
    os.environ["TORCHINDUCTOR_ENABLE_PDL"] = "1"
    logger.info_once(
        "Enabling PDL (TRTLLM_ENABLE_PDL=1, TORCHINDUCTOR_ENABLE_PDL=1) "
        "while capturing full CUDA graphs on SM90+ NVIDIA GPU.",
        scope="local",
    )
    try:
        yield
    finally:
        os.environ.pop("TRTLLM_ENABLE_PDL", None)
        os.environ.pop("TORCHINDUCTOR_ENABLE_PDL", None)

Alternatives

Alternatively, we can just leave it up to the caller to provide these flags or add them to the --performance-mode latency feature set.

Either way, we'll need to document this.

Additional context

No response

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.