[Kernel] Scale elementwise kernel launches with the device, not just batch size

[XiaobingSuper]

Overview

silu_and_mul and the other act_and_mul variants, the scalar activations,
swigluoai_and_mul, situ_and_mul and rotary_embedding all launch grid(num_tokens)
with one block per token, so the grid is only as wide as the batch however much work a
token carries. At decode that leaves most of the device idle. This splits a token's work
across gridDim.y when the batch alone cannot fill the device.

Claims

• Up to 7.05x on situ_and_mul, 4.36x on swigluoai_and_mul, 3.42x on silu_and_mul
  and 1.91x on rotary_embedding at decode-sized batches.
• Medians 1.31x (activations), 1.19x (rope) and 1.58x (gated activations).
• Unchanged at batches that already fill the device: the helper returns the previous
  configuration, and measurements there are 0.994x to 1.004x.
• Output is bit-identical to the unpatched kernels on every tested shape.
• One shared helper, no API change, no new kernel entry point.

Validation

Bitwise: 288 shapes (silu_and_mul and rotary_embedding, bf16 and fp16, 1 to 2048
tokens, neox and non-neox, several head counts) compared with torch.equal against the
unpatched build — not a tolerance.

Performance, MI355X (gfx950, 256 CUs), ROCm 7.2.4, HIP-graph replay:

kernel	shape	before	after
situ_and_mul	d=28672, 1-2 tokens	14.9-15.3us	2.1-2.2us	7.0x
swigluoai_and_mul	d=28672, 1 token	9.0us	2.1us	4.4x
silu_and_mul	d=28672, 1-16 tokens	7.2-8.0us	2.2-2.5us	up to 3.4x
silu_and_mul	d=14336, 64 tokens	4.52us	2.82us	1.60x
rotary_embedding	32 heads, 8-32 tokens	4.3-4.4us	2.3us	up to 1.91x

Medians: 1.31x over 288 activation/rope shapes, 1.58x over 108 gated-activation shapes.
Where the batch already fills the device the configuration is unchanged by construction
and the measurements agree — 0.994x to 1.004x for the first group, 0.98x worst case for
the second.

situ_and_mul gains the most because its loop is scalar, so it had no other source of
parallelism within a token.

This change is not ROCm-specific and has also been verified on NVIDIA hardware, where it
likewise shows a gain.

Details

Root cause. These kernels launch one block per token and stride that block over the
hidden dim, so a batch of 8 launches 8 blocks. They are nowhere near memory bound at that
point — they simply have no blocks to run.

Change. vllm_elementwise_launch_config() in csrc/launch_bounds_utils.h, shared by
every call site in activation_kernels.cu and pos_encoding_kernels.cu. When num_tokens >= sm_count it
returns the existing configuration (chunks == 1) so large-batch launches are untouched;
below that it uses smaller blocks and splits a token's work over chunks blocks in
gridDim.y. Every element is still handled by exactly one thread reading the same inputs,
which is why results do not move.

Where this helper fits. It suits kernels whose loop strides over the hidden dim, so
any (threads, chunks) split is valid. It does not suit kernels that give each
blockIdx.y a contiguous chunk and vectorize inside it — there a small chunk leaves too
few elements per thread to form a wide load. All five kernels here are the first kind.

Scope. The activation and rope ops are CustomOps, and CustomOp.default_on() is false when
Inductor is in use, so in the default compiled configuration they are replaced by
Inductor-generated Triton kernels and this patch does not run. A Llama-3.1-8B decode
profile confirms it: the activation and norm work appears as triton_poi_fused_* with no
vllm::act_and_mul_kernel at all. Where the patch does apply — enforce_eager, or
-O.custom_ops=+silu_and_mul / +rotary_embedding — the same profile shows these kernels
on the decode path at 4.6us and 4.4us per call, which is the regime measured above. So
this is a kernel-level fix for the configurations that use these kernels, not an
end-to-end win for the default Llama path.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: One helper in an existing header, reused by both call sites; no API
  or caller change, and the previous configuration is what it returns at large batch.

• Testing and Validation: Bitwise sweep above; performance numbers need an MI355X
  and are not reproducible in CI.

• Code Quality and Style: pre-commit clean; comments say why, not what.

• Pull Request Contents: Claims, evidence, root cause and scope limits above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[XiaobingSuper]

@tjtanaa please help review it. Thanks,