[AvivBick]

rotary_embedding_fwdbwd allocated its output with torch.empty_like, but the Triton kernel only stores rows whose rotary position o_cs is in range (m_t = (o_t>=0)&(o_t<T)&(o_cs>=0)&(o_cs<TR)). Under left-padding the per-sequence seqlen_offset is negative, so padding rows get o_cs < 0, are masked out of the store, and keep uninitialized memory. That garbage then leaks into attention on paths that do not unpad (e.g. torch SDPA): when it is NaN/Inf it propagates across the batch and RoPE models emit empty generations. (The flash-attn varlen path unpads padding rows before attention, which is why it only surfaces without flash-attn.)

Fix: allocate the output with torch.zeros_like so the kernel-skipped rows are defined rather than uninitialized.

Tests (tests/modules/test_rotary.py):

• test_rotary_left_padding / test_rotary_partial: correctness on the negative-offset (left-padding) and partial-rotary paths, which were previously untested (existing tests cover only positive offsets); output is checked against rotary_embedding_ref.
• test_rotary_left_padding_no_uninit_leak: regression guard for this bug -- seeds the allocator free list with NaN blocks the output buffer would reuse, then requires the result to be finite. Fails on the empty_like version, passes with the fix.

[gemini-code-assist]

Code Review

This pull request addresses an issue where uninitialized memory could leak into attention under left-padding by replacing torch.empty_like with torch.zeros_like in the rotary embedding module. It also introduces comprehensive tests covering left-padding, partial rotary dimensions, and a regression test for uninitialized memory leaks. The reviewer suggests using x.clone() instead of torch.zeros_like(x) to maintain behavioral consistency between the in-place and out-of-place execution paths for out-of-range rows, which also simplifies the code. Additionally, the reviewer recommends adding assertions to the tests to verify that padded prefixes are correctly preserved.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using torch.zeros_like(x) to initialize the output tensor introduces a behavioral discrepancy between inplace=True and inplace=False for the out-of-range (padded) rows. When inplace=True, the out-of-range rows in x are untouched and retain their original values. When inplace=False, the out-of-range rows in y are initialized to 0 for the rotated part (:R2), but retain their original values for the non-rotated part (R2:). To ensure consistent behavior and mathematical equivalence between the in-place and out-of-place versions, we can use x.clone() instead of torch.zeros_like(x). This also allows us to completely remove the explicit slice copy y[..., R2:].copy_(x[..., R2:]) when R2 < D, simplifying the code and potentially improving performance.

    # clone, not empty_like: the kernel does not write out-of-range rows (negative o_cs under left-padding), which must be defined so uninitialized memory cannot leak into attention.
    # Using clone also ensures that inplace=True and inplace=False behave identically for out-of-range rows.
    y = x.clone() if not inplace else x

[gemini-code-assist]

To verify that the out-of-range (padded) tokens are correctly preserved and consistent with the input (matching the behavior of inplace=True), we should also assert that the padded prefix of the output matches the input.

Suggested change

	for i, pad in enumerate(pads.tolist()):
	ref_q = rotary_embedding_ref(q[i:i+1, pad:].float(), rotary._cos_cached[:T-pad], rotary._sin_cached[:T-pad]).to(dtype=dtype)
	ref_k = rotary_embedding_ref(k[i:i+1, pad:].float(), rotary._cos_cached[:T-pad], rotary._sin_cached[:T-pad]).to(dtype=dtype)
	assert_close(f" q[{i}]", ref_q, tri_q[i:i+1, pad:], ratio=1e-5)
	assert_close(f" k[{i}]", ref_k, tri_k[i:i+1, pad:], ratio=1e-5)
	for i, pad in enumerate(pads.tolist()):
	ref_q = rotary_embedding_ref(q[i:i+1, pad:].float(), rotary._cos_cached[:T-pad], rotary._sin_cached[:T-pad]).to(dtype=dtype)
	ref_k = rotary_embedding_ref(k[i:i+1, pad:].float(), rotary._cos_cached[:T-pad], rotary._sin_cached[:T-pad]).to(dtype=dtype)
	assert_close(f" q[{i}]", ref_q, tri_q[i:i+1, pad:], ratio=1e-5)
	assert_close(f" k[{i}]", ref_k, tri_k[i:i+1, pad:], ratio=1e-5)
	if pad > 0:
	assert_close(f" q_pad[{i}]", q[i:i+1, :pad], tri_q[i:i+1, :pad], ratio=1e-5)
	assert_close(f" k_pad[{i}]", k[i:i+1, :pad], tri_k[i:i+1, :pad], ratio=1e-5)