Rope Utils

namespace trt_edgellm

Argument bundles for the CuTe DSL FMHA launchers.

Each public run* entry point fills one of these once from its own arguments and members, and the launcher then spreads it across the generated descriptors. One struct per wrapper shape rather than a single superset: a superset would let a field that matters on one path (kvCacheCapacity on the dense path, tokensPerPage on the paged one) sit silently at zero on another.

These carry no generated types, so this header is safe to include from translation units that use the required FMHA-v2 runner, the optional CUTE_DSL_FMHA_BLACKWELL_ENABLED runner, or both. The descriptor-filling machinery itself lives in cuteDslTensorDescriptors.h, which stays free of any FMHA concept.

Helpers for populating the tensor descriptors emitted by the CuTe DSL C exporter.

Every AOT variant exports its own nominally distinct but layout-identical descriptor structs (fmha_d64_Tensor_q_tensor_t vs fmha_d128_Tensor_q_tensor_t) plus a cute_dsl_<variant>_wrapper entry point. There is no umbrella C type, so descriptor types are recovered here from the signature of the wrapper that consumes them: a call site names only the wrapper and the kernel module, and pairing a descriptor with the wrong variant is not expressible.

The exporter (cutlass/cute/export/c_header_generator.py) always names the members data, dynamic_shapes and dynamic_strides, but emits each array only when its dynamic mask is non-empty. A rank-1 descriptor therefore has no dynamic_strides member at all and needs makeCuSeqLenTensor() rather than the strided builders below.

Functions

inline int64_t getRotaryDim(

nlohmann::json const &configJson,

int64_t headDim

)

Return the RoPE cos/sin cache width for one attention head.

headDim is the full per-head Q/K dimension. rotaryDim is the width of the runtime rope_rotary_cos_sin cache consumed by the attention plugin. For the usual partial-RoPE representation, it is also the number of head channels that receive RoPE values; channels outside rotaryDim bypass rotation in the attention kernel.

rope_scaling is the normalized runtime/HuggingFace object that selects the RoPE variant and optional scaling parameters. It is not Gemma-specific: Gemma4 sliding/full RoPE configs, LongRoPE, MRoPE, dynamic scaling, and proportional RoPE all use this common field shape after export normalization.

Most RoPE variants use partial_rotary_factor to shrink rotaryDim to the rotated prefix of the head. Proportional RoPE is the exception: it keeps a full-head cache (rotaryDim == headDim) and consumes partial_rotary_factor in collectRopeConfig() as the fraction of full-head angle slots that receive non-identity cos/sin values. The inactive slots are materialized as identity values (cos=1, sin=0). The attention plugin does not need a Gemma-specific path for this case because it reads the cache width from the binding shape and already accepts rotaryDim <= headDim.