core.inference.contexts.zero_padding_rows_kernel

Triton kernel that zeroes the padding-token rows of a token-major activation.

Dynamic inference pads the token dimension, both on eager steps and up to a captured size on CUDA-graph steps. Padding rows still feed the per-tensor amax of FP8/FP4 GEMMs, so they are zeroed first. The boundary between real and padding rows is read from a fixed-address int32[1] GPU tensor, so the kernel is safe to capture in a CUDA graph: only the value behind the pointer changes between replays. It stores zeros to the padding rows only and never reads the activation.

Module Contents

Functions

_zero_padding_rows_kernel	Store zeros to rows whose token index is at or past the real token count.
zero_padding_rows	Zero, in place, the rows of tensor whose token index is at or past the real count.

API

core.inference.contexts.zero_padding_rows_kernel._zero_padding_rows_kernel(

x_ptr,

real_token_count_ptr,

num_rows,

row_numel,

first_token_idx,

BLOCK_M: triton.language.constexpr,

BLOCK_N: triton.language.constexpr,

)

Store zeros to rows whose token index is at or past the real token count.

core.inference.contexts.zero_padding_rows_kernel.zero_padding_rows(

tensor: torch.Tensor,

real_token_count: torch.Tensor,

first_token_idx: int = 0,

) → None

Zero, in place, the rows of tensor whose token index is at or past the real count.

Parameters:

• tensor – Activation whose first dimension is the padded token dimension.

• real_token_count – [1] int32 GPU tensor holding the number of real tokens. It is read on device, so the boundary follows the value at CUDA-graph replay.

• first_token_idx – Token index of row 0 of tensor, e.g. the offset of this rank’s sequence-parallel shard.