[Feature][Kernel][DSR1]: Make`fused_grouped_topk` more fused (integrate TRT-LLM kernel)

[robertgshaw2-redhat]

🚀 The feature, motivation and pitch

Right now, this operation kernel has multiple steps (

vllm/vllm/model_executor/layers/fused_moe/fused_moe.py

Lines 1320 to 1350 in 4ea62b7

	def fused_grouped_topk(
	hidden_states: torch.Tensor,
	gating_output: torch.Tensor,
	topk: int,
	renormalize: bool,
	e_score_correction_bias: torch.Tensor,
	num_expert_group: int = 0,
	topk_group: int = 0,
	scoring_func: str = "softmax",
	routed_scaling_factor: float = 1.0,
	) -> tuple[torch.Tensor, torch.Tensor]:
	assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
	if scoring_func == "softmax":
	scores = torch.softmax(gating_output, dim=-1)
	elif scoring_func == "sigmoid":
	scores = gating_output.sigmoid()
	else:
	raise ValueError(f"Unsupported scoring function: {scoring_func}")
	scores_with_bias = scores + e_score_correction_bias.unsqueeze(0)
	topk_values, topk_indices = ops.grouped_topk(
	scores,
	scores_with_bias.to(scores.dtype),
	num_expert_group,
	topk_group,
	topk,
	renormalize,
	routed_scaling_factor,
	)
	return topk_values.to(torch.float32), topk_indices.to(torch.int32)

)

def fused_grouped_topk(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    e_score_correction_bias: torch.Tensor,
    num_expert_group: int = 0,
    topk_group: int = 0,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"

    if scoring_func == "softmax":
        scores = torch.softmax(gating_output, dim=-1)
    elif scoring_func == "sigmoid":
        scores = gating_output.sigmoid()
    else:
        raise ValueError(f"Unsupported scoring function: {scoring_func}")

    scores_with_bias = scores + e_score_correction_bias.unsqueeze(0)
    topk_values, topk_indices = ops.grouped_topk(
        scores,
        scores_with_bias.to(scores.dtype),
        num_expert_group,
        topk_group,
        topk,
        renormalize,
        routed_scaling_factor,
    )
    return topk_values.to(torch.float32), topk_indices.to(torch.int32)

We should make this one single kernel to do the:

• sigmoid
• addition
• output types

Alternatives

See below. we should pull in the kernel from trt-llm

Additional context

See below. we should pull in the kernel from trt-llm

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.

[robertgshaw2-redhat]

perhaps this kernel has what we need:

• https://github.com/NVIDIA/TensorRT-LLM/blob/eeb56c2848a23bb60acd82c158d086c4305b249b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp#L37
• which is used here: https://github.com/NVIDIA/TensorRT-LLM/blob/eeb56c2848a23bb60acd82c158d086c4305b249b/tensorrt_llm/_torch/models/modeling_deepseekv3.py#L638-L640

[yujiaoliang]

Hi, I’m planning to work on fusing the scoring and bias addition into a single kernel (similar to noaux_tc_op in TRT-LLM).
Will open a draft PR soon for discussion.

[robertgshaw2-redhat]

Hi, I’m planning to work on fusing the scoring and bias addition into a single kernel (similar to noaux_tc_op in TRT-LLM). Will open a draft PR soon for discussion.

Thanks! Any reason why not to just pull in the kernel from TRT-LLM?