Make prefix block size configurable via env var

docs/architecture.md

@@ -81,7 +81,7 @@ These components are maintained in the `llm-d-inference-scheduler` repository an
 | Scorer           | Description                                | Env Vars |
 |------------------|--------------------------------------------|----------|
 | Session-aware    | Prefers pods from same session             | `ENABLE_SESSION_AWARE_SCORER`, `SESSION_AWARE_SCORER_WEIGHT`, `PREFILL_ENABLE_SESSION_AWARE_SCORER`, `PREFILL_SESSION_AWARE_SCORER_WEIGHT` |
-| Prefix-aware     | Matches prompt prefix                      | `ENABLE_PREFIX_AWARE_SCORER`, `PREFIX_AWARE_SCORER_WEIGHT`, `PREFILL_ENABLE_PREFIX_AWARE_SCORER`, `PREFILL_PREFIX_AWARE_SCORER_WEIGHT` |
+| Prefix-aware     | Matches prompt prefix                      | `ENABLE_PREFIX_AWARE_SCORER`, `PREFIX_AWARE_SCORER_WEIGHT`, `PREFILL_ENABLE_PREFIX_AWARE_SCORER`, `PREFILL_PREFIX_AWARE_SCORER_WEIGHT`, `PREFIX_SCORER_BLOCK_SIZE`|
 | KVCache-aware    | Optimizes for KV reuse                     | `ENABLE_KVCACHE_AWARE_SCORER`, `KVCACHE_INDEXER_REDIS_ADDR`, `PREFILL_ENABLE_KVCACHE_AWARE_SCORER`, `PREFILL_KVCACHE_INDEXER_REDIS_ADDR`, `HF_TOKEN`, `KVCACHE_INDEXER_REDIS_ADDR` |
 | Load-aware       | Avoids busy pods                           | `ENABLE_LOAD_AWARE_SCORER`, `LOAD_AWARE_SCORER_WEIGHT`, `PREFILL_ENABLE_LOAD_AWARE_SCORER`, `PREFILL_LOAD_AWARE_SCORER_WEIGHT` |
 

pkg/config/config.go

@@ -47,6 +47,9 @@ const (
 	pdEnabledEnvKey             = "PD_ENABLED"
 	pdPromptLenThresholdEnvKey  = "PD_PROMPT_LEN_THRESHOLD"
 	pdPromptLenThresholdDefault = 100
+
+	prefixScorerBlockSizeEnvKey  = "PREFIX_SCORER_BLOCK_SIZE"
+	prefixScorerBlockSizeDefault = 256
 )
 
 // Config contains scheduler configuration, currently configuration is loaded from environment variables
@@ -55,8 +58,9 @@ type Config struct {
 	DecodeSchedulerPlugins  map[string]int
 	PrefillSchedulerPlugins map[string]int
 
-	PDEnabled   bool
-	PDThreshold int
+	PDEnabled       bool
+	PDThreshold     int
+	PrefixBlockSize int
 }
 
 // NewConfig creates a new instance if Config
@@ -67,6 +71,7 @@ func NewConfig(logger logr.Logger) *Config {
 		PrefillSchedulerPlugins: map[string]int{},
 		PDEnabled:               false,
 		PDThreshold:             math.MaxInt,
+		PrefixBlockSize:         prefixScorerBlockSizeDefault,
 	}
 }
 
@@ -86,6 +91,7 @@ func (c *Config) LoadConfig() {
 
 	c.PDEnabled = env.GetEnvString(pdEnabledEnvKey, "false", c.logger) == "true"
 	c.PDThreshold = env.GetEnvInt(pdPromptLenThresholdEnvKey, pdPromptLenThresholdDefault, c.logger)
+	c.PrefixBlockSize = env.GetEnvInt(prefixScorerBlockSizeEnvKey, prefixScorerBlockSizeDefault, c.logger)
 }
 
 func (c *Config) loadPluginInfo(plugins map[string]int, prefill bool, pluginNames ...string) {

pkg/scheduling/pd/scheduler.go

@@ -58,11 +58,14 @@ type Datastore interface {
 // NewScheduler returns a new disaggregated Prefill/Decode filter, using the
 // provided configuration.
 func NewScheduler(ctx context.Context, schedCfg *config.Config, ds Datastore) (*Scheduler, error) {
+	prefixConfig := scorer.DefaultPrefixStoreConfig()
+	prefixConfig.BlockSize = schedCfg.PrefixBlockSize
+
 	scheduler := &Scheduler{
 		threshold:    schedCfg.PDThreshold,
 		pdEnabled:    schedCfg.PDEnabled,
 		store:        ds,
-		prefixScorer: scorer.NewPrefixAwareScorer(ctx, nil),
+		prefixScorer: scorer.NewPrefixAwareScorer(ctx, prefixConfig),
 	}
 
 	scheduler.prefill = scheduling.NewSchedulerWithConfig(

pkg/scheduling/plugins/scorer/prefix_aware.go

@@ -38,6 +38,10 @@ var _ plugins.Scorer = &PrefixAwareScorer{} // validate interface conformance
 // NewPrefixAwareScorer creates a new PrefixAwareScorer with the given
 // PrefixStoreConfig. If the config is nil, default is used.
 func NewPrefixAwareScorer(ctx context.Context, config *PrefixStoreConfig) *PrefixAwareScorer {
+	if config == nil {
+		config = DefaultPrefixStoreConfig()
+	}
+
 	scorer := &PrefixAwareScorer{
 		prefixStore:     NewPrefixStore(config),
 		podToPromptHits: sync.Map{},