Add --disable-sgmv flag

.github/workflows/build.yaml

@@ -1,6 +1,7 @@
 name: Build and push docker image to ghcr
 
 on:
+  pull_request:
   workflow_dispatch:
   push:
     branches:

launcher/src/main.rs

@@ -457,6 +457,9 @@ struct Args {
     /// The embedding dimension to use for the model.
     #[clap(long, env)]
     embedding_dim: Option<usize>,
+
+    #[clap(long, env)]
+    disable_sgmv: bool,
 }
 
 #[derive(Debug)]
@@ -500,6 +503,7 @@ fn shard_manager(
     shutdown: Arc<AtomicBool>,
     _shutdown_sender: mpsc::Sender<()>,
     embedding_dim: Option<usize>,
+    disable_sgmv: bool,
 ) {
     // Enter shard-manager tracing span
     let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();
@@ -608,6 +612,7 @@ fn shard_manager(
         shard_args.push(embedding_dim.to_string())
     }
 
+
     // Copy current process env
     let mut envs: Vec<(OsString, OsString)> = env::vars_os().collect();
 
@@ -640,6 +645,10 @@ fn shard_manager(
         envs.push(("FLASH_INFER".into(), "1".into()));
     }
 
+    if disable_sgmv {
+        envs.push(("DISABLE_SGMV".into(), "1".into()))
+    }
+
     // Safetensors load fast
     envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
 

server/lorax_server/utils/sgmv.py

@@ -10,6 +10,7 @@
     import punica_kernels as _kernels
 
     HAS_SGMV = not bool(os.environ.get("DISABLE_SGMV", ""))
+    print("HAS_SGMV: ", HAS_SGMV)
 except ImportError:
     warnings.warn("Could not import SGMV kernel from Punica, falling back to loop.")
     _kernels = None