Fix nccl blocking case (#1228)

EricLBuehler · web-flow · commit 085900049f6c · 2025-03-22T14:54:42.000-04:00
diff --git a/mistralrs-core/src/attention.rs b/mistralrs-core/src/attention.rs
@@ -340,7 +340,7 @@ impl Sdpa {
         let k = repeat_kv(k.clone(), sdpa_params.n_kv_groups)?;
         let v = repeat_kv(v.clone(), sdpa_params.n_kv_groups)?;
 
-        if mask.is_some_and(|x| x.rank() == 2) {
+        if mask.is_some_and(|x| x.rank() == 2) || mistralrs_quant::distributed::use_nccl() {
             return naive_sdpa(q, &k, &v, mask, sdpa_params);
         }
 

Original file line number	Diff line number	Diff line change
`@@ -340,7 +340,7 @@ impl Sdpa {`
`340`	`340`	`let k = repeat_kv(k.clone(), sdpa_params.n_kv_groups)?;`
`341`	`341`	`let v = repeat_kv(v.clone(), sdpa_params.n_kv_groups)?;`
`342`	`342`
`343`		`- if mask.is_some_and(\|x\| x.rank() == 2) {`
	`343`	`+ if mask.is_some_and(\|x\| x.rank() == 2) \|\| mistralrs_quant::distributed::use_nccl() {`
`344`	`344`	`return naive_sdpa(q, &k, &v, mask, sdpa_params);`
`345`	`345`	`}`
`346`	`346`