diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py index 45123e37..89dfad34 100644 --- a/cosmos_framework/model/generator/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -353,8 +353,15 @@ def multi_control_two_way_attention( - ctrl_i output: from pass i only - noisy output: w_1 * noisy_out_1 + ... + w_N * noisy_out_N (weighted sum) - All SDPA calls are maskless → Flash Attention is always active. - N=1, w=1.0 → identical to ``two_way_attention``. + All SDPA calls are maskless. N=1, w=1.0 → identical to ``two_way_attention``. + + Dense vs varlen: + Single-sample inference uses the dense attention API (no cumulative/max + seqlen args), matching ``two_way_attention``. That lets the chooser keep + cuDNN on sm_120 (workstation Blackwell), where Cosmos prefers cuDNN but + the current cuDNN integration rejects varlen and would otherwise fall back + to NATTEN kernels that lack sm_120 images. Training / multi-sample packs + keep the varlen path. Padding safety: Both ``get_causal_seq`` and ``get_full_only_seq`` can return padded rows. @@ -374,22 +381,34 @@ def multi_control_two_way_attention( noisy_s, noisy_e = split_info.noisy_token_range weights = split_info.control_weights - # ── 1. Text self-attention (causal, unchanged) ─────────────────────────── + # ── 1. Text self-attention (causal) ────────────────────────────────────── causal_q, causal_q_offsets = get_causal_seq(packed_query_states) causal_k, causal_k_offsets = get_causal_seq(packed_key_states) causal_v, _ = get_causal_seq(packed_value_states) + # Mirror two_way_attention: dense FMHA for single-sample inference so cuDNN + # remains eligible on sm_120; varlen only when training / multi-sample. + sample_offsets = packed_query_states["sample_offsets"] + num_samples = sample_offsets.shape[0] - 1 + use_dense = num_samples == 1 and not torch.is_grad_enabled() + use_dont_care_mask = causal_q_offsets is causal_k_offsets + if use_dense: + causal_varlen_kwargs = {} + else: + causal_varlen_kwargs = dict( + cumulative_seqlen_Q=causal_q_offsets, + cumulative_seqlen_KV=causal_k_offsets, + max_seqlen_Q=packed_query_states["max_causal_len"], + max_seqlen_KV=packed_query_states["max_causal_len"], + ) causal_res = attention( causal_q.unsqueeze(0), causal_k.unsqueeze(0), causal_v.unsqueeze(0), - cumulative_seqlen_Q=causal_q_offsets, - cumulative_seqlen_KV=causal_k_offsets, - max_seqlen_Q=packed_query_states["max_causal_len"], - max_seqlen_KV=packed_query_states["max_causal_len"], is_causal=True, causal_type=CausalType.DontCare if use_dont_care_mask else CausalType.TopLeft, + **causal_varlen_kwargs, ) causal_out = causal_res.squeeze(0).flatten(-2, -1) # [N_text, Hq*D] @@ -432,30 +451,37 @@ def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: torch._check(k.shape[0] == v.shape[0]) n_q, n_kv = q.shape[0], k.shape[0] # These lengths come from data-dependent unpadding, so they are unbacked - # symints under torch.compile. The selected attention backend (NATTEN) - # validates varlen inputs with `max_seqlen == 0` / `max_seqlen < 1` - # guards; without a positivity fact Dynamo cannot discharge `Eq(n, 0)`. - # Every control/noisy segment always has at least one token, so assert it. + # symints under torch.compile. Varlen backends validate with + # `max_seqlen == 0` / `max_seqlen < 1` guards; without a positivity fact + # Dynamo cannot discharge `Eq(n, 0)`. Every control/noisy segment always + # has at least one token, so assert it (needed for the varlen path). torch._check(n_q > 0) torch._check(n_kv > 0) - # Pass cumulative_seqlen_{Q,KV} + max_seqlen_{Q,KV} directly instead of + # Dense (inference, single sample): omit varlen args so cuDNN stays + # eligible on sm_120. Varlen (training / multi-sample): pass + # cumulative_seqlen_{Q,KV} + max_seqlen_{Q,KV} directly instead of # seqlens_{Q,KV}. The frontend derives cumulative offsets from seqlens via # `generate_varlen_parameters`, which calls `.max().item()` (a device-host # sync) and is explicitly disallowed inside a torch.compile region. Each - # pass here is a single (batch=1) packed sequence, so the cumulative - # offsets are simply [0, n]. Building them ourselves keeps the whole path - # inside the compiled graph. - zero = torch.zeros(1, dtype=torch.int32, device=q.device) - cu_seqlens_q = torch.cat([zero, torch.tensor([n_q], dtype=torch.int32, device=q.device)]) - cu_seqlens_kv = torch.cat([zero, torch.tensor([n_kv], dtype=torch.int32, device=q.device)]) + # pass here is a single packed sequence, so the cumulative offsets are + # simply [0, n]. Building them ourselves keeps the path compiled. + if use_dense: + varlen_kwargs = {} + else: + zero = torch.zeros(1, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.cat([zero, torch.tensor([n_q], dtype=torch.int32, device=q.device)]) + cu_seqlens_kv = torch.cat([zero, torch.tensor([n_kv], dtype=torch.int32, device=q.device)]) + varlen_kwargs = dict( + cumulative_seqlen_Q=cu_seqlens_q, + cumulative_seqlen_KV=cu_seqlens_kv, + max_seqlen_Q=n_q, + max_seqlen_KV=n_kv, + ) res = attention( q.unsqueeze(0), # [1, N_q, Hq, D] k.unsqueeze(0), # [1, N_kv, Hkv, D] v.unsqueeze(0), # [1, N_kv, Hkv, D] - cumulative_seqlen_Q=cu_seqlens_q, - cumulative_seqlen_KV=cu_seqlens_kv, - max_seqlen_Q=n_q, - max_seqlen_KV=n_kv, + **varlen_kwargs, ) # [1, N_q, Hq, D] return res.squeeze(0).flatten(-2, -1) # [N_q, Hq*D] diff --git a/cosmos_framework/model/generator/mot/attention_test.py b/cosmos_framework/model/generator/mot/attention_test.py index 27da6dc8..9670b28f 100644 --- a/cosmos_framework/model/generator/mot/attention_test.py +++ b/cosmos_framework/model/generator/mot/attention_test.py @@ -400,6 +400,62 @@ def fake_two_way_attention(*args: object, **kwargs: object) -> object: assert kv_to_store is None +@pytest.mark.L0 +@pytest.mark.CPU +def test_multi_control_matches_two_way_dense_inference_convention(monkeypatch: pytest.MonkeyPatch) -> None: + """Single-sample inference is dense; grad-enabled execution remains varlen.""" + + def make_pack(x: torch.Tensor): + return build_packed_sequence( + "two_way", + packed_sequence=x, + attn_modes=["causal", "full"], + split_lens=[2, 4], + sample_lens=[6], + packed_und_token_indexes=torch.tensor([0, 1], dtype=torch.long), + packed_gen_token_indexes=torch.tensor([2, 3, 4, 5], dtype=torch.long), + num_heads=1, + head_dim=8, + num_layers=1, + ) + + q_pack, split_info, _ = make_pack(torch.arange(48, dtype=torch.float32).reshape(6, 1, 8)) + raw_k_pack, _, _ = make_pack(torch.zeros(6, 1, 8)) + v_pack, _, _ = make_pack(torch.ones(6, 1, 8)) + split_info.control_stream_token_ranges = [(0, 2)] + split_info.noisy_token_range = (2, 4) + split_info.control_weights = [1.0] + + calls: list[dict[str, object]] = [] + + def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: object) -> torch.Tensor: + calls.append(kwargs) + return q + + monkeypatch.setattr(attention, "attention", fake_attention) + + with torch.no_grad(): + attention.multi_control_two_way_attention( + q_pack, + raw_k_pack, + v_pack, + split_info, + ) + + assert len(calls) == 3 # causal text, control, noisy + assert all("cumulative_seqlen_Q" not in call for call in calls) + + calls.clear() + attention.multi_control_two_way_attention( + q_pack, + raw_k_pack, + v_pack, + split_info, + ) + assert len(calls) == 3 + assert all("cumulative_seqlen_Q" in call for call in calls) + + @pytest.mark.L0 def test_dispatch_attention_rejects_incomplete_split_info() -> None: foreign_split_info = _foreign_split_info() diff --git a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py index cdb2cbcf..25a8b1e8 100644 --- a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py @@ -1010,8 +1010,8 @@ def forward( # one per control. For each pass i, KV = [text | ctrl_i | noisy]. # The final noisy output is the weighted sum of the N pass outputs: # noisy_out = w_1 * noisy_out_1 + ... + w_N * noisy_out_N - # All SDPA calls are maskless → Flash Attention always active. - # N=1, w=1.0 → identical to two_way_attention. + # Single-sample inference uses dense FMHA (cuDNN-eligible on sm_120); + # see multi_control_two_way_attention. N=1, w=1.0 → identical to two_way_attention. # # CP compatibility: control_stream_token_ranges are gen-relative global # offsets computed here, before CP sharding. Ulysses CP restores the full