[NO MERGE] Integrate CK varlen cross attention for small-seq (s_q=1, s_kv<=16)

[VeeraRajasekhar]

Integrate the CK team's unfused variable-length attention HIP kernels from varlen_attn/ into Transformer Engine's ROCm fused-attn path as a specialized path for specialized cross-attention (Q length 1, KV length 2-16, large batch)..

• Add fused_attn_smallseq.hpp and fused_attn_smallseq.cpp under fused_attn_rocm/: declarations and implementation adapted from varlen_attn/attn_fwd.cpp and attn_bwd.cpp (scores, mask+softmax, output; grad_V, grad_attn, softmax bwd, grad_Q/grad_K). Runtime dispatch over max_seqlen_kv in {2,4,6,8,12,16}, head_dim 128, BF16.

• Add fused_attn_smallseq.cpp to the ROCm fused-attn build in transformer_engine/common/CMakeLists.txt.

• In fused_attn_ck_fwd: when THD and no bias, branch to small-seq path when max_seqlen_q==1 and 2<=max_seqlen_kv<=16. On shape query (Aux_CTX_Tensors->size == 0) skip get_runtime_max_seqlen (cu_seqlens pointers are null); use host max_seqlen_kv and set output_S to attention-weights shape {max_tokens_q, h_q, 1, runtime_max_seqlen_kv} and dtype QKV_type. On real run (size >= 2) call get_runtime_max_seqlen then fused_attn_smallseq_fwd. Use sequence count b_varlen = max_tokens_q (not segment count b) for get_runtime_max_seqlen, output_S shape, workspace size, and small-seq fwd so varlen kernel indexing matches Q and cu_seqlens_kv (THD may pass segment-level cu_seqlens; varlen kernel expects sequence-level batch).

• In fused_attn_ck_bwd: same THD/small-seq condition. On workspace query (workspace->data.dptr == nullptr) skip get_runtime_max_seqlen and use host max_seqlen_kv; on real run call get_runtime_max_seqlen then fused_attn_smallseq_bwd. Use b_varlen = max_tokens_q_bwd for get_runtime_max_seqlen, workspace size, and small-seq bwd.

• Reuse softmax LSE auxiliary buffer for attention weights in the small-seq path (forward write, backward read);

• JAX attention.py: in NVTE_CK block, when THD and q_max_seqlen==1 and kv_max_seqlen<=16 set softmax_shape = (*batch_shape, attn_heads, q_max_seqlen, kv_max_seqlen) and softmax_dtype = q_dtype so Python aux buffer matches C++ attention-weights convention.

• Add test_ck_unfused_smallseq_backend in tests/jax/test_fused_attn.py (parametrized s_kv in {2,4,6,8,12,16}, b=30720, s_q=1, THD_THD_THD, SeqDescFormat.Seqlens) and optional NVTE_LOG_CK_SMALLSEQ debug logging in C++.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[wangye805]

Let's make this PR work for jax extension first. Later we can support pytorch.

One key difference btw jax and pytorch fused-attn dispatch is that pytorch can calculate, request, and allocate softmax_aux, workspace during runtime with actual cu_seqlen_q/kv data. However, in jax extension, softmax_aux and workspace calculation is done in

TransformerEngine/transformer_engine/jax/cpp_extensions/attention.py

Lines 364 to 375 in b685686

	if backend == NVTE_Fused_Attn_Backend.NVTE_AOTriton:
	softmax_shape = (*batch_shape, attn_heads, q_max_seqlen, config.max_segments_per_seq)
	softmax_dtype = dtypes.canonicalize_dtype(jnp.float32)
	elif backend == NVTE_Fused_Attn_Backend.NVTE_CK:
	if config.qkv_layout.is_thd():
	softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1)
	else:
	softmax_shape = (*batch_shape, attn_heads, q_max_seqlen, 1)
	softmax_dtype = dtypes.canonicalize_dtype(jnp.float32)
	else:
	raise ValueError(f"Unsupported {backend=}")
	softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype)

and

TransformerEngine/transformer_engine/jax/cpp_extensions/attention.py

Line 394 in b685686

wkspace_info = transformer_engine_jax.get_fused_attn_fwd_workspace_sizes(

, without knowing actual runtime cu_seqlen_q/kv. Aux tensors are prepared in

TransformerEngine/transformer_engine/jax/csrc/extensions/attention.cpp

Line 36 in b685686

void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch,

, also without the knowledge of runtime cu_seqlen_q/kv

General guideline:
1). Pre-allocate large enough softmax_aux and workspace ahead of time. Do not modify the aux preparation function or the c++ level aux workspace calculation/preparation, since we know our softmax aux and workspace size will be large enough for both flow, and the special flow only need a valid start pointer address.
2). During actual kernel dispatch, we do a seqlen_q/kv check, if it satisfy the special cross-attn condition, we launch it here
3). Use an env to guard this new flow and disable it when CP is used

[wangye805]

Probably no need for this cast. cudaStream_t will be hipified correctly to hipStream_t

[VeeraRajasekhar]

https://github.com/ROCm/TransformerEngine/actions/runs/22507932253/job/65210673064

CI run, no issues